Update inference.py

inference.py

@@ -1,54 +1,51 @@
+# inference.py
+
 import torch
-import openai
-import os
+import torch.nn.functional as F
 from transformers import AutoTokenizer
-from evo_model import EvoTransformerV22
-from rag_utils import extract_text_from_file
-from search_utils import web_search
+from model import EvoTransformerV22
+import openai
 
+# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+
+# Load EvoTransformer model
 model = EvoTransformerV22()
 model.load_state_dict(torch.load("evo_hellaswag.pt", map_location="cpu"))
 model.eval()
 
-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()
-
-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(encoded["input_ids"])
-    probs = torch.softmax(logits, dim=1).squeeze()
-    pred_index = torch.argmax(probs).item()
-    confidence = probs[pred_index].item()
-
-    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(question, context, options):
-    openai.api_key = os.getenv("OPENAI_API_KEY", "")
-    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 reasoning assistant."},
-                {"role": "user", "content": prompt}
-            ]
-        )
-        return response['choices'][0]['message']['content']
-    except Exception as e:
-        return f"⚠️ GPT error: {str(e)}"
+# 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)
+
+    better = option1 if score1 > score2 else option2
+    confidence = max(score1, score2)
+
+    return better, confidence, score1, score2
+
+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
+    )
+    return response.choices[0].message.content.strip()