Update inference.py

inference.py

@@ -2,6 +2,7 @@ import openai
 from transformers import BertTokenizer
 from evo_model import EvoTransformerForClassification
 import torch
+import os
 
 # Load Evo model
 model = EvoTransformerForClassification.from_pretrained("trained_model")
@@ -11,7 +12,6 @@ model.eval()
 tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
 
 # Set OpenAI key (assumes you have it set as ENV VAR or replace directly)
-import os
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 def query_gpt35(prompt):
@@ -31,6 +31,10 @@ def generate_response(goal, option1, option2):
     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)
+
     with torch.no_grad():
         logits = model(**inputs)
     pred = torch.argmax(logits, dim=1).item()