Update inference.py

inference.py

@@ -1,33 +1,39 @@
 import torch
 from transformers import AutoTokenizer
 from evo_model import EvoTransformerV22
+from retriever import retrieve
 from openai import OpenAI
 import os
 
-# Load Evo model
+# --- Load Evo Model ---
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 evo_model = EvoTransformerV22()
 evo_model.load_state_dict(torch.load("trained_model_evo_hellaswag.pt", map_location=device))
 evo_model.to(device)
 evo_model.eval()
 
-# Load tokenizer
+# --- Load Tokenizer ---
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
-# 🧠 Evo logic (binary classification with sigmoid)
-def get_evo_response(query, context):
-    combined = query + " " + context
+# --- EvoRAG Inference ---
+def evo_rag_response(query):
+    # Step 1: retrieve document chunks
+    rag_context = retrieve(query)
+
+    # Step 2: combine query with retrieved context
+    combined = query + " " + rag_context
     inputs = tokenizer(combined, return_tensors="pt", truncation=True, padding="max_length", max_length=128)
     input_ids = inputs["input_ids"].to(device)
 
+    # Step 3: predict using Evo
     with torch.no_grad():
         logits = evo_model(input_ids)
         pred = int(torch.sigmoid(logits).item() > 0.5)
 
     return f"Evo suggests: Option {pred + 1}"
 
-# 🤖 GPT-3.5 comparison using openai>=1.0.0
-openai_api_key = os.environ.get("OPENAI_API_KEY", "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")  # Replace with real key or set via HF secrets
+# --- GPT-3.5 Inference (OpenAI >= 1.0.0) ---
+openai_api_key = os.environ.get("OPENAI_API_KEY", "sk-...")  # Replace or use HF secret
 client = OpenAI(api_key=openai_api_key)
 
 def get_gpt_response(query, context):
@@ -35,9 +41,7 @@ def get_gpt_response(query, context):
         prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
         response = client.chat.completions.create(
             model="gpt-3.5-turbo",
-            messages=[
-                {"role": "user", "content": prompt}
-            ],
+            messages=[{"role": "user", "content": prompt}],
             temperature=0.3
         )
         return response.choices[0].message.content.strip()