Update app.py

app.py

@@ -3,11 +3,10 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import spaces  # Import the spaces library
 
-# Model IDs from Hugging Face Hub (same as before)
+# Model IDs from Hugging Face Hub (now only 1.5B and 7B)
 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"
 }
 
 # Function to load model and tokenizer (slightly adjusted device_map)
@@ -21,7 +20,7 @@ def load_model_and_tokenizer(model_id):
     )
     return model, tokenizer
 
-# Load all three models and tokenizers (loaded once at app startup - potentially on CPU initially)
+# Load the selected models and tokenizers
 models = {}
 tokenizers = {}
 for size, model_id in model_ids.items():
@@ -81,29 +80,22 @@ def swarm_agent_sequential_rag(user_prompt):
     print(f"7B Response:\n{response_7b}")
     store_in_memory(f"7B Model Elaborated Response: {response_7b[:200]}...")
 
-    # 14B Model - Final Reasoning and Refinement
-    print("\n[14B Model - Final Refinement] - GPU Accelerated") # Added GPU indication
-    retrieved_memory_14b = retrieve_from_memory(response_7b)
-    context_14b = "\n".join([f"- {mem}" for mem in retrieved_memory_14b]) if retrieved_memory_14b else "No relevant context found in memory."
-    prompt_14b = f"Context from Shared Memory:\n{context_14b}\n\nYou are a high-level reasoner and refiner.  Take the following elaborated response and refine it to be a final, well-reasoned, and polished answer, considering the context above. \n\nElaborated Response:\n{response_7b}\n\nFinal Answer:"
-    input_ids_14b = tokenizers["14B"].encode(prompt_14b, return_tensors="pt").to(models["14B"].device)
-    output_14b = models["14B"].generate(input_ids_14b, max_new_tokens=400, temperature=0.6, do_sample=True) # Reverted to original max_new_tokens
-    response_14b = tokenizers["14B"].decode(output_14b[0], skip_special_tokens=True)
-    print(f"14B Response (Final):\n{response_14b}")
+    # No 14B Model Stage anymore
 
-    return response_14b
+    return response_7b # Now returns the 7B model's response as final
 
 
-# --- Gradio Interface --- (Same as before)
+# --- Gradio Interface --- (Modified to reflect 2-model output)
 def gradio_interface(user_prompt):
-    return swarm_agent_sequential_rag(user_prompt)
+    final_response = swarm_agent_sequential_rag(user_prompt) # Get the final response (from 7B now)
+    return final_response
 
 iface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Textbox(lines=5, placeholder="Enter your task here..."),
     outputs=gr.Textbox(lines=10, placeholder="Agent Swarm Output will appear here..."),
-    title="DeepSeek Agent Swarm (ZeroGPU Demo)",
-    description="Agent swarm using DeepSeek-R1-Distill models (1.5B, 7B, 14B) with shared memory. **GPU accelerated using ZeroGPU!** (Requires Pro Space)", # Updated description
+    title="DeepSeek Agent Swarm (ZeroGPU Demo - 2 Models)", # Updated title
+    description="Agent swarm using DeepSeek-R1-Distill models (1.5B, 7B) with shared memory. **GPU accelerated using ZeroGPU!** (Requires Pro Space)", # Updated description
 )
 
 if __name__ == "__main__":