Update inference.py

inference.py

@@ -22,15 +22,13 @@ def predict(model, tokenizer, prompt, option1, option2, device):
     encoded = tokenizer(inputs, padding=True, truncation=True, return_tensors="pt").to(device)
 
     with torch.no_grad():
-        outputs = model(encoded["input_ids"])
-    
-    # Simple linear classifier logic
-    logits = torch.nn.functional.linear(outputs, model.classifier.weight, model.classifier.bias)
-    probs = torch.softmax(logits, dim=1)
+        logits = model(encoded["input_ids"])  # shape: [2, 1]
+
+    probs = torch.softmax(logits.squeeze(), dim=0)  # shape: [2]
     best = torch.argmax(probs).item()
 
     return {
         "choice": best,
-        "confidence": probs[0][best].item(),
-        "scores": probs[0].tolist(),
+        "confidence": probs[best].item(),
+        "scores": probs.tolist(),
     }