Update app.py

app.py

@@ -1,16 +1,52 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
 import spaces  # Import the spaces library
 
-# Model IDs from Hugging Face Hub (now 1.5B, 7B, and 14B)
+# Model IDs from Hugging Face Hub (Fixed to 7B and 32B Unsloth)
 model_ids = {
-    "1.5B": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
     "7B": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
-    "14B": "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B", # Added 14B back
+    "32B-Unsloth": "unsloth/DeepSeek-R1-Distill-Qwen-32B-bnb-4bit", # Unsloth 32B model
 }
 
-# Revised Default Prompts (as defined above)
+models = {}  # Keep models as a dictionary, but initially empty
+tokenizers = {} # Keep tokenizers as a dictionary, initially empty
+
+# BitsAndBytesConfig for 4-bit quantization (for the 32B model)
+bnb_config_4bit = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16, # Or torch.float16 if needed
+)
+
+
+def get_model_and_tokenizer(size): # Function to load model on demand
+    if size not in models: # Load only if not already loaded
+        model_id = model_ids[size]
+        print(f"Loading {size} model: {model_id} on demand")
+        tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+        if size == "32B-Unsloth": # Apply 4-bit config for 32B model
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                quantization_config=bnb_config_4bit,
+                torch_dtype=torch.bfloat16, # Or torch.float16 if needed
+                device_map='auto',
+                trust_remote_code=True
+            )
+        else: # 7B model - standard loading
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                torch_dtype=torch.bfloat16, # Or torch.float16 if needed
+                device_map='auto',
+                trust_remote_code=True
+            )
+        models[size] = model
+        tokenizers[size] = tokenizer
+        print(f"Loaded {size} model on demand.")
+    return models[size], tokenizers[size]
+
+
+# Revised Default Prompts (as defined previously - these are still good)
 default_prompt_1_5b = """**Code Analysis Task**
 As a Senior Code Analyst, analyze this programming problem:
 
@@ -23,8 +59,8 @@ As a Senior Code Analyst, analyze this programming problem:
 **Analysis Required:**
 1. Briefly break down the problem, including key constraints and edge cases.
 2. Suggest 2-3 potential approach options (algorithms/data structures).
-3. Recommend ONE primary strategy and briefly justify your choice. 
-4. Provide a very brief initial pseudocode sketch of the core logic.""" 
+3. Recommend ONE primary strategy and briefly justify your choice.
+4. Provide a very brief initial pseudocode sketch of the core logic."""
 
 
 default_prompt_7b = """**Code Implementation Task**
@@ -43,25 +79,6 @@ As a Principal Software Engineer, provide production-ready Streamlit/Python code
     """
 
 
-# Function to load model and tokenizer (same)
-def load_model_and_tokenizer(model_id):
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.bfloat16, # Or torch.float16 if you prefer
-        device_map='auto', # Let accelerate decide (will use GPU when @spaces.GPU active)
-        trust_remote_code=True
-    )
-    return model, tokenizer
-
-# Load the selected models and tokenizers (now loads 1.5B, 7B, 14B)
-models = {}
-tokenizers = {}
-for size, model_id in model_ids.items():
-    print(f"Loading {size} model: {model_id}")
-    models[size], tokenizers[size] = load_model_and_tokenizer(model_id)
-    print(f"Loaded {size} model.")
-
 # --- Shared Memory Implementation --- (Same)
 shared_memory = []
 
@@ -84,62 +101,46 @@ def retrieve_from_memory(query, top_k=2):
     return relevant_memories[:top_k]
 
 
-# --- Swarm Agent Function with Model Swapping ---
+# --- Swarm Agent Function - Fixed Models (7B and 32B Unsloth) ---
 @spaces.GPU  # <----  GPU DECORATOR ADDED HERE!
-def swarm_agent_sequential_rag(user_prompt, prompt_1_5b_template, prompt_7b_template, final_model_size="7B", temperature=0.5, top_p=0.9, max_new_tokens=300): # Added final_model_size
+def swarm_agent_sequential_rag(user_prompt, prompt_1_5b_template, prompt_7b_template, temperature=0.5, top_p=0.9, max_new_tokens=300): # Removed final_model_size
     global shared_memory
     shared_memory = [] # Clear memory for each new request
 
-    print(f"\n--- Swarm Agent Processing with Shared Memory (RAG) - GPU ACCELERATED - Final Model: {final_model_size} ---") # Updated message
+    print(f"\n--- Swarm Agent Processing with Shared Memory (RAG) - GPU ACCELERATED - Final Model: 32B Unsloth ---") # Updated message
 
-    # 1.5B Model - Brainstorming/Initial Draft (same logic)
-    print("\n[1.5B Model - Brainstorming] - GPU Accelerated")
-    retrieved_memory_1_5b = retrieve_from_memory(user_prompt)
-    context_1_5b = "\n".join([f"- {mem}" for mem in retrieved_memory_1_5b]) if retrieved_memory_1_5b else "No relevant context found in memory."
+    # 7B Model - Brainstorming/Initial Draft (Lazy Load and get model)
+    print("\n[7B Model - Brainstorming] - GPU Accelerated") # Now 7B is brainstorming
+    model_7b, tokenizer_7b = get_model_and_tokenizer("7B") # Lazy load 7B
+    retrieved_memory_7b = retrieve_from_memory(user_prompt)
+    context_7b = "\n".join([f"- {mem}" for mem in retrieved_memory_7b]) if retrieved_memory_7b else "No relevant context found in memory."
 
-    # Use user-provided prompt template for 1.5B model
-    prompt_1_5b = prompt_1_5b_template.format(user_prompt=user_prompt, context_1_5b=context_1_5b)
+    # Use user-provided prompt template for 7B model (as brainstorming model now)
+    prompt_7b_brainstorm = prompt_1_5b_template.format(user_prompt=user_prompt, context_1_5b=context_7b) # Reusing 1.5B template - adjust if needed
 
-    input_ids_1_5b = tokenizers["1.5B"].encode(prompt_1_5b, return_tensors="pt").to(models["1.5B"].device)
-    output_1_5b = models["1.5B"].generate(
-        input_ids_1_5b,
+    input_ids_7b = tokenizer_7b.encode(prompt_7b_brainstorm, return_tensors="pt").to(model_7b.device)
+    output_7b = model_7b.generate(
+        input_ids_7b,
         max_new_tokens=max_new_tokens, # Use user-defined max_new_tokens
         temperature=temperature, # Use user-defined temperature
         top_p=top_p,           # Use user-defined top_p
         do_sample=True
     )
-    response_1_5b = tokenizers["1.5B"].decode(output_1_5b[0], skip_special_tokens=True)
-    print(f"1.5B Response:\n{response_1_5b}")
-    store_in_memory(f"1.5B Model Initial Response: {response_1_5b[:200]}...")
-
-    # Final Stage Model Selection (7B or 14B)
-    if final_model_size == "7B":
-        final_model = models["7B"]
-        final_tokenizer = tokenizers["7B"]
-        print("\n[7B Model - Final Code Generation] - GPU Accelerated") # Model-specific message
-        model_stage_name = "7B Model - Final Code"
-        final_max_new_tokens = max_new_tokens + 100 # Slightly more tokens for 7B
-
-    elif final_model_size == "14B":
-        final_model = models["14B"]
-        final_tokenizer = tokenizers["14B"]
-        print("\n[14B Model - Final Code Generation] - GPU Accelerated") # Model-specific message
-        model_stage_name = "14B Model - Final Code"
-        final_max_new_tokens = max_new_tokens + 200 # Even more tokens for 14B
-
-    else: # Default to 7B if selection is somehow invalid
-        final_model = models["7B"]
-        final_tokenizer = tokenizers["7B"]
-        print("\n[7B Model - Final Code Generation] - GPU Accelerated (Default)")
-        model_stage_name = "7B Model - Final Code (Default)"
-        final_max_new_tokens = max_new_tokens + 100
-
-
-    retrieved_memory_final = retrieve_from_memory(response_1_5b)
+    response_7b = tokenizer_7b.decode(output_7b[0], skip_special_tokens=True)
+    print(f"7B Response (Brainstorming):\n{response_7b}") # Updated message
+    store_in_memory(f"7B Model Initial Response: {response_7b[:200]}...")
+
+    # 32B Unsloth Model - Final Code Generation (Lazy Load and get model)
+    final_model, final_tokenizer = get_model_and_tokenizer("32B-Unsloth") # Lazy load 32B Unsloth
+    print("\n[32B Unsloth Model - Final Code Generation] - GPU Accelerated") # Model-specific message
+    model_stage_name = "32B Unsloth Model - Final Code"
+    final_max_new_tokens = max_new_tokens + 200 # More tokens for 32B model
+
+    retrieved_memory_final = retrieve_from_memory(response_7b) # Memory from 7B brainstorm
     context_final = "\n".join([f"- {mem}" for mem in retrieved_memory_final]) if retrieved_memory_final else "No relevant context found in memory."
 
-    # Use user-provided prompt template for final model (currently using 7B prompt for both 7B and 14B for simplicity, you can create a separate 14B prompt if needed)
-    prompt_final = prompt_7b_template.format(response_1_5b=response_1_5b, context_7b=context_final) # Using prompt_7b_template for final stage for now
+    # Use user-provided prompt template for final model (using 7B template)
+    prompt_final = prompt_7b_template.format(response_1_5b=response_7b, context_7b=context_final) # Using prompt_7b_template for final stage
 
 
     input_ids_final = final_tokenizer.encode(prompt_final, return_tensors="pt").to(final_model.device)
@@ -157,14 +158,13 @@ def swarm_agent_sequential_rag(user_prompt, prompt_1_5b_template, prompt_7b_temp
     return response_final # Returns final model's response
 
 
-# --- Gradio ChatInterface --- (with Model Selection Dropdown)
-def gradio_interface(message, history, temp, top_p, max_tokens, prompt_1_5b_text, prompt_7b_text, final_model_selector): # Added final_model_selector
+# --- Gradio ChatInterface --- (No Model Selection Dropdown anymore)
+def gradio_interface(message, history, temp, top_p, max_tokens, prompt_1_5b_text, prompt_7b_text): # Removed final_model_selector
     # history is automatically managed by ChatInterface
     response = swarm_agent_sequential_rag(
         message,
         prompt_1_5b_template=prompt_1_5b_text, # Pass prompt templates
         prompt_7b_template=prompt_7b_text,
-        final_model_size=final_model_selector, # Pass model selection
         temperature=temp,
         top_p=top_p,
         max_new_tokens=int(max_tokens) # Ensure max_tokens is an integer
@@ -173,17 +173,16 @@ def gradio_interface(message, history, temp, top_p, max_tokens, prompt_1_5b_text
 
 iface = gr.ChatInterface( # Using ChatInterface now
     fn=gradio_interface,
-    # Define additional inputs for settings, prompts, and model selection
+    # Define additional inputs for settings and prompts (NO model dropdown)
     additional_inputs=[
         gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.5, label="Temperature"), # Lowered default temp to 0.5
         gr.Slider(minimum=0.01, maximum=1.0, step=0.05, value=0.9, label="Top P"),
         gr.Number(value=300, label="Max Tokens", precision=0), # Use Number for integer tokens
-        gr.Textbox(value=default_prompt_1_5b, lines=10, label="1.5B Model Prompt Template"), # Textbox for 1.5B prompt
-        gr.Textbox(value=default_prompt_7b, lines=10, label="7B Model Prompt Template"),   # Textbox for 7B prompt
-        gr.Dropdown(choices=["7B", "14B"], value="7B", label="Final Stage Model (7B or 14B)") # Model selection dropdown
+        gr.Textbox(value=default_prompt_1_5b, lines=10, label="Brainstorming Model Prompt Template (7B Model)"), # Updated label - 7B now brainstormer
+        gr.Textbox(value=default_prompt_7b, lines=10, label="Code Generation Prompt Template (32B Unsloth Model)"),   # Updated label - 32B is code generator
     ],
-    title="DeepSeek Agent Swarm Chat (ZeroGPU Demo - 2 Models + Model Swap)", # Updated title
-    description="Chat with a DeepSeek agent swarm (1.5B + 7B/14B selectable) with shared memory, adjustable settings, **customizable prompts, and model swapping!** **GPU accelerated using ZeroGPU!** (Requires Pro Space)", # Updated description
+    title="DeepSeek Agent Swarm Chat (ZeroGPU Demo - Fixed Models: 7B + 32B Unsloth)", # Updated title
+    description="Chat with a DeepSeek agent swarm (7B + 32B Unsloth) with shared memory, adjustable settings, **and customizable prompts!** **GPU accelerated using ZeroGPU!** (Requires Pro Space)", # Updated description
 )
 
 if __name__ == "__main__":