Update inference.py

inference.py

@@ -1,18 +1,33 @@
-from evo_model import EvoTransformerForClassification
 from transformers import AutoTokenizer
+from evo_model import EvoTransformerForClassification
 import torch
 
-# Load model and tokenizer once
-model = EvoTransformerForClassification.from_pretrained("trained_model")
+# Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+model = EvoTransformerForClassification.from_pretrained("trained_model")
 model.eval()
 
 def generate_response(goal, sol1, sol2):
-    prompt = f"Goal: {goal} Option 1: {sol1} Option 2: {sol2}"
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+    input_text = f"Goal: {goal}\nOption A: {sol1}\nOption B: {sol2}"
+    
+    # Tokenize input
+    inputs = tokenizer(
+        input_text,
+        return_tensors="pt",
+        padding=True,
+        truncation=True,
+    )
+    
+    # ✅ Filter out unwanted keys (e.g., token_type_ids)
+    filtered_inputs = {
+        k: v for k, v in inputs.items()
+        if k in ["input_ids", "attention_mask"]
+    }
+
+    # Predict
     with torch.no_grad():
-        outputs = model(**inputs)
-        logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
-    prediction = torch.argmax(logits, dim=1).item()
-    return "Solution 1" if prediction == 0 else "Solution 2"
+        outputs = model(**filtered_inputs)
+        logits = outputs.logits
+        predicted_class = torch.argmax(logits, dim=1).item()
 
+    return "A" if predicted_class == 0 else "B"