Update backend.py

backend.py

@@ -1,4 +1,4 @@
-# backend.py — ENHANCED WITH AGENT CONVERSATION LOGS + RAM/VRAM MONITORING
+# backend.py — FINAL VERSION — PHASE-3 MINI INSTRUCT + CACHING + LOGGING
 import sqlite3
 import os
 import json
@@ -152,18 +152,18 @@ def get_vram_usage_gb():
     except:
         return 0.0
 
-# ------------------------------ MODEL LOADING ------------------------------
+# ------------------------------ MODEL LOADING (PER ROLE) ------------------------------
 MODEL_REGISTRY = {
     "understander": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
     "architect": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
-    "ceo": "google/gemma-2-2b-it",
+    "ceo": "microsoft/Phi-3-mini-4k-instruct",
     "manager": "microsoft/Phi-3-mini-4k-instruct",
     "worker": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
     "reviewer": "microsoft/Phi-3-mini-4k-instruct",
     "editor": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
     "tester": "microsoft/Phi-3-mini-4k-instruct",
     "publisher": "Qwen/Qwen2.5-Coder-0.6B-Instruct",
-    "final_ceo": "google/gemma-2-2b-it",
+    "final_ceo": "microsoft/Phi-3-mini-4k-instruct",
 }
 
 _MODEL_CACHE = {}
@@ -177,7 +177,8 @@ def load_model(model_name):
         model_name,
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
         device_map="auto",
-        trust_remote_code=True
+        trust_remote_code=True,
+        attn_implementation="eager"  # ✅ Avoids flash-attn requirement — safe on all hardware
     )
     _MODEL_CACHE[model_name] = (tokenizer, model)
     return tokenizer, model
@@ -418,51 +419,35 @@ def zip_project(user_id, project_name: str):
                 zf.write(full_path, arcname)
     return zip_path
 
-# ------------------------------ LLM GENERATOR (WITH LOGGING) ------------------------------
+# ------------------------------ LLM GENERATOR (WITH CHAT TEMPLATE FOR PHI-3) ------------------------------
 def generate_with_model(role: str, prompt: str, context: dict = {}) -> str:
     try:
         model_name = MODEL_REGISTRY[role]
         tokenizer, model = load_model(model_name)
 
-        # Format message based on role
-        if role == "understander":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["understander"]},
-                        {"role": "user", "content": prompt}]
-        elif role == "architect":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["architect"]},
-                        {"role": "user", "content": prompt}]
-        elif role == "ceo":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["ceo"]},
-                        {"role": "user", "content": prompt}]
-        elif role == "manager":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["manager"]},
-                        {"role": "user", "content": prompt}]
-        elif role == "worker":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["worker"].format(file=context.get("file"), instructions=context.get("instructions"))},
-                        {"role": "user", "content": ""}]
-        elif role == "reviewer":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["reviewer"]},
-                        {"role": "user", "content": prompt}]
-        elif role == "editor":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["editor"]},
-                        {"role": "user", "content": prompt}]
-        elif role == "tester":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["tester"]},
-                        {"role": "user", "content": prompt}]
-        elif role == "publisher":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["publisher"]},
-                        {"role": "user", "content": prompt}]
-        elif role == "final_ceo":
-            messages = [{"role": "system", "content": ROLE_PROMPTS["final_ceo"]},
-                        {"role": "user", "content": prompt}]
+        # Format message using Phi-3 or Qwen chat template
+        if "Phi-3" in model_name:
+            # Phi-3 format
+            messages = [
+                {"role": "system", "content": ROLE_PROMPTS[role]},
+                {"role": "user", "content": prompt}
+            ]
+            text = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
         else:
-            messages = [{"role": "user", "content": prompt}]
-
-        text = tokenizer.apply_chat_template(
-            messages,
-            tokenize=False,
-            add_generation_prompt=True
-        )
+            # Qwen format
+            messages = [
+                {"role": "system", "content": ROLE_PROMPTS[role]},
+                {"role": "user", "content": prompt}
+            ]
+            text = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
 
         inputs = tokenizer(text, return_tensors="pt").to(model.device)