Update inference.py

inference.py

@@ -27,20 +27,30 @@ def query_gpt35(prompt):
         return f"[GPT-3.5 Error] {e}"
 
 def generate_response(goal, option1, option2):
-    # Evo prediction
-    prompt = f"Goal: {goal}\nOption 1: {option1}\nOption 2: {option2}\nWhich is better?"
-    inputs = tokenizer([goal + " " + option1, goal + " " + option2],
-                       return_tensors="pt", padding=True, truncation=True)
-    
-    # ✅ Remove token_type_ids if it exists
-    inputs.pop("token_type_ids", None)
+    # Build inputs for option 1 and option 2
+    text1 = goal + " " + option1
+    text2 = goal + " " + option2
+
+    # Tokenize separately
+    input1 = tokenizer(text1, return_tensors="pt", padding=True, truncation=True)
+    input2 = tokenizer(text2, return_tensors="pt", padding=True, truncation=True)
+
+    # Remove token_type_ids to avoid forward() issues
+    input1.pop("token_type_ids", None)
+    input2.pop("token_type_ids", None)
 
     with torch.no_grad():
-        logits = model(**inputs)
-    pred = torch.argmax(logits, dim=1).item()
-    evo_result = option1 if pred == 0 else option2
+        logit1 = model(**input1)
+        logit2 = model(**input2)
+
+    # Get logits[0][0] since we only expect 1 class output vector per input
+    score1 = logit1[0][0].item()
+    score2 = logit2[0][0].item()
+
+    evo_result = option1 if score1 > score2 else option2
 
     # GPT-3.5 prediction
+    prompt = f"Goal: {goal}\nOption 1: {option1}\nOption 2: {option2}\nWhich is better?"
     gpt_result = query_gpt35(prompt)
 
     return {