Update app.py

app.py

@@ -1,4 +1,4 @@
-# Updated multi-agent UI generation system with all agents using public base models (Gemma used for Product Manager too)
+# Multi-agent UI generator with Hugging Face token-based gated model support
 
 import streamlit as st
 import time
@@ -8,9 +8,12 @@ from langgraph.graph import StateGraph, END
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel, PeftConfig
 import torch
+import os
 
 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": {
@@ -22,7 +25,7 @@ AGENT_MODEL_CONFIG = {
         "adapter": "spandana30/project-manager-gemma"
     },
     "software_architect": {
-        "base": "cohere/command-r",  # update if you have a local base version
+        "base": "cohere/command-r",
         "adapter": "spandana30/software-architect-cohere"
     },
     "software_engineer": {
@@ -39,10 +42,10 @@ AGENT_MODEL_CONFIG = {
 
 def load_agent_model(base_id, adapter_id):
     base_model = AutoModelForCausalLM.from_pretrained(
-        base_id, torch_dtype=torch.float16, device_map="auto"
+        base_id, torch_dtype=torch.float16, device_map="auto", token=HF_TOKEN
     )
-    model = PeftModel.from_pretrained(base_model, adapter_id)
-    tokenizer = AutoTokenizer.from_pretrained(adapter_id)
+    model = PeftModel.from_pretrained(base_model, adapter_id, token=HF_TOKEN)
+    tokenizer = AutoTokenizer.from_pretrained(adapter_id, token=HF_TOKEN)
     return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=1024)
 
 AGENT_PIPELINES = {
@@ -177,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()