Update inference.py

inference.py

@@ -1,62 +1,53 @@
 import torch
-from transformers import AutoTokenizer
 from evo_model import EvoTransformer
+from transformers import AutoTokenizer
 from rag_utils import extract_text_from_file
-from search_utils import web_search_and_format
+from search_utils import web_search
 
-# Load Evo model and tokenizer
-model_path = "evo_hellaswag.pt"
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# Load Evo model
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 model = EvoTransformer()
-model.load_state_dict(torch.load(model_path, map_location=device))
-model.to(device)
+model.load_state_dict(torch.load("evo_hellaswag.pt", map_location="cpu"))
 model.eval()
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-
-def get_evo_response(query, context="", file=None, enable_search=True):
-    rag_context = ""
-
-    if file is not None:
-        rag_context += extract_text_from_file(file)
 
+def get_evo_response(query, context=None, enable_search=True):
+    search_snippets = ""
     if enable_search:
-        search_context = web_search_and_format(query)
-        rag_context += "\n" + search_context
+        snippets = web_search(query)
+        if snippets:
+            search_snippets = "\n".join(snippets)
 
-    full_context = f"{context}\n{rag_context}".strip()
+    full_context = f"{context or ''}\n\n{search_snippets}".strip()
+    input_1 = f"{query} Option 1"
+    input_2 = f"{query} Option 2"
 
-    # Define hypothetical options (can be more sophisticated later)
-    option1 = "Yes, take action."
-    option2 = "No, do not take action."
-
-    inputs = [
-        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)
+    inputs = tokenizer([input_1, input_2], padding=True, truncation=True, return_tensors="pt")
     with torch.no_grad():
-        logits = model(encoded["input_ids"]).squeeze(-1)
+        logits = model(inputs["input_ids"]).squeeze(-1)
         probs = torch.softmax(logits, dim=0)
-        best = torch.argmax(probs).item()
+        best_idx = torch.argmax(probs).item()
 
-    answer = option1 if best == 0 else option2
+    suggestion = f"Option {best_idx + 1}"
     reasoning = (
-        f"✅ Evo suggests: **{answer}**\n\n"
-        f"🧠 Confidence: {probs[best]:.2f}\n"
-        f"📖 Context used:\n{full_context[:1000]}..."  # limit to 1000 chars
+        f"Evo suggests: **{suggestion}** (Confidence: {probs[best_idx]:.2f})\n\n"
+        f"Context used:\n{full_context}"
     )
+    return suggestion, reasoning
 
-    return answer, reasoning
-
-def get_gpt_response(query, context=""):
+def get_gpt_response(query, context=None):
     import openai
-    openai.api_key = "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"  # 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()
+    openai.api_key = os.getenv("OPENAI_API_KEY", "")
+    context = context or "None"
+
+    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
+        )
+        return response["choices"][0]["message"]["content"].strip()
+    except Exception as e:
+        return f"⚠️ GPT error: {str(e)}"