Update app.py

app.py

@@ -9,11 +9,19 @@ model = AutoModelForSequenceClassification.from_pretrained(model_name)  # Use th
 
 # Define a function to generate responses
 def generate_response(prompt):
+    # Tokenize the input prompt
     inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
     with torch.no_grad():
+        # Get the model output
         outputs = model(**inputs)
-    # Customize the response generation as per your model's output
-    response = tokenizer.decode(outputs.logits.argmax(dim=-1), skip_special_tokens=True)
+    
+    # Process the output logits
+    logits = outputs.logits
+    predicted_class_id = logits.argmax().item()
+    
+    # Generate a response based on the predicted class
+    response = f"Predicted class ID: {predicted_class_id}"
+    
     return response
 
 # Create a Gradio interface
@@ -28,8 +36,3 @@ iface = gr.Interface(
 # Launch the Gradio app
 if __name__ == "__main__":
     iface.launch()
-
-
-
-
-