Update inference.py

inference.py

@@ -5,11 +5,10 @@ from transformers import AutoTokenizer
 from evo_model import EvoTransformerV22
 from search_utils import web_search
 import openai
-import time
 import psutil
 import platform
 
-# 🔐 Load OpenAI API Key securely
+# 🔐 Load OpenAI API Key
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 # 📦 Constants
@@ -18,7 +17,7 @@ tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 model = None
 last_mod_time = 0
 
-# 🔁 Reload model if updated on disk
+# 🔁 Load or reload Evo model
 def load_model():
     global model, last_mod_time
     try:
@@ -28,86 +27,73 @@ def load_model():
             model.load_state_dict(torch.load(MODEL_PATH, map_location="cpu"))
             model.eval()
             last_mod_time = current_mod_time
-            print("🔁 Evo model reloaded.")
+            print("✅ Evo model loaded.")
     except Exception as e:
-        print(f"❌ Error loading Evo model: {e}")
+        print(f"❌ Error loading model: {e}")
         model = None
     return model
 
-# 🧠 Evo inference logic
-def get_evo_response(query, options, user_context=""):
+# 🧠 Evo response engine
+def evo_infer(question, option1, option2, user_context=""):
     model = load_model()
     if model is None:
-        return "Error", 0.0, "Model failed to load", ""
+        return "Evo failed", 0.0, "Model not loaded", ""
 
-    # Context = web + user
-    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]
+    # 🌐 Retrieve context
+    context_blobs = web_search(question) + ([user_context] if user_context else [])
+    context = "\n".join(context_blobs)
 
-    # Score each option
+    # ⛓ Format input pairs
+    inputs = [f"{question} [SEP] {opt} [CTX] {context}" for opt in [option1, option2]]
     scores = []
-    for pair in input_pairs:
+
+    for pair in inputs:
         encoded = tokenizer(pair, return_tensors="pt", padding="max_length", truncation=True, max_length=128)
         with torch.no_grad():
             logits = model(encoded["input_ids"])
             score = torch.sigmoid(logits).item()
             scores.append(score)
 
-    best_idx = int(scores[1] > scores[0])
-    return (
-        options[best_idx],                      # ✅ Evo's best answer
-        max(scores),                            # ✅ Confidence
-        f"{options[0]}: {scores[0]:.3f} vs {options[1]}: {scores[1]:.3f}",  # ✅ Reasoning trace
-        context_str                             # ✅ Used context
-    )
+    best = option2 if scores[1] > scores[0] else option1
+    trace = f"{option1}: {scores[0]:.3f} vs {option2}: {scores[1]:.3f}"
+
+    return best, max(scores), trace, context
 
-# 🔄 GPT-3.5 response
-def get_gpt_response(query, user_context=""):
+# 🔄 GPT backup
+def gpt_infer(question, user_context=""):
     try:
-        context_block = f"\n\nContext:\n{user_context}" if user_context else ""
+        block = f"\n\nContext:\n{user_context}" if user_context else ""
         response = openai.chat.completions.create(
             model="gpt-3.5-turbo",
-            messages=[{"role": "user", "content": query + context_block}],
+            messages=[{"role": "user", "content": question + block}],
             temperature=0.7,
         )
         return response.choices[0].message.content.strip()
     except Exception as e:
-        return f"⚠️ GPT error:\n\n{str(e)}"
-
-# 🎯 For EvoRAG app UI
-def evo_chat_predict(history, query, options):
-    context = "\n".join(history[-6:]) if history else ""
-    evo_ans, evo_score, evo_reason, evo_ctx = get_evo_response(query, options, context)
-    return {
-        "answer": evo_ans,
-        "confidence": round(evo_score, 3),
-        "reasoning": evo_reason,
-        "context_used": evo_ctx
-    }
+        return f"⚠️ GPT error: {str(e)}"
 
-# 📊 Evo architecture stats
+# 📊 Evo architecture info
 def get_model_config():
     return {
         "num_layers": 6,
         "num_heads": 8,
         "ffn_dim": 1024,
         "memory_enabled": True,
-        "param_count": sum(p.numel() for p in model.parameters() if p.requires_grad) if model else "N/A"
+        "total_params": sum(p.numel() for p in model.parameters()) if model else "N/A"
     }
 
-# 💻 Hardware and system stats
+# 💻 System stats
 def get_system_stats():
-    gpu_info = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None
-    memory = psutil.virtual_memory()
+    gpu = torch.cuda.get_device_properties(0) if torch.cuda.is_available() else None
+    mem = psutil.virtual_memory()
 
     return {
         "device": "GPU" if torch.cuda.is_available() else "CPU",
         "cpu_usage_percent": psutil.cpu_percent(),
-        "memory_used_gb": round(memory.used / (1024 ** 3), 2),
-        "memory_total_gb": round(memory.total / (1024 ** 3), 2),
-        "gpu_name": gpu_info.name if gpu_info else "N/A",
-        "gpu_memory_total_gb": round(gpu_info.total_memory / (1024 ** 3), 2) if gpu_info else "N/A",
-        "gpu_memory_used_gb": round(torch.cuda.memory_allocated() / (1024 ** 3), 2) if gpu_info else "N/A",
+        "memory_used_gb": round(mem.used / (1024 ** 3), 2),
+        "memory_total_gb": round(mem.total / (1024 ** 3), 2),
+        "gpu_name": gpu.name if gpu else "N/A",
+        "gpu_memory_total_gb": round(gpu.total_memory / (1024 ** 3), 2) if gpu else "N/A",
+        "gpu_memory_used_gb": round(torch.cuda.memory_allocated() / (1024 ** 3), 2) if gpu else "N/A",
         "platform": platform.platform()
     }