Update app.py
Browse files
app.py
CHANGED
@@ -1,7 +1,110 @@
|
|
1 |
import gradio as gr
|
|
|
|
|
|
|
2 |
|
3 |
-
|
4 |
-
|
|
|
|
|
|
|
|
|
5 |
|
6 |
-
|
7 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
import gradio as gr
|
2 |
+
from transformers import AutoModelForCausalLM, AutoTokenizer
|
3 |
+
import torch
|
4 |
+
import spaces # Import the spaces library
|
5 |
|
6 |
+
# Model IDs from Hugging Face Hub (same as before)
|
7 |
+
model_ids = {
|
8 |
+
"1.5B": "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
|
9 |
+
"7B": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
|
10 |
+
"14B": "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
|
11 |
+
}
|
12 |
|
13 |
+
# Function to load model and tokenizer (slightly adjusted device_map)
|
14 |
+
def load_model_and_tokenizer(model_id):
|
15 |
+
tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
|
16 |
+
model = AutoModelForCausalLM.from_pretrained(
|
17 |
+
model_id,
|
18 |
+
torch_dtype=torch.bfloat16, # Or torch.float16 if you prefer
|
19 |
+
device_map='auto', # Let accelerate decide (will use GPU when @spaces.GPU active)
|
20 |
+
trust_remote_code=True
|
21 |
+
)
|
22 |
+
return model, tokenizer
|
23 |
+
|
24 |
+
# Load all three models and tokenizers (loaded once at app startup - potentially on CPU initially)
|
25 |
+
models = {}
|
26 |
+
tokenizers = {}
|
27 |
+
for size, model_id in model_ids.items():
|
28 |
+
print(f"Loading {size} model: {model_id}")
|
29 |
+
models[size], tokenizers[size] = load_model_and_tokenizer(model_id)
|
30 |
+
print(f"Loaded {size} model.")
|
31 |
+
|
32 |
+
# --- Shared Memory Implementation --- (Same as before)
|
33 |
+
shared_memory = []
|
34 |
+
|
35 |
+
def store_in_memory(memory_item):
|
36 |
+
shared_memory.append(memory_item)
|
37 |
+
print(f"\n[Memory Stored]: {memory_item[:50]}...")
|
38 |
+
|
39 |
+
def retrieve_from_memory(query, top_k=2):
|
40 |
+
relevant_memories = []
|
41 |
+
query_lower = query.lower()
|
42 |
+
for memory_item in shared_memory:
|
43 |
+
if query_lower in memory_item.lower():
|
44 |
+
relevant_memories.append(memory_item)
|
45 |
+
|
46 |
+
if not relevant_memories:
|
47 |
+
print("\n[Memory Retrieval]: No relevant memories found.")
|
48 |
+
return []
|
49 |
+
|
50 |
+
print(f"\n[Memory Retrieval]: Found {len(relevant_memories)} relevant memories.")
|
51 |
+
return relevant_memories[:top_k]
|
52 |
+
|
53 |
+
|
54 |
+
# --- Swarm Agent Function with Shared Memory (RAG) - DECORATED with @spaces.GPU ---
|
55 |
+
@spaces.GPU # <---- GPU DECORATOR ADDED HERE!
|
56 |
+
def swarm_agent_sequential_rag(user_prompt):
|
57 |
+
global shared_memory
|
58 |
+
shared_memory = [] # Clear memory for each new request
|
59 |
+
|
60 |
+
print("\n--- Swarm Agent Processing with Shared Memory (RAG) - GPU ACCELERATED ---") # Updated message
|
61 |
+
|
62 |
+
# 1.5B Model - Brainstorming/Initial Draft
|
63 |
+
print("\n[1.5B Model - Brainstorming] - GPU Accelerated") # Added GPU indication
|
64 |
+
retrieved_memory_1_5b = retrieve_from_memory(user_prompt)
|
65 |
+
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."
|
66 |
+
prompt_1_5b = f"Context from Shared Memory:\n{context_1_5b}\n\nYou are a quick idea generator. Generate an initial response to the following user request, considering the context above:\n\nUser Request: {user_prompt}\n\nInitial Response:"
|
67 |
+
input_ids_1_5b = tokenizers["1.5B"].encode(prompt_1_5b, return_tensors="pt").to(models["1.5B"].device)
|
68 |
+
output_1_5b = models["1.5B"].generate(input_ids_1_5b, max_new_tokens=200, temperature=0.7, do_sample=True) # Reverted to original max_new_tokens (can adjust)
|
69 |
+
response_1_5b = tokenizers["1.5B"].decode(output_1_5b[0], skip_special_tokens=True)
|
70 |
+
print(f"1.5B Response:\n{response_1_5b}")
|
71 |
+
store_in_memory(f"1.5B Model Initial Response: {response_1_5b[:200]}...")
|
72 |
+
|
73 |
+
# 7B Model - Elaboration and Detail
|
74 |
+
print("\n[7B Model - Elaboration] - GPU Accelerated") # Added GPU indication
|
75 |
+
retrieved_memory_7b = retrieve_from_memory(response_1_5b)
|
76 |
+
context_7b = "\n".join([f"- {mem}" for mem in retrieved_memory_7b]) if retrieved_memory_7b else "No relevant context found in memory."
|
77 |
+
prompt_7b = f"Context from Shared Memory:\n{context_7b}\n\nYou are a detailed elaborator. Take the following initial response and elaborate on it, adding more detail and reasoning, considering the context above. \n\nInitial Response:\n{response_1_5b}\n\nElaborated Response:"
|
78 |
+
input_ids_7b = tokenizers["7B"].encode(prompt_7b, return_tensors="pt").to(models["7B"].device)
|
79 |
+
output_7b = models["7B"].generate(input_ids_7b, max_new_tokens=300, temperature=0.7, do_sample=True) # Reverted to original max_new_tokens
|
80 |
+
response_7b = tokenizers["7B"].decode(output_7b[0], skip_special_tokens=True)
|
81 |
+
print(f"7B Response:\n{response_7b}")
|
82 |
+
store_in_memory(f"7B Model Elaborated Response: {response_7b[:200]}...")
|
83 |
+
|
84 |
+
# 14B Model - Final Reasoning and Refinement
|
85 |
+
print("\n[14B Model - Final Refinement] - GPU Accelerated") # Added GPU indication
|
86 |
+
retrieved_memory_14b = retrieve_from_memory(response_7b)
|
87 |
+
context_14b = "\n".join([f"- {mem}" for mem in retrieved_memory_14b]) if retrieved_memory_14b else "No relevant context found in memory."
|
88 |
+
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:"
|
89 |
+
input_ids_14b = tokenizers["14B"].encode(prompt_14b, return_tensors="pt").to(models["14B"].device)
|
90 |
+
output_14b = models["14B"].generate(input_ids_14b, max_new_tokens=400, temperature=0.6, do_sample=True) # Reverted to original max_new_tokens
|
91 |
+
response_14b = tokenizers["14B"].decode(output_14b[0], skip_special_tokens=True)
|
92 |
+
print(f"14B Response (Final):\n{response_14b}")
|
93 |
+
|
94 |
+
return response_14b
|
95 |
+
|
96 |
+
|
97 |
+
# --- Gradio Interface --- (Same as before)
|
98 |
+
def gradio_interface(user_prompt):
|
99 |
+
return swarm_agent_sequential_rag(user_prompt)
|
100 |
+
|
101 |
+
iface = gr.Interface(
|
102 |
+
fn=gradio_interface,
|
103 |
+
inputs=gr.Textbox(lines=5, placeholder="Enter your task here..."),
|
104 |
+
outputs=gr.Textbox(lines=10, placeholder="Agent Swarm Output will appear here..."),
|
105 |
+
title="DeepSeek Agent Swarm (ZeroGPU Demo)",
|
106 |
+
description="Agent swarm using DeepSeek-R1-Distill models (1.5B, 7B, 14B) with shared memory. **GPU accelerated using ZeroGPU!** (Requires Pro Space)", # Updated description
|
107 |
+
)
|
108 |
+
|
109 |
+
if __name__ == "__main__":
|
110 |
+
iface.launch() # Only launch locally if running this script directly
|