Update app.py

app.py

@@ -8,30 +8,17 @@ Edit the CONFIG below, then deploy.
 # ============================================================================
 
 CONFIG = {
-    # This node's identity
     "node_id": "head-main",
-    
-    # Which transformer blocks this node runs (0-indexed)
-    # Sam-large-2 has 12 blocks (0-11)
     "layer_start": 0,
-    "layer_end": 6,  # exclusive, so this runs blocks 0,1,2,3,4,5
-    
-    # Worker Space URLs (in order of execution)
-    # Leave empty [] for standalone mode (all layers on this node)
-    "worker_urls": [
-        # "https://YOUR-WORKER-SPACE.hf.space",
-    ],
-    
-    # Shared secret for worker communication
+    "layer_end": 6,
+    "worker_urls": [],
     "secret_token": "sam2-distributed-secret-change-me",
-    
-    # Model settings
     "model_repo": "Smilyai-labs/Sam-large-2",
     "cache_dir": "./model_cache",
 }
 
 # ============================================================================
-# CPU Optimization - MUST be before TensorFlow import
+# CPU Optimization
 # ============================================================================
 
 import os
@@ -45,7 +32,6 @@ os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 
 import json
 import time
-import threading
 import io
 import base64
 from typing import Dict, List, Optional, Tuple, Any
@@ -204,13 +190,10 @@ class ModelState:
         self.config = None
         self.tokenizer = None
         self.eos_token_id = 50256
-        
-        # Model components
         self.embedding = None
         self.blocks: List = []
         self.final_norm = None
         self.lm_head = None
-        
         self.my_block_start = 0
         self.my_block_end = 0
 
@@ -245,7 +228,6 @@ def deserialize_kv_cache(data):
 # ============================================================================
 
 def call_worker(url: str, hidden_states: tf.Tensor, past_kv=None, use_cache=False) -> Tuple[tf.Tensor, Any]:
-    """Send hidden states to worker and get result."""
     try:
         response = requests.post(
             f"{url.rstrip('/')}/api/forward",
@@ -273,16 +255,13 @@ def call_worker(url: str, hidden_states: tf.Tensor, past_kv=None, use_cache=Fals
 # ============================================================================
 
 def load_model():
-    """Load model and extract components for this node."""
     print("🚀 Loading model...")
     
-    # Load config
     config_path = hf_hub_download(CONFIG["model_repo"], "config.json", cache_dir=CONFIG["cache_dir"])
     with open(config_path, 'r') as f:
         model_config = json.load(f)
     STATE.config = model_config
     
-    # Load tokenizer
     from transformers import AutoTokenizer
     from tokenizers import Tokenizer
     
@@ -294,10 +273,8 @@ def load_model():
     STATE.tokenizer = Tokenizer.from_file("./temp_tokenizer/tokenizer.json")
     STATE.eos_token_id = model_config.get('eos_token_id', 50256)
     
-    # Load weights
     weights_path = hf_hub_download(CONFIG["model_repo"], "ckpt.weights.h5", cache_dir=CONFIG["cache_dir"])
     
-    # Build full model to load weights
     n_layers = model_config['num_hidden_layers']
     d_model = model_config['hidden_size']
     n_heads = model_config['num_attention_heads']
@@ -306,14 +283,12 @@ def load_model():
     rope_theta = model_config['rope_theta']
     vocab_size = model_config['vocab_size']
     
-    # Temporary full model
     embedding = keras.layers.Embedding(vocab_size, d_model, name="embed_tokens")
     blocks = [TransformerBlock(d_model, n_heads, ff_dim, 0.0, max_len, rope_theta, i, name=f"block_{i}") 
               for i in range(n_layers)]
     final_norm = RMSNorm(name="final_norm")
     lm_head = keras.layers.Dense(vocab_size, use_bias=False, name="lm_head")
     
-    # Build
     dummy = tf.zeros((1, 16), dtype=tf.int32)
     x = embedding(dummy)
     for block in blocks:
@@ -321,7 +296,6 @@ def load_model():
     x = final_norm(x)
     _ = lm_head(x)
     
-    # Load weights into a temp model structure
     class TempModel(keras.Model):
         def __init__(self):
             super().__init__()
@@ -340,25 +314,19 @@ def load_model():
     temp_model.load_weights(weights_path)
     print("✅ Weights loaded")
     
-    # Extract components for this node
     STATE.my_block_start = CONFIG["layer_start"]
     STATE.my_block_end = CONFIG["layer_end"] if CONFIG["layer_end"] > 0 else n_layers
     
-    # HEAD always has embedding
     STATE.embedding = embedding
-    
-    # Extract our blocks
     STATE.blocks = blocks[STATE.my_block_start:STATE.my_block_end]
     print(f"✅ Loaded blocks {STATE.my_block_start} to {STATE.my_block_end - 1}")
     
-    # HEAD has final norm and lm_head only if no workers OR we handle last block
     has_workers = len(CONFIG["worker_urls"]) > 0
     if not has_workers:
         STATE.final_norm = final_norm
         STATE.lm_head = lm_head
         print("✅ Loaded final norm and LM head (standalone mode)")
     
-    # Warmup
     print("🔥 Warming up...")
     dummy = tf.constant([[1, 2, 3]], dtype=tf.int32)
     x = STATE.embedding(dummy)
@@ -375,14 +343,8 @@ def load_model():
 # ============================================================================
 
 def forward_pass(input_ids: tf.Tensor, past_kv_local=None, past_kv_workers=None, use_cache=False):
-    """
-    Full forward pass through HEAD + all workers.
-    Returns logits and updated KV caches.
-    """
-    # Embedding
     x = STATE.embedding(input_ids)
     
-    # Local blocks
     new_local_kv = [] if use_cache else None
     for i, block in enumerate(STATE.blocks):
         block_past = past_kv_local[i] if past_kv_local else None
@@ -390,7 +352,6 @@ def forward_pass(input_ids: tf.Tensor, past_kv_local=None, past_kv_workers=None,
         if use_cache:
             new_local_kv.append(kv)
     
-    # Workers
     new_worker_kv = {} if use_cache else None
     for worker_url in CONFIG["worker_urls"]:
         worker_past = past_kv_workers.get(worker_url) if past_kv_workers else None
@@ -398,12 +359,9 @@ def forward_pass(input_ids: tf.Tensor, past_kv_local=None, past_kv_workers=None,
         if use_cache:
             new_worker_kv[worker_url] = worker_kv
     
-    # Final (only if standalone or last worker returned to us)
-    # In distributed mode, the last worker applies final_norm + lm_head
     if STATE.lm_head:
         logits = STATE.lm_head(STATE.final_norm(x))
     else:
-        # x should already be logits from last worker
         logits = x
     
     return logits, new_local_kv, new_worker_kv
@@ -465,7 +423,6 @@ def generate_stream(prompt: str, max_tokens=512, temperature=0.8, top_k=40, top_
     
     start = time.time()
     
-    # Prefill
     input_tensor = tf.constant([input_ids], dtype=tf.int32)
     try:
         logits, local_kv, worker_kv = forward_pass(input_tensor, None, None, use_cache=True)
@@ -477,7 +434,6 @@ def generate_stream(prompt: str, max_tokens=512, temperature=0.8, top_k=40, top_
     prefill_time = time.time() - start
     print(f"⚡ Prefill: {len(input_ids)} tokens in {prefill_time:.2f}s")
     
-    # Generate
     decode_start = time.time()
     tokens_generated = 0
     
@@ -496,7 +452,6 @@ def generate_stream(prompt: str, max_tokens=512, temperature=0.8, top_k=40, top_
         tokens_generated += 1
         yield generated
         
-        # Next step
         next_input = tf.constant([[next_id]], dtype=tf.int32)
         try:
             logits, local_kv, worker_kv = forward_pass(next_input, local_kv, worker_kv, use_cache=True)
@@ -506,7 +461,6 @@ def generate_stream(prompt: str, max_tokens=512, temperature=0.8, top_k=40, top_
         
         next_logits = logits[0, -1, :].numpy()
     
-    # Stats
     if tokens_generated > 0:
         total = time.time() - start
         tps = tokens_generated / (time.time() - decode_start)
@@ -519,10 +473,12 @@ def generate_stream(prompt: str, max_tokens=512, temperature=0.8, top_k=40, top_
 
 def format_prompt(message: str, history: list, reasoning: bool) -> str:
     prompt = ""
-    for user, assistant in history:
-        prompt += f"<|im_start|>user\n{user}<|im_end|>\n"
-        if assistant:
-            prompt += f"<|im_start|>assistant\n{assistant.split('*[')[0].strip()}<|im_end|>\n"
+    for msg in history:
+        if msg["role"] == "user":
+            prompt += f"<|im_start|>user\n{msg['content']}<|im_end|>\n"
+        elif msg["role"] == "assistant":
+            content = msg['content'].split('*[')[0].strip()
+            prompt += f"<|im_start|>assistant\n{content}<|im_end|>\n"
     prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
     if reasoning:
         prompt += "<think>"
@@ -536,21 +492,27 @@ def chat_respond(message, history, max_tokens, temp, top_k, top_p, rep_pen, reas
     
     prompt = format_prompt(message, history, reasoning)
     
+    # Add user message to history
+    history = history + [{"role": "user", "content": message}]
+    
     for text in generate_stream(prompt, max_tokens, temp, top_k, top_p, rep_pen):
         display = text
+        
+        # Clean stop tags
         for tag in ["<|im_end|>", "<im end for model tun>"]:
             if tag in display:
                 idx = display.find(tag)
                 stats = display.find("\n\n*[")
                 display = display[:idx] + (display[stats:] if stats > idx else "")
         
+        # Format reasoning
         if reasoning and '<think>' in display and '</think>' in display:
             s, e = display.find('<think>'), display.find('</think>')
             if s < e:
                 thought = display[s+7:e].strip()
                 display = display[:s] + f'<details><summary>🧠 Reasoning</summary><p>{thought}</p></details>' + display[e+8:]
         
-        yield history + [[message, display.strip()]]
+        yield history + [{"role": "assistant", "content": display.strip()}]
 
 
 def stop():
@@ -575,7 +537,11 @@ def create_ui():
             gr.Markdown("**Workers:** " + ", ".join(f"`{w}`" for w in workers))
         
         reasoning = gr.State(False)
-        chatbot = gr.Chatbot(height=500)
+        
+        chatbot = gr.Chatbot(
+            height=500,
+            type="messages"  # Use new messages format
+        )
         
         with gr.Row():
             reason_btn = gr.Button("💡", size="sm", scale=0)
@@ -600,7 +566,7 @@ def create_ui():
         click = send.click(chat_respond, inputs, chatbot).then(lambda: "", outputs=msg)
         stop_btn.click(stop, cancels=[submit, click])
         
-        gr.Button("🗑️ Clear").click(lambda: ([], ""), outputs=[chatbot, msg])
+        gr.Button("🗑️ Clear").click(lambda: [], outputs=[chatbot])
     
     return app