Update inference.py

inference.py

@@ -1,51 +1,47 @@
-# inference.py
-
 import torch
-import torch.nn.functional as F
 from transformers import AutoTokenizer
-from evo_model import EvoTransformerV22
+from model import EvoTransformerV22
+from search_utils import web_search
 import openai
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
-
-# Load EvoTransformer model
+# Load Evo model and tokenizer
 model = EvoTransformerV22()
 model.load_state_dict(torch.load("evo_hellaswag.pt", map_location="cpu"))
 model.eval()
 
-# GPT-3.5 API
-openai.api_key = "sk-..."  # Replace with your key
-
-def get_evo_response(question, option1, option2):
-    pair1 = f"{question} {option1}"
-    pair2 = f"{question} {option2}"
-
-    def score(pair):
-        encoded = tokenizer(pair, return_tensors="pt", padding=True, truncation=True, max_length=128)
-        with torch.no_grad():
-            logits = model(encoded["input_ids"])
-            prob = torch.sigmoid(logits).item()
-        return prob
-
-    score1 = score(pair1)
-    score2 = score(pair2)
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
-    better = option1 if score1 > score2 else option2
-    confidence = max(score1, score2)
+# GPT Setup
+openai.api_key = "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"  # 🔑 Set your actual key securely
 
-    return better, confidence, score1, score2
+def get_evo_response(query, options, user_context=""):
+    context_texts = web_search(query) + ([user_context] if user_context else [])
+    context_str = "\n".join(context_texts)
+    input_pairs = [f"{query} [SEP] {opt} [CTX] {context_str}" for opt in options]
 
-def get_gpt_response(question, option1, option2):
-    prompt = (
-        f"Question: {question}\n"
-        f"Option 1: {option1}\n"
-        f"Option 2: {option2}\n"
-        f"Which option makes more sense and why?"
-    )
-    response = openai.ChatCompletion.create(
-        model="gpt-3.5-turbo",
-        messages=[{"role": "user", "content": prompt}],
-        temperature=0.7
+    scores = []
+    for pair in input_pairs:
+        encoded = tokenizer(pair, return_tensors="pt", truncation=True, padding="max_length", max_length=128)
+        with torch.no_grad():
+            output = model(encoded["input_ids"])
+            score = torch.sigmoid(output).item()
+            scores.append(score)
+
+    best_idx = int(scores[1] > scores[0])
+    return (
+        options[best_idx],
+        f"{options[0]}: {scores[0]:.3f} vs {options[1]}: {scores[1]:.3f}",
+        max(scores),
+        context_str
     )
-    return response.choices[0].message.content.strip()
+
+def get_gpt_response(query, user_context=""):
+    try:
+        context_block = f"\n\nContext:\n{user_context}" if user_context else ""
+        completion = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": query + context_block}]
+        )
+        return completion.choices[0].message.content.strip()
+    except Exception as e:
+        return f"⚠️ GPT error: {str(e)}"