Update app.py

app.py

@@ -8,12 +8,12 @@ model = AutoModelForSequenceClassification.from_pretrained("MarkAdamsMSBA24/ADRv
 
 # Define the prediction function
 def get_prediction(text):
-    inputs = tokenizer(text, return_tensors="pt", max_length=512, truncation=True, padding=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    prediction_scores = outputs.logits
-    predicted_class = torch.argmax(prediction_scores, dim=-1).item()
-    return f"Predicted Class: {predicted_class}", prediction_scores.tolist()
+    X_test = str(text).lower()
+    encoded_input = tokenizer(X_test, return_tensors='pt')
+    output = model(**encoded_input)
+    scores = output[0][0].detach()
+    scores = torch.nn.functional.softmax(scores)
+    return {"Severe Reaction": float(scores.numpy()[1]), "Non-severe Reaction": float(scores.numpy()[0])}
 
 iface = gr.Interface(
     fn=get_prediction,