Update app.py

app.py

@@ -1,4 +1,4 @@
-# Multi-agent UI generator with token handling and model load debug support
+# Multi-agent UI generator with lazy model loading to reduce memory usage
 
 import streamlit as st
 import time
@@ -6,7 +6,7 @@ import base64
 from typing import Dict, List, TypedDict
 from langgraph.graph import StateGraph, END
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
-from peft import PeftModel, PeftConfig
+from peft import PeftModel
 import torch
 import os
 
@@ -14,7 +14,6 @@ st.set_page_config(page_title="Multi-Agent Collaboration", layout="wide")
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
-# Agent model loading config
 AGENT_MODEL_CONFIG = {
     "product_manager": {
         "base": "google/gemma-1.1-7b-it",
@@ -29,7 +28,7 @@ AGENT_MODEL_CONFIG = {
         "adapter": "spandana30/software-architect-cohere"
     },
     "software_engineer": {
-        "base": "codellama/CodeLlama-7b-Instruct-hf",
+        "base": "codellama/CodeLLaMA-7b-Instruct-hf",
         "adapter": "spandana30/software-engineer-codellama"
     },
     "qa": {
@@ -40,28 +39,20 @@ AGENT_MODEL_CONFIG = {
 
 @st.cache_resource
 
-def load_agent_model(base_id, adapter_id):
+def get_text_pipeline(role: str):
+    cfg = AGENT_MODEL_CONFIG[role]
     try:
-        st.write(f"🔍 Loading: {base_id} with adapter {adapter_id}")
-        st.write(f"🔐 Using token: {'Yes' if HF_TOKEN else 'No'}")
-
+        st.write(f"🔍 Loading model for {role}: {cfg['base']} + {cfg['adapter']}")
         base_model = AutoModelForCausalLM.from_pretrained(
-            base_id, torch_dtype=torch.float16, device_map="auto", token=HF_TOKEN
+            cfg["base"], torch_dtype=torch.float16, device_map="auto", token=HF_TOKEN
         )
-        model = PeftModel.from_pretrained(base_model, adapter_id, token=HF_TOKEN)
-        tokenizer = AutoTokenizer.from_pretrained(adapter_id, token=HF_TOKEN)
-
+        model = PeftModel.from_pretrained(base_model, cfg["adapter"], token=HF_TOKEN)
+        tokenizer = AutoTokenizer.from_pretrained(cfg["adapter"], token=HF_TOKEN)
         return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=1024)
-
     except Exception as e:
-        st.error(f"❌ Model load failed!\nBase: {base_id}\nAdapter: {adapter_id}\nError: {e}")
+        st.error(f"❌ Failed to load model for {role}\nError: {e}")
         raise
 
-AGENT_PIPELINES = {
-    role: load_agent_model(cfg["base"], cfg["adapter"])
-    for role, cfg in AGENT_MODEL_CONFIG.items()
-}
-
 class AgentState(TypedDict):
     messages: List[Dict[str, str]]
     user_request: str
@@ -75,8 +66,8 @@ class AgentState(TypedDict):
     timings: Dict[str, float]
 
 def run_pipeline(role: str, prompt: str):
-    response = AGENT_PIPELINES[role](prompt, do_sample=False)[0]['generated_text']
-    return response.strip()
+    pipe = get_text_pipeline(role)
+    return pipe(prompt, do_sample=False)[0]["generated_text"].strip()
 
 PROMPTS = {
     "product_manager": """You're a Product Manager. Refine and clarify this request:
@@ -189,4 +180,4 @@ def main():
                 st.write(f"🧩 {stage.replace('_', ' ').title()} Time: {final_state['timings'].get(stage, 0):.2f}s")
 
 if __name__ == "__main__":
-    main()
+    main()