Update inference.py

inference.py

@@ -1,63 +1,62 @@
 import torch
-from evo_model import EvoTransformer
 from transformers import AutoTokenizer
+from evo_model import EvoTransformer
 from rag_utils import extract_text_from_file
-from search_utils import web_search
-
+from search_utils import web_search_and_format
+
+# Load Evo model and tokenizer
+model_path = "evo_hellaswag.pt"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = EvoTransformer()
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.to(device)
+model.eval()
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
-def load_model(model_path="evo_hellaswag.pt", device=None):
-    if device is None:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = EvoTransformer()
-    model.load_state_dict(torch.load(model_path, map_location=device))
-    model.to(device)
-    model.eval()
-    return model, device
-
-evo_model, device = load_model()
+def get_evo_response(query, context="", file=None, enable_search=True):
+    rag_context = ""
 
-def get_evo_response(query, file=None, enable_search=True):
-    context = ""
-
-    if file:
-        try:
-            context += extract_text_from_file(file)[:800]
-        except:
-            pass
+    if file is not None:
+        rag_context += extract_text_from_file(file)
 
     if enable_search:
-        search_snippets = web_search(query)
-        context += "\n".join(search_snippets)
+        search_context = web_search_and_format(query)
+        rag_context += "\n" + search_context
+
+    full_context = f"{context}\n{rag_context}".strip()
 
-    combined_prompt = f"{query}\nContext:\n{context}"
+    # Define hypothetical options (can be more sophisticated later)
+    option1 = "Yes, take action."
+    option2 = "No, do not take action."
 
     inputs = [
-        f"{combined_prompt} Option 1:",
-        f"{combined_prompt} Option 2:",
+        f"Q: {query} Context: {full_context} A: {option1}",
+        f"Q: {query} Context: {full_context} A: {option2}",
     ]
 
     encoded = tokenizer(inputs, padding=True, truncation=True, return_tensors="pt").to(device)
     with torch.no_grad():
-        outputs = evo_model(encoded["input_ids"]).squeeze(-1)
-        probs = torch.softmax(outputs, dim=0)
+        logits = model(encoded["input_ids"]).squeeze(-1)
+        probs = torch.softmax(logits, dim=0)
         best = torch.argmax(probs).item()
 
-    return f"Option {best + 1} with {probs[best]:.2%} confidence.\n\nReasoning:\n{inputs[best]}"
+    answer = option1 if best == 0 else option2
+    reasoning = (
+        f"✅ Evo suggests: **{answer}**\n\n"
+        f"🧠 Confidence: {probs[best]:.2f}\n"
+        f"📖 Context used:\n{full_context[:1000]}..."  # limit to 1000 chars
+    )
+
+    return answer, reasoning
 
 def get_gpt_response(query, context=""):
     import openai
-    openai.api_key = os.getenv("OPENAI_API_KEY")
-
-    prompt = f"{query}\nContext:\n{context}\nGive a thoughtful recommendation with reasons."
-
-    try:
-        response = openai.ChatCompletion.create(
-            model="gpt-3.5-turbo",
-            messages=[{"role": "user", "content": prompt}],
-            max_tokens=300,
-            temperature=0.7,
-        )
-        return response.choices[0].message.content.strip()
-    except Exception as e:
-        return f"Error: {str(e)}"
+    openai.api_key = "sk-..."  # Make sure to secure this
+
+    prompt = f"Q: {query}\nContext: {context}\nA:"
+    response = openai.ChatCompletion.create(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": prompt}],
+        temperature=0.3
+    )
+    return response["choices"][0]["message"]["content"].strip()