Update inference.py

inference.py

@@ -1,80 +1,63 @@
 import torch
 from evo_model import EvoTransformer
-from transformers import AutoTokenizer, pipeline
-from rag_utils import RAGRetriever, extract_text_from_file
-import os
+from transformers import AutoTokenizer
+from rag_utils import extract_text_from_file
+from search_utils import web_search
 
-# Load Evo model
-def load_evo_model(model_path="evo_hellaswag.pt", device=None):
+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_evo_model()
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+evo_model, device = load_model()
 
-# Load GPT-3.5 (via OpenAI API)
-import openai
-openai.api_key = os.getenv("OPENAI_API_KEY")
+def get_evo_response(query, file=None, enable_search=True):
+    context = ""
 
-# RAG Retriever
-retriever = RAGRetriever()
+    if file:
+        try:
+            context += extract_text_from_file(file)[:800]
+        except:
+            pass
 
-def get_context_from_file(file_obj):
-    file_path = file_obj.name
-    text = extract_text_from_file(file_path)
-    retriever.add_document(text)
-    return text
+    if enable_search:
+        search_snippets = web_search(query)
+        context += "\n".join(search_snippets)
 
-# Evo prediction
-def get_evo_response(prompt, file=None):
-    # Step 1: augment context if document is uploaded
-    context = ""
-    if file is not None:
-        context_list = retriever.retrieve(prompt)
-        context = "\n".join(context_list)
+    combined_prompt = f"{query}\nContext:\n{context}"
 
-    full_prompt = f"{prompt}\n{context}"
-
-    # Step 2: use Evo to predict
-    options = ["Yes, proceed with the action.", "No, maintain current strategy."]
-    inputs = [f"{full_prompt} {opt}" for opt in options]
+    inputs = [
+        f"{combined_prompt} Option 1:",
+        f"{combined_prompt} Option 2:",
+    ]
 
     encoded = tokenizer(inputs, padding=True, truncation=True, return_tensors="pt").to(device)
-
     with torch.no_grad():
-        logits = evo_model(encoded["input_ids"]).squeeze(-1)
-        probs = torch.softmax(logits, dim=0)
+        outputs = evo_model(encoded["input_ids"]).squeeze(-1)
+        probs = torch.softmax(outputs, dim=0)
         best = torch.argmax(probs).item()
 
-    return f"Evo suggests: {options[best]} (Confidence: {probs[best]:.2f})"
-
-# GPT-3.5 response
-def get_gpt_response(prompt, file=None):
-    context = ""
-    if file is not None:
-        context_list = retriever.retrieve(prompt)
-        context = "\n".join(context_list)
-
-    full_prompt = (
-        f"Question: {prompt}\n"
-        f"Relevant Context:\n{context}\n"
-        f"Answer like a financial advisor."
-    )
+    return f"Option {best + 1} with {probs[best]:.2%} confidence.\n\nReasoning:\n{inputs[best]}"
 
-    response = openai.ChatCompletion.create(
-        model="gpt-3.5-turbo",
-        messages=[
-            {"role": "user", "content": full_prompt}
-        ],
-        temperature=0.4,
-    )
+def get_gpt_response(query, context=""):
+    import openai
+    openai.api_key = os.getenv("OPENAI_API_KEY")
 
-    return response.choices[0].message.content.strip()
+    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)}"