Update inference.py

inference.py

@@ -45,6 +45,11 @@ def generate_response(goal, option1, option2):
         enc1 = tokenizer(text1, return_tensors="pt", truncation=True, padding="max_length", max_length=128)
         enc2 = tokenizer(text2, return_tensors="pt", truncation=True, padding="max_length", max_length=128)
 
+        # Remove token_type_ids to avoid crash in EvoTransformer
+        enc1.pop("token_type_ids", None)
+        enc2.pop("token_type_ids", None)
+
+        # Move tensors to device
         for k in enc1:
             enc1[k] = enc1[k].to(device)
             enc2[k] = enc2[k].to(device)
@@ -56,13 +61,11 @@ def generate_response(goal, option1, option2):
         logits1 = get_logits(out1)
         logits2 = get_logits(out2)
 
-        # Sanity check shape
         if logits1.shape[-1] < 2 or logits2.shape[-1] < 2:
             raise ValueError("Model did not return 2-class logits.")
 
         score1 = logits1[0][1].item()
         score2 = logits2[0][1].item()
-
         evo_result = option1 if score1 > score2 else option2
 
     except Exception as e: