Update inference.py

inference.py

@@ -2,6 +2,7 @@ import torch
 from transformers import AutoTokenizer
 from evo_model import EvoTransformerV22
 from retriever import retrieve
+from websearch import web_search
 from openai import OpenAI
 import os
 
@@ -15,17 +16,20 @@ evo_model.eval()
 # --- Load Tokenizer ---
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
-# --- EvoRAG Inference ---
+# --- EvoRAG+ Inference ---
 def evo_rag_response(query):
-    # Step 1: retrieve document chunks
+    # Step 1: get document context (from uploaded file)
     rag_context = retrieve(query)
 
-    # Step 2: combine query with retrieved context
-    combined = query + " " + rag_context
+    # Step 2: get online info (search/web)
+    web_context = web_search(query)
+
+    # Step 3: combine all into one input
+    combined = query + "\n\n" + rag_context + "\n\n" + web_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
+    # Step 4: Evo prediction
     with torch.no_grad():
         logits = evo_model(input_ids)
         pred = int(torch.sigmoid(logits).item() > 0.5)
@@ -33,7 +37,7 @@ def evo_rag_response(query):
     return f"Evo suggests: Option {pred + 1}"
 
 # --- GPT-3.5 Inference (OpenAI >= 1.0.0) ---
-openai_api_key = os.environ.get("OPENAI_API_KEY", "sk-...")  # Replace or use HF secret
+openai_api_key = os.environ.get("OPENAI_API_KEY", "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")  # Replace or use HF secret
 client = OpenAI(api_key=openai_api_key)
 
 def get_gpt_response(query, context):