Update app.py

app.py

@@ -3,13 +3,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import spaces  # Import the spaces library
 
-# Model IDs from Hugging Face Hub (now only 1.5B and 7B)
+# Model IDs from Hugging Face Hub (now 1.5B, 7B, and 14B)
 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
 }
 
-# Revised Default Prompts
+# Revised Default Prompts (as defined above)
 default_prompt_1_5b = """**Code Analysis Task**
 As a Senior Code Analyst, analyze this programming problem:
 
@@ -22,12 +23,12 @@ 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 a primary strategy and explain your reasoning concisely.
-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**
-As a Principal Software Engineer, develop a solution based on this analysis:
+As a Principal Software Engineer, provide production-ready Streamlit/Python code based on this analysis:
 
 **Initial Analysis:**
 {response_1_5b}
@@ -35,11 +36,11 @@ As a Principal Software Engineer, develop a solution based on this analysis:
 **Relevant Context:**
 {context_7b}
 
-**Solution Development Requirements:**
-1. Present an optimized solution approach, justifying your algorithm choices.
-2. Provide production-grade code in [Python/JS/etc.] (infer language). Include error handling and comments.
-3. Outline a testing plan with key test cases.
-4. Briefly suggest optimization opportunities and debugging tips."""
+**Code Requirements:**
+1.  Generate concise, production-grade Python code for a Streamlit app.
+2.  Include necessary imports, UI elements, and basic functionality.
+3.  Add comments for clarity.
+    """
 
 
 # Function to load model and tokenizer (same)
@@ -53,7 +54,7 @@ def load_model_and_tokenizer(model_id):
     )
     return model, tokenizer
 
-# Load the selected models and tokenizers (same)
+# Load the selected models and tokenizers (now loads 1.5B, 7B, 14B)
 models = {}
 tokenizers = {}
 for size, model_id in model_ids.items():
@@ -83,16 +84,16 @@ def retrieve_from_memory(query, top_k=2):
     return relevant_memories[:top_k]
 
 
-# --- Swarm Agent Function with Shared Memory (RAG) - DECORATED with @spaces.GPU ---
+# --- Swarm Agent Function with Model Swapping ---
 @spaces.GPU  # <----  GPU DECORATOR ADDED HERE!
-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): # Lowered default temperature
+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
     global shared_memory
     shared_memory = [] # Clear memory for each new request
 
-    print("\n--- Swarm Agent Processing with Shared Memory (RAG) - GPU ACCELERATED ---") # Updated message
+    print(f"\n--- Swarm Agent Processing with Shared Memory (RAG) - GPU ACCELERATED - Final Model: {final_model_size} ---") # Updated message
 
     # 1.5B Model - Brainstorming/Initial Draft (same logic)
-    print("\n[1.5B Model - Brainstorming] - GPU Accelerated") # Added GPU indication
+    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."
 
@@ -111,37 +112,59 @@ def swarm_agent_sequential_rag(user_prompt, prompt_1_5b_template, prompt_7b_temp
     print(f"1.5B Response:\n{response_1_5b}")
     store_in_memory(f"1.5B Model Initial Response: {response_1_5b[:200]}...")
 
-    # 7B Model - Elaboration and Detail (same logic)
-    print("\n[7B Model - Elaboration] - GPU Accelerated") # Added GPU indication
-    retrieved_memory_7b = retrieve_from_memory(response_1_5b)
-    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 7B model
-    prompt_7b = prompt_7b_template.format(response_1_5b=response_1_5b, context_7b=context_7b)
-
-
-    input_ids_7b = tokenizers["7B"].encode(prompt_7b, return_tensors="pt").to(models["7B"].device)
-    output_7b = models["7B"].generate(
-        input_ids_7b,
-        max_new_tokens=max_new_tokens + 100, # Slightly more tokens for 7B
-        temperature=temperature, # Use user-defined temperature
-        top_p=top_p,           # Use user-defined top_p
+    # 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)
+    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
+
+
+    input_ids_final = final_tokenizer.encode(prompt_final, return_tensors="pt").to(final_model.device)
+    output_final = final_model.generate(
+        input_ids_final,
+        max_new_tokens=final_max_new_tokens,
+        temperature=temperature,
+        top_p=top_p,
         do_sample=True
     )
-    response_7b = tokenizers["7B"].decode(output_7b[0], skip_special_tokens=True)
-    print(f"7B Response:\n{response_7b}")
-    store_in_memory(f"7B Model Elaborated Response: {response_7b[:200]}...")
+    response_final = final_tokenizer.decode(output_final[0], skip_special_tokens=True)
+    print(f"{model_stage_name} Response:\n{response_final}")
+    store_in_memory(f"{model_stage_name} Response: {response_final[:200]}...")
 
-    return response_7b # Now returns the 7B model's response as final
+    return response_final # Returns final model's response
 
 
-# --- Gradio ChatInterface --- (same interface definition)
-def gradio_interface(message, history, temp, top_p, max_tokens, prompt_1_5b_text, prompt_7b_text): # Accept prompt textboxes
+# --- 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
     # 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
@@ -150,16 +173,17 @@ 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 and prompts
+    # Define additional inputs for settings, prompts, and model selection
     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
     ],
-    title="DeepSeek Agent Swarm Chat (ZeroGPU Demo - 2 Models) - PROMPT CUSTOMIZATION", # Updated title
-    description="Chat with a DeepSeek agent swarm (1.5B, 7B) with shared memory, adjustable settings, **and customizable prompts!** **GPU accelerated using ZeroGPU!** (Requires Pro Space)", # Updated description
+    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
 )
 
 if __name__ == "__main__":