Update inference.py

inference.py

@@ -1,34 +1,25 @@
-import torch
 from transformers import AutoTokenizer
 from evo_model import EvoTransformerForClassification
+from init_save import initialize_and_save_model  # Ensure this line is added
+import torch
 
-# Load tokenizer and model
+# ✅ Ensure model is initialized and saved BEFORE loading
+initialize_and_save_model("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, option1, option2):
-    prompt1 = f"Goal: {goal}\nOption: {option1}"
-    prompt2 = f"Goal: {goal}\nOption: {option2}"
+def generate_response(goal, sol1, sol2):
+    prompt = f"Goal: {goal}\nOption 1: {sol1}\nOption 2: {sol2}\nWhich is better?"
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
 
-    inputs1 = tokenizer(prompt1, return_tensors="pt", padding=True, truncation=True)
-    inputs2 = tokenizer(prompt2, return_tensors="pt", padding=True, truncation=True)
+    if 'token_type_ids' in inputs:
+        del inputs['token_type_ids']  # Evo doesn't use this
 
     with torch.no_grad():
-        output1 = model(**inputs1)
-        output2 = model(**inputs2)
-
-    logits1 = output1["logits"]
-    logits2 = output2["logits"]
-
-    prob1 = torch.softmax(logits1, dim=1)[0][1].item()
-    prob2 = torch.softmax(logits2, dim=1)[0][1].item()
-
-    if prob1 > prob2:
-        suggestion = "✅ Option 1 is more aligned with the goal."
-    elif prob2 > prob1:
-        suggestion = "✅ Option 2 is more aligned with the goal."
-    else:
-        suggestion = "⚖️ Both options are equally likely."
-
-    return suggestion
+        logits = model(**inputs)
+    
+    predicted = torch.argmax(logits, dim=1).item()
+    return f"Option {predicted + 1} seems more reasonable based on EvoTransformer."