Update inference.py

inference.py

@@ -1,54 +1,54 @@
-import os
 import torch
-from evo_model import EvoTransformer
+import openai
+import os
 from transformers import AutoTokenizer
+from evo_model import EvoTransformerV22
 from rag_utils import extract_text_from_file
 from search_utils import web_search
 
-# Load Evo model
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-model = EvoTransformer()
+model = EvoTransformerV22()
 model.load_state_dict(torch.load("evo_hellaswag.pt", map_location="cpu"))
 model.eval()
 
-def get_evo_response(query, context=None, enable_search=True):
-    search_snippets = ""
-    if enable_search:
-        snippets = web_search(query)
-        if snippets:
-            search_snippets = "\n".join(snippets)
-
-    full_context = f"{context or ''}\n\n{search_snippets}".strip()
-    input_1 = f"{query} Option 1"
-    input_2 = f"{query} Option 2"
+def format_input(question, options, context, web_results):
+    prompt = f"{question}\n"
+    if context:
+        prompt += f"\nContext:\n{context}\n"
+    if web_results:
+        prompt += f"\nWeb Search Results:\n" + "\n".join(web_results)
+    prompt += "\nOptions:\n"
+    for idx, opt in enumerate(options):
+        prompt += f"{idx+1}. {opt}\n"
+    return prompt.strip()
 
-    inputs = tokenizer([input_1, input_2], padding=True, truncation=True, return_tensors="pt")
+def get_evo_response(question, context, options, enable_search=True):
+    web_results = web_search(question) if enable_search else []
+    input_text = format_input(question, options, context, web_results)
+    encoded = tokenizer(input_text, return_tensors="pt", padding=True, truncation=True, max_length=256)
     with torch.no_grad():
-        logits = model(inputs["input_ids"]).squeeze(-1)
-        probs = torch.softmax(logits, dim=0)
-        best_idx = torch.argmax(probs).item()
+        logits = model(encoded["input_ids"])
+    probs = torch.softmax(logits, dim=1).squeeze()
+    pred_index = torch.argmax(probs).item()
+    confidence = probs[pred_index].item()
 
-    suggestion = f"Option {best_idx + 1}"
-    reasoning = (
-        f"Evo suggests: **{suggestion}** (Confidence: {probs[best_idx]:.2f})\n\n"
-        f"Context used:\n{full_context}"
-    )
-    return suggestion, reasoning
+    suggestion = options[pred_index] if pred_index < len(options) else "N/A"
+    evo_reasoning = f"Evo suggests: **{suggestion}** (Confidence: {confidence:.2f})\n\nContext used:\n" + "\n".join(web_results)
+    return suggestion, evo_reasoning
 
-def get_gpt_response(query, context=None):
-    import openai
+def get_gpt_response(question, context, options):
     openai.api_key = os.getenv("OPENAI_API_KEY", "")
-    context = context or "None"
+    formatted_options = "\n".join([f"{i+1}. {opt}" for i, opt in enumerate(options)])
+    prompt = f"Question: {question}\n\nContext:\n{context}\n\nOptions:\n{formatted_options}\n\nWhich option makes the most sense and why?"
 
     try:
         response = openai.ChatCompletion.create(
             model="gpt-3.5-turbo",
             messages=[
-                {"role": "system", "content": "You are a helpful expert advisor."},
-                {"role": "user", "content": f"Context: {context}\n\nQuestion: {query}"}
-            ],
-            max_tokens=250
+                {"role": "system", "content": "You are a helpful reasoning assistant."},
+                {"role": "user", "content": prompt}
+            ]
         )
-        return response["choices"][0]["message"]["content"].strip()
+        return response['choices'][0]['message']['content']
     except Exception as e:
         return f"⚠️ GPT error: {str(e)}"