Update app.py

app.py

@@ -1,19 +1,19 @@
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
-# 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)
+# Load the model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("huggingface_user/my-model-name")  # Use your model name here
+model = AutoModelForSequenceClassification.from_pretrained("huggingface_user/my-model-name")
 
-# Define your prediction function
-def predict(text):
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model(**inputs)
-    logits = outputs.logits
-    prediction = logits.argmax(dim=-1).item()
-    return f"Predicted Class: {prediction}"
+# Function to classify input text
+def classify_text(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        logits = model(**inputs).logits
+    predicted_class = logits.argmax().item()  # Get the predicted class
+    return f"Predicted class: {predicted_class}"
 
-# Gradio UI
-demo = gr.Interface(fn=predict, inputs="text", outputs="text")
+# Gradio Interface
+demo = gr.Interface(fn=classify_text, inputs="text", outputs="text")
 demo.launch()