wuhp commited on
Commit
ea0faa1
·
verified ·
1 Parent(s): 5e61a38

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +107 -4
app.py CHANGED
@@ -1,7 +1,110 @@
1
  import gradio as gr
 
 
 
2
 
3
- def greet(name):
4
- return "Hello " + name + "!!"
 
 
 
 
5
 
6
- demo = gr.Interface(fn=greet, inputs="text", outputs="text")
7
- demo.launch()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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