Update app.py

app.py

@@ -1,10 +1,10 @@
 """
-SAM-Z-1 Smart Worker Node
-Supports both full generation and gen/decode split modes
+SAM-Z-1 Distributed Worker Node v4.0
+Optimized for distributed gen/decode pipeline
 """
 
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, HTMLResponse
 from pydantic import BaseModel
 import tensorflow as tf
 import keras
@@ -17,10 +17,10 @@ import time
 from typing import List, Optional
 import asyncio
 
-app = FastAPI(title="SAM-Z-1 Smart Worker", version="3.0.0")
+app = FastAPI(title="SAM-Z-1 Distributed Worker", version="4.0.0")
 
 # ============================================================================
-# Model Architecture (same as before)
+# Model Architecture
 # ============================================================================
 
 @keras.saving.register_keras_serializable()
@@ -201,7 +201,7 @@ class SAM1Model(keras.Model):
         return base_config
 
 # ============================================================================
-# Global Variables
+# Global State
 # ============================================================================
 
 model = None
@@ -213,6 +213,14 @@ fast_forward = None
 MODEL_REPO = "Smilyai-labs/Sam-Z-1-tensorflow"
 CACHE_DIR = "./model_cache"
 
+# Stats
+worker_stats = {
+    "total_requests": 0,
+    "total_tokens": 0,
+    "decode_requests": 0,
+    "uptime_start": time.time()
+}
+
 # ============================================================================
 # Request Models
 # ============================================================================
@@ -225,7 +233,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
+    return_token_ids: bool = False
 
 class ChatMessage(BaseModel):
     role: str
@@ -239,11 +247,14 @@ class ChatRequest(BaseModel):
     top_p: float = 0.9
     repetition_penalty: float = 1.1
     stream: bool = False
-    return_token_ids: bool = False  # NEW
+    return_token_ids: bool = False
 
 class DecodeRequest(BaseModel):
     token_ids: List[int]
 
+class BatchDecodeRequest(BaseModel):
+    batches: List[List[int]]
+
 # ============================================================================
 # Generation Functions
 # ============================================================================
@@ -257,11 +268,7 @@ def generate_tokens(
     repetition_penalty: float = 1.1,
     return_token_ids: bool = False
 ):
-    """
-    Core generation function
-    If return_token_ids=True, yields (token_id, None)
-    If return_token_ids=False, yields (token_id, token_text)
-    """
+    """Core generation - yields (token_id, token_text or None)"""
     global model, tokenizer, config, eos_token_id, fast_forward
     
     input_ids = [i for i in tokenizer.encode(prompt).ids if i != eos_token_id]
@@ -314,7 +321,6 @@ def generate_tokens(
         
         token_freq[next_token_id] = token_freq.get(next_token_id, 0) + 1
         
-        # Yield token ID and optionally decoded text
         if return_token_ids:
             yield (next_token_id, None)
         else:
@@ -327,7 +333,6 @@ def generate_tokens(
             input_tensor = input_tensor[:, -config['max_position_embeddings']:]
 
 def format_chat_prompt(messages: List[ChatMessage]) -> str:
-    """Format chat messages into prompt"""
     prompt = ""
     for msg in messages:
         if msg.role == "user":
@@ -339,60 +344,262 @@ def format_chat_prompt(messages: List[ChatMessage]) -> str:
     return prompt
 
 # ============================================================================
-# API Endpoints
+# Status Page
 # ============================================================================
 
-@app.get("/")
-async def root():
-    """Worker info"""
-    return {
-        "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"
+@app.get("/", response_class=HTMLResponse)
+async def status_page():
+    """Worker status page"""
+    return """
+<!DOCTYPE html>
+<html>
+<head>
+    <title>SAM-Z-1 Worker Node</title>
+    <style>
+        * { margin: 0; padding: 0; box-sizing: border-box; }
+        body {
+            font-family: 'Courier New', monospace;
+            background: linear-gradient(135deg, #1a1f3a 0%, #0a0e27 100%);
+            color: #00bfff;
+            padding: 20px;
+            min-height: 100vh;
         }
-    }
+        .container {
+            max-width: 900px;
+            margin: 0 auto;
+        }
+        .header {
+            text-align: center;
+            padding: 30px;
+            background: rgba(0, 191, 255, 0.1);
+            border: 2px solid #00bfff;
+            border-radius: 10px;
+            margin-bottom: 30px;
+            box-shadow: 0 0 20px rgba(0, 191, 255, 0.3);
+        }
+        .header h1 {
+            font-size: 2.5em;
+            text-transform: uppercase;
+            letter-spacing: 3px;
+            animation: glow 2s ease-in-out infinite alternate;
+        }
+        @keyframes glow {
+            from { text-shadow: 0 0 10px #00bfff; }
+            to { text-shadow: 0 0 20px #00bfff, 0 0 30px #00bfff; }
+        }
+        .badge {
+            display: inline-block;
+            padding: 5px 15px;
+            border-radius: 15px;
+            font-size: 0.9em;
+            margin-top: 10px;
+        }
+        .badge-ready {
+            background: rgba(0, 255, 136, 0.2);
+            border: 1px solid #00ff88;
+            color: #00ff88;
+        }
+        .badge-loading {
+            background: rgba(255, 165, 0, 0.2);
+            border: 1px solid #ffa500;
+            color: #ffa500;
+        }
+        .stats-grid {
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+            gap: 20px;
+            margin-bottom: 30px;
+        }
+        .stat-card {
+            background: rgba(0, 191, 255, 0.05);
+            border: 1px solid #00bfff;
+            border-radius: 8px;
+            padding: 20px;
+            text-align: center;
+        }
+        .stat-label {
+            font-size: 0.8em;
+            opacity: 0.7;
+            text-transform: uppercase;
+            margin-bottom: 10px;
+        }
+        .stat-value {
+            font-size: 2em;
+            font-weight: bold;
+        }
+        .features {
+            background: rgba(0, 191, 255, 0.05);
+            border: 1px solid #00bfff;
+            border-radius: 8px;
+            padding: 20px;
+        }
+        .features h3 {
+            margin-bottom: 15px;
+        }
+        .feature-list {
+            list-style: none;
+            padding: 0;
+        }
+        .feature-list li {
+            padding: 10px;
+            margin: 5px 0;
+            background: rgba(0, 191, 255, 0.1);
+            border-radius: 5px;
+        }
+        .feature-list li:before {
+            content: "⚡ ";
+            color: #00ff88;
+        }
+        .timestamp {
+            text-align: center;
+            margin-top: 20px;
+            opacity: 0.5;
+        }
+    </style>
+</head>
+<body>
+    <div class="container">
+        <div class="header">
+            <h1>⚙️ WORKER NODE ⚙️</h1>
+            <div>SAM-Z-1 Distributed Worker v4.0</div>
+            <div class="badge" id="status-badge">CHECKING STATUS...</div>
+        </div>
+        
+        <div class="stats-grid" id="stats">
+            <div class="stat-card">
+                <div class="stat-label">Total Requests</div>
+                <div class="stat-value" id="total-req">--</div>
+            </div>
+            <div class="stat-card">
+                <div class="stat-label">Total Tokens</div>
+                <div class="stat-value" id="total-tokens">--</div>
+            </div>
+            <div class="stat-card">
+                <div class="stat-label">Decode Requests</div>
+                <div class="stat-value" id="decode-req">--</div>
+            </div>
+            <div class="stat-card">
+                <div class="stat-label">Uptime</div>
+                <div class="stat-value" id="uptime">--</div>
+            </div>
+        </div>
+        
+        <div class="features">
+            <h3>🚀 CAPABILITIES</h3>
+            <ul class="feature-list">
+                <li>Full Text Generation</li>
+                <li>Token-Only Mode (for distributed pipeline)</li>
+                <li>High-Speed Batch Decoding</li>
+                <li>Chat Completion</li>
+                <li>Streaming & Non-Streaming</li>
+            </ul>
+        </div>
+        
+        <div class="timestamp" id="timestamp">Initializing...</div>
+    </div>
+    
+    <script>
+        async function updateStats() {
+            try {
+                const response = await fetch('/health');
+                const data = await response.json();
+                
+                const badge = document.getElementById('status-badge');
+                if (data.model_loaded) {
+                    badge.textContent = '✅ READY FOR INFERENCE';
+                    badge.className = 'badge badge-ready';
+                } else {
+                    badge.textContent = '⏳ LOADING MODEL...';
+                    badge.className = 'badge badge-loading';
+                }
+                
+                // Fetch stats
+                const statsRes = await fetch('/stats');
+                const stats = await statsRes.json();
+                
+                document.getElementById('total-req').textContent = stats.total_requests;
+                document.getElementById('total-tokens').textContent = stats.total_tokens;
+                document.getElementById('decode-req').textContent = stats.decode_requests;
+                
+                const uptime = Math.floor(stats.uptime);
+                const h = Math.floor(uptime / 3600);
+                const m = Math.floor((uptime % 3600) / 60);
+                const s = uptime % 60;
+                document.getElementById('uptime').textContent = `${h}h ${m}m ${s}s`;
+                
+                document.getElementById('timestamp').textContent = 
+                    `Last update: ${new Date().toLocaleTimeString()}`;
+            } catch (e) {
+                console.error('Failed to update stats:', e);
+            }
+        }
+        
+        // Update every second
+        setInterval(updateStats, 1000);
+        updateStats();
+    </script>
+</body>
+</html>
+    """
+
+# ============================================================================
+# API Endpoints
+# ============================================================================
 
 @app.get("/health")
 async def health():
-    """Health check"""
     return {
         "status": "healthy" if model is not None else "loading",
         "model_loaded": model is not None
     }
 
+@app.get("/stats")
+async def stats():
+    uptime = time.time() - worker_stats["uptime_start"]
+    return {
+        "total_requests": worker_stats["total_requests"],
+        "total_tokens": worker_stats["total_tokens"],
+        "decode_requests": worker_stats["decode_requests"],
+        "uptime": uptime,
+        "tokens_per_second": worker_stats["total_tokens"] / uptime if uptime > 0 else 0
+    }
+
 @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!
-    """
+    """Fast single decode"""
     if tokenizer is None:
         raise HTTPException(status_code=503, detail="Tokenizer not loaded")
     
     try:
+        worker_stats["decode_requests"] += 1
         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("/decode/batch")
+async def batch_decode(request: BatchDecodeRequest):
+    """Optimized batch decoding for distributed pipeline"""
+    if tokenizer is None:
+        raise HTTPException(status_code=503, detail="Tokenizer not loaded")
+    
+    try:
+        worker_stats["decode_requests"] += len(request.batches)
+        results = [tokenizer.decode(batch) for batch in request.batches]
+        return {"texts": results}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Batch decode error: {str(e)}")
+
 @app.post("/generate")
 async def generate(request: GenerateRequest):
-    """Generate text - supports both full gen and token-only mode"""
+    """Generate text"""
     if model is None:
-        raise HTTPException(status_code=503, detail="Model not loaded yet")
+        raise HTTPException(status_code=503, detail="Model not loaded")
     
+    worker_stats["total_requests"] += 1
     start_time = time.time()
     
     if request.stream:
-        # Streaming response
         async def stream_tokens():
             generated_text = ""
             token_count = 0
@@ -408,12 +615,11 @@ async def generate(request: GenerateRequest):
                     return_token_ids=request.return_token_ids
                 ):
                     token_count += 1
+                    worker_stats["total_tokens"] += 1
                     
                     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"
                     
@@ -428,7 +634,6 @@ async def generate(request: GenerateRequest):
         return StreamingResponse(stream_tokens(), media_type="text/event-stream")
     
     else:
-        # Non-streaming
         generated_text = ""
         token_count = 0
         
@@ -445,6 +650,7 @@ async def generate(request: GenerateRequest):
                 if not request.return_token_ids:
                     generated_text += token_text
                 token_count += 1
+                worker_stats["total_tokens"] += 1
             
             elapsed = time.time() - start_time
             
@@ -460,10 +666,11 @@ async def generate(request: GenerateRequest):
 
 @app.post("/chat")
 async def chat(request: ChatRequest):
-    """Chat completion - supports both modes"""
+    """Chat completion"""
     if model is None:
-        raise HTTPException(status_code=503, detail="Model not loaded yet")
+        raise HTTPException(status_code=503, detail="Model not loaded")
     
+    worker_stats["total_requests"] += 1
     prompt = format_chat_prompt(request.messages)
     start_time = time.time()
     
@@ -483,6 +690,7 @@ async def chat(request: ChatRequest):
                     return_token_ids=request.return_token_ids
                 ):
                     token_count += 1
+                    worker_stats["total_tokens"] += 1
                     
                     if request.return_token_ids:
                         yield f"data: {json.dumps({'token_id': token_id})}\n\n"
@@ -527,6 +735,7 @@ async def chat(request: ChatRequest):
                         break
                 
                 token_count += 1
+                worker_stats["total_tokens"] += 1
             
             elapsed = time.time() - start_time
             
@@ -544,12 +753,11 @@ async def chat(request: ChatRequest):
             raise HTTPException(status_code=500, detail=f"Generation error: {str(e)}")
 
 # ============================================================================
-# Startup: Load Model
+# Model Loading
 # ============================================================================
 
 @app.on_event("startup")
 async def load_model():
-    """Load model on startup"""
     global model, tokenizer, config, eos_token_id, fast_forward
     
     print("🚀 Loading SAM-Z-1 Model...")
@@ -619,7 +827,12 @@ async def load_model():
         
         fast_forward = optimized_forward
         
-        print("✅ SAM-Z-1 Smart Worker ready! 🚀")
+        print("✅ SAM-Z-1 Distributed Worker ready! 🚀")
+        print("🔥 Features enabled:")
+        print("   - Full text generation")
+        print("   - Token-only mode (distributed pipeline)")
+        print("   - Batch decoding optimization")
+        print("   - Streaming support")
     
     except Exception as e:
         print(f"❌ Failed to load model: {e}")