Update app.py

app.py

@@ -1,6 +1,6 @@
 """
-SAM-Z-1 Worker Node - Complete Implementation
-Loads model and processes generation requests
+SAM-Z-1 Smart Worker Node
+Supports both full generation and gen/decode split modes
 """
 
 from fastapi import FastAPI, HTTPException
@@ -14,13 +14,13 @@ import os
 from tokenizers import Tokenizer
 import numpy as np
 import time
-from typing import List
+from typing import List, Optional
 import asyncio
 
-app = FastAPI(title="SAM-Z-1 Worker", version="1.0.0")
+app = FastAPI(title="SAM-Z-1 Smart Worker", version="3.0.0")
 
 # ============================================================================
-# Model Architecture Definitions
+# Model Architecture (same as before)
 # ============================================================================
 
 @keras.saving.register_keras_serializable()
@@ -36,7 +36,6 @@ class RotaryEmbedding(keras.layers.Layer):
         super().build(input_shape)
     
     def _build_cache(self):
-        """Build RoPE cache on first forward pass"""
         if not self.built_cache:
             inv_freq = 1.0 / (self.theta ** (tf.range(0, self.dim, 2, dtype=tf.float32) / self.dim))
             t = tf.range(self.max_len, dtype=tf.float32)
@@ -53,7 +52,6 @@ class RotaryEmbedding(keras.layers.Layer):
     
     def call(self, q, k):
         self._build_cache()
-        
         seq_len = tf.shape(q)[2]
         dtype = q.dtype
         cos = tf.cast(self.cos_cached[:seq_len, :], dtype)[None, None, :, :]
@@ -69,7 +67,6 @@ class RotaryEmbedding(keras.layers.Layer):
         config.update({"dim": self.dim, "max_len": self.max_len, "theta": self.theta})
         return config
 
-
 @keras.saving.register_keras_serializable()
 class RMSNorm(keras.layers.Layer):
     def __init__(self, epsilon=1e-5, **kwargs):
@@ -88,7 +85,6 @@ class RMSNorm(keras.layers.Layer):
         config.update({"epsilon": self.epsilon})
         return config
 
-
 @keras.saving.register_keras_serializable()
 class TransformerBlock(keras.layers.Layer):
     def __init__(self, d_model, n_heads, ff_dim, dropout, max_len, rope_theta, layer_idx=0, **kwargs):
@@ -122,7 +118,6 @@ class TransformerBlock(keras.layers.Layer):
         B, T, D = tf.shape(x)[0], tf.shape(x)[1], self.d_model
         dtype = x.dtype
         
-        # Attention
         res = x
         y = self.pre_attn_norm(x)
         
@@ -133,7 +128,6 @@ class TransformerBlock(keras.layers.Layer):
         q, k = self.rope(q, k)
         
         scores = tf.matmul(q, k, transpose_b=True) / tf.sqrt(tf.cast(self.head_dim, dtype))
-        
         mask = tf.where(
             tf.linalg.band_part(tf.ones([T, T], dtype=dtype), -1, 0) == 0,
             tf.constant(-1e9, dtype=dtype),
@@ -145,7 +139,6 @@ class TransformerBlock(keras.layers.Layer):
         attn = tf.reshape(tf.transpose(attn, [0, 2, 1, 3]), [B, T, D])
         x = res + self.dropout(self.out_proj(attn), training=training)
         
-        # FFN (SwiGLU)
         res = x
         y = self.pre_ffn_norm(x)
         ffn = self.down_proj(keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
@@ -165,7 +158,6 @@ class TransformerBlock(keras.layers.Layer):
         })
         return config
 
-
 @keras.saving.register_keras_serializable()
 class SAM1Model(keras.Model):
     def __init__(self, **kwargs):
@@ -199,10 +191,8 @@ class SAM1Model(keras.Model):
     
     def call(self, input_ids, training=None):
         x = self.embed(input_ids)
-        
         for block in self.blocks:
             x = block(x, training=training)
-        
         return self.lm_head(self.norm(x))
     
     def get_config(self):
@@ -235,6 +225,7 @@ class GenerateRequest(BaseModel):
     top_p: float = 0.9
     repetition_penalty: float = 1.1
     stream: bool = False
+    return_token_ids: bool = False  # NEW: for gen/decode split
 
 class ChatMessage(BaseModel):
     role: str
@@ -248,6 +239,10 @@ class ChatRequest(BaseModel):
     top_p: float = 0.9
     repetition_penalty: float = 1.1
     stream: bool = False
+    return_token_ids: bool = False  # NEW
+
+class DecodeRequest(BaseModel):
+    token_ids: List[int]
 
 # ============================================================================
 # Generation Functions
@@ -259,12 +254,16 @@ def generate_tokens(
     temperature: float = 0.8,
     top_k: int = 40,
     top_p: float = 0.9,
-    repetition_penalty: float = 1.1
+    repetition_penalty: float = 1.1,
+    return_token_ids: bool = False
 ):
-    """Core generation function (yields token IDs)"""
+    """
+    Core generation function
+    If return_token_ids=True, yields (token_id, None)
+    If return_token_ids=False, yields (token_id, token_text)
+    """
     global model, tokenizer, config, eos_token_id, fast_forward
     
-    # Tokenize
     input_ids = [i for i in tokenizer.encode(prompt).ids if i != eos_token_id]
     
     if len(input_ids) == 0:
@@ -277,26 +276,21 @@ def generate_tokens(
     token_freq = {}
     
     for step in range(max_tokens):
-        # Get logits
         logits = fast_forward(input_tensor)
         next_token_logits = logits[0, -1, :].numpy()
         
-        # Temperature
         next_token_logits = next_token_logits / temperature
         
-        # Repetition penalty
         if repetition_penalty != 1.0:
             for token_id, freq in token_freq.items():
                 if token_id < len(next_token_logits):
                     next_token_logits[token_id] /= (repetition_penalty ** freq)
         
-        # Top-k filtering
         if top_k > 0:
             top_k_indices = np.argpartition(next_token_logits, -top_k)[-top_k:]
             top_k_logits = next_token_logits[top_k_indices]
             top_k_probs = tf.nn.softmax(top_k_logits).numpy()
             
-            # Top-p sampling
             if top_p < 1.0:
                 sorted_indices = np.argsort(top_k_probs)[::-1]
                 cumsum = np.cumsum(top_k_probs[sorted_indices])
@@ -315,16 +309,18 @@ def generate_tokens(
             probs = tf.nn.softmax(next_token_logits).numpy()
             next_token_id = np.random.choice(len(probs), p=probs)
         
-        # Stop on EOS
         if next_token_id == eos_token_id:
             break
         
         token_freq[next_token_id] = token_freq.get(next_token_id, 0) + 1
         
-        # Yield token
-        yield next_token_id
+        # Yield token ID and optionally decoded text
+        if return_token_ids:
+            yield (next_token_id, None)
+        else:
+            token_text = tokenizer.decode([next_token_id])
+            yield (next_token_id, token_text)
         
-        # Update input
         input_tensor = tf.concat([input_tensor, [[next_token_id]]], axis=1)
         
         if input_tensor.shape[1] > config['max_position_embeddings']:
@@ -350,12 +346,15 @@ def format_chat_prompt(messages: List[ChatMessage]) -> str:
 async def root():
     """Worker info"""
     return {
-        "name": "SAM-Z-1 Worker",
+        "name": "SAM-Z-1 Smart Worker",
+        "version": "3.0.0",
         "status": "ready" if model is not None else "loading",
         "model": MODEL_REPO,
+        "features": ["full_generation", "token_only_mode", "decode_only_mode"],
         "endpoints": {
             "generate": "/generate",
             "chat": "/chat",
+            "decode": "/decode",
             "health": "/health"
         }
     }
@@ -368,11 +367,27 @@ async def health():
         "model_loaded": model is not None
     }
 
+@app.post("/decode")
+async def decode(request: DecodeRequest):
+    """
+    DECODE ONLY endpoint
+    Takes token IDs and returns decoded text
+    This is the bottleneck we're parallelizing!
+    """
+    if tokenizer is None:
+        raise HTTPException(status_code=503, detail="Tokenizer not loaded")
+    
+    try:
+        text = tokenizer.decode(request.token_ids)
+        return {"text": text}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Decode error: {str(e)}")
+
 @app.post("/generate")
 async def generate(request: GenerateRequest):
-    """Generate text from prompt"""
+    """Generate text - supports both full gen and token-only mode"""
     if model is None:
-        raise HTTPException(status_code=503, detail="Model not loaded yet, please wait")
+        raise HTTPException(status_code=503, detail="Model not loaded yet")
     
     start_time = time.time()
     
@@ -383,27 +398,29 @@ async def generate(request: GenerateRequest):
             token_count = 0
             
             try:
-                for token_id in generate_tokens(
+                for token_id, token_text in generate_tokens(
                     request.prompt,
                     max_tokens=request.max_tokens,
                     temperature=request.temperature,
                     top_k=request.top_k,
                     top_p=request.top_p,
-                    repetition_penalty=request.repetition_penalty
+                    repetition_penalty=request.repetition_penalty,
+                    return_token_ids=request.return_token_ids
                 ):
-                    token_text = tokenizer.decode([token_id])
-                    generated_text += token_text
                     token_count += 1
                     
-                    # Send chunk
-                    yield f"data: {json.dumps({'text': token_text, 'total': generated_text})}\n\n"
+                    if request.return_token_ids:
+                        # TOKEN-ONLY mode for gen/decode split
+                        yield f"data: {json.dumps({'token_id': token_id})}\n\n"
+                    else:
+                        # FULL mode with text
+                        generated_text += token_text
+                        yield f"data: {json.dumps({'text': token_text, 'total': generated_text})}\n\n"
                     
-                    # Small delay
                     await asyncio.sleep(0.001)
                 
-                # Send final stats
                 elapsed = time.time() - start_time
-                yield f"data: {json.dumps({'done': True, 'tokens': token_count, 'time': elapsed, 'tokens_per_sec': token_count/elapsed if elapsed > 0 else 0})}\n\n"
+                yield f"data: {json.dumps({'done': True, 'tokens': token_count, 'time': elapsed})}\n\n"
             
             except Exception as e:
                 yield f"data: {json.dumps({'error': str(e)})}\n\n"
@@ -411,21 +428,22 @@ async def generate(request: GenerateRequest):
         return StreamingResponse(stream_tokens(), media_type="text/event-stream")
     
     else:
-        # Non-streaming response
+        # Non-streaming
         generated_text = ""
         token_count = 0
         
         try:
-            for token_id in generate_tokens(
+            for token_id, token_text in generate_tokens(
                 request.prompt,
                 max_tokens=request.max_tokens,
                 temperature=request.temperature,
                 top_k=request.top_k,
                 top_p=request.top_p,
-                repetition_penalty=request.repetition_penalty
+                repetition_penalty=request.repetition_penalty,
+                return_token_ids=request.return_token_ids
             ):
-                token_text = tokenizer.decode([token_id])
-                generated_text += token_text
+                if not request.return_token_ids:
+                    generated_text += token_text
                 token_count += 1
             
             elapsed = time.time() - start_time
@@ -442,44 +460,45 @@ async def generate(request: GenerateRequest):
 
 @app.post("/chat")
 async def chat(request: ChatRequest):
-    """Chat completion"""
+    """Chat completion - supports both modes"""
     if model is None:
-        raise HTTPException(status_code=503, detail="Model not loaded yet, please wait")
+        raise HTTPException(status_code=503, detail="Model not loaded yet")
     
-    # Format prompt
     prompt = format_chat_prompt(request.messages)
-    
     start_time = time.time()
     
     if request.stream:
-        # Streaming
         async def stream_tokens():
             generated_text = ""
             token_count = 0
             
             try:
-                for token_id in generate_tokens(
+                for token_id, token_text in generate_tokens(
                     prompt,
                     max_tokens=request.max_tokens,
                     temperature=request.temperature,
                     top_k=request.top_k,
                     top_p=request.top_p,
-                    repetition_penalty=request.repetition_penalty
+                    repetition_penalty=request.repetition_penalty,
+                    return_token_ids=request.return_token_ids
                 ):
-                    token_text = tokenizer.decode([token_id])
-                    generated_text += token_text
                     token_count += 1
                     
-                    # Stop at end tag
-                    if "<|im_end|>" in generated_text:
-                        generated_text = generated_text.split("<|im_end|>")[0]
-                        break
+                    if request.return_token_ids:
+                        yield f"data: {json.dumps({'token_id': token_id})}\n\n"
+                    else:
+                        generated_text += token_text
+                        
+                        if "<|im_end|>" in generated_text:
+                            generated_text = generated_text.split("<|im_end|>")[0]
+                            break
+                        
+                        yield f"data: {json.dumps({'delta': token_text, 'content': generated_text})}\n\n"
                     
-                    yield f"data: {json.dumps({'delta': token_text, 'content': generated_text})}\n\n"
                     await asyncio.sleep(0.001)
                 
                 elapsed = time.time() - start_time
-                yield f"data: {json.dumps({'done': True, 'tokens': token_count, 'time': elapsed, 'tokens_per_sec': token_count/elapsed if elapsed > 0 else 0})}\n\n"
+                yield f"data: {json.dumps({'done': True, 'tokens': token_count, 'time': elapsed})}\n\n"
             
             except Exception as e:
                 yield f"data: {json.dumps({'error': str(e)})}\n\n"
@@ -487,26 +506,27 @@ async def chat(request: ChatRequest):
         return StreamingResponse(stream_tokens(), media_type="text/event-stream")
     
     else:
-        # Non-streaming
         generated_text = ""
         token_count = 0
         
         try:
-            for token_id in generate_tokens(
+            for token_id, token_text in generate_tokens(
                 prompt,
                 max_tokens=request.max_tokens,
                 temperature=request.temperature,
                 top_k=request.top_k,
                 top_p=request.top_p,
-                repetition_penalty=request.repetition_penalty
+                repetition_penalty=request.repetition_penalty,
+                return_token_ids=request.return_token_ids
             ):
-                token_text = tokenizer.decode([token_id])
-                generated_text += token_text
-                token_count += 1
+                if not request.return_token_ids:
+                    generated_text += token_text
+                    
+                    if "<|im_end|>" in generated_text:
+                        generated_text = generated_text.split("<|im_end|>")[0]
+                        break
                 
-                if "<|im_end|>" in generated_text:
-                    generated_text = generated_text.split("<|im_end|>")[0]
-                    break
+                token_count += 1
             
             elapsed = time.time() - start_time
             
@@ -535,10 +555,8 @@ async def load_model():
     print("🚀 Loading SAM-Z-1 Model...")
     
     try:
-        # Download model files
         config_path = hf_hub_download(MODEL_REPO, "config.json", cache_dir=CACHE_DIR)
         
-        # Try checkpoint first
         try:
             weights_path = hf_hub_download(MODEL_REPO, "ckpt.weights.h5", cache_dir=CACHE_DIR)
             print("✅ Found checkpoint weights")
@@ -548,13 +566,11 @@ async def load_model():
             model_path = hf_hub_download(MODEL_REPO, "model.keras", cache_dir=CACHE_DIR)
             use_checkpoint = False
         
-        # Load config
         with open(config_path, 'r') as f:
             config = json.load(f)
         
         print(f"📦 Config loaded: {config['num_hidden_layers']} layers")
         
-        # Create tokenizer
         print("📦 Creating tokenizer...")
         from transformers import AutoTokenizer
         
@@ -570,11 +586,9 @@ async def load_model():
         
         print(f"✅ Tokenizer ready: vocab size {tokenizer.get_vocab_size()}")
         
-        # Load model
         print("🔄 Loading model...")
         
         if use_checkpoint:
-            # Build from config
             model_config = {
                 'vocab_size': config['vocab_size'],
                 'd_model': config['hidden_size'],
@@ -587,30 +601,25 @@ async def load_model():
             }
             
             model = SAM1Model(config=model_config)
-            
-            # Build with dummy input
             dummy_input = tf.zeros((1, config['max_position_embeddings']), dtype=tf.int32)
             _ = model(dummy_input, training=False)
             
             print(f"✅ Architecture built: {model.count_params():,} parameters")
             
-            # Load weights
             model.load_weights(weights_path)
             print("✅ Weights loaded!")
         
         else:
-            # Load full model
             model = keras.models.load_model(model_path, compile=False)
             print("✅ Model loaded!")
         
-        # Create optimized inference function
         @tf.function(reduce_retracing=True)
         def optimized_forward(input_tensor):
             return model(input_tensor, training=False)
         
         fast_forward = optimized_forward
         
-        print("✅ SAM-Z-1 Worker ready for inference! 🚀")
+        print("✅ SAM-Z-1 Smart Worker ready! 🚀")
     
     except Exception as e:
         print(f"❌ Failed to load model: {e}")