Update app.py

app.py

@@ -1,27 +1,19 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import gradio as gr
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
-import torch
 
-# Load the model and tokenizer
-model_name = "willco-afk/my-model-name"
+# Load the model and tokenizer from your Hugging Face model repository
+model_name = "willco-afk/my-model-name"  # Replace with your actual model repo name
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-def classify_text(text):
-    # Preprocess text
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+# Define your prediction function
+def predict(text):
+    inputs = tokenizer(text, return_tensors="pt")
     outputs = model(**inputs)
-    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    label = torch.argmax(probs, dim=1).item()
-    labels = ["english", "spanish", "tagalog"]
-    return labels[label]
-
-# Define the Gradio interface
-with gr.Blocks() as demo:
-    gr.Markdown("# Slang Translation Classifier")
-    input_text = gr.Textbox(label="Enter slang text", lines=1)
-    output_label = gr.Textbox(label="Predicted Language", interactive=False)
-    submit_button = gr.Button("Classify")
-    submit_button.click(classify_text, inputs=[input_text], outputs=[output_label])
+    logits = outputs.logits
+    prediction = logits.argmax(dim=-1).item()
+    return f"Predicted Class: {prediction}"
 
+# Gradio UI
+demo = gr.Interface(fn=predict, inputs="text", outputs="text")
 demo.launch()