Update inference.py

inference.py

@@ -1,33 +1,34 @@
+import torch
 from transformers import AutoTokenizer
 from evo_model import EvoTransformerForClassification
-import torch
 
 # 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):
-    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
+def generate_response(goal, option1, option2):
+    prompt1 = f"Goal: {goal}\nOption: {option1}"
+    prompt2 = f"Goal: {goal}\nOption: {option2}"
+
+    inputs1 = tokenizer(prompt1, return_tensors="pt", padding=True, truncation=True)
+    inputs2 = tokenizer(prompt2, return_tensors="pt", padding=True, truncation=True)
+
     with torch.no_grad():
-        outputs = model(**filtered_inputs)
-        logits = outputs.logits
-        predicted_class = torch.argmax(logits, dim=1).item()
+        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 "A" if predicted_class == 0 else "B"
+    return suggestion