Update app.py

app.py

@@ -2,21 +2,26 @@ import gradio as gr
 from transformers import AutoModelForSequenceClassification, AutoTokenizer
 import torch
 
-# Load pre-trained model and tokenizer
-model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased", num_labels=3)
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+# Load the model and tokenizer
+model_name = "willco-afk/my-model-name"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# Define a function to make predictions using the model
-def predict(text):
+def classify_text(text):
+    # Preprocess text
     inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    logits = outputs.logits
-    predicted_class = torch.argmax(logits, dim=-1).item()
-    return predicted_class
+    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]
 
-# Create Gradio interface
-iface = gr.Interface(fn=predict, inputs="text", outputs="text", live=True)
+# 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])
 
-# Launch the Gradio app
-iface.launch()
+demo.launch()