Update app.py

app.py

@@ -7,12 +7,9 @@ NUM_CORES = os.cpu_count() or 4
 
 os.environ['TF_NUM_INTEROP_THREADS'] = str(NUM_CORES)
 os.environ['TF_NUM_INTRAOP_THREADS'] = str(NUM_CORES)
-os.environ['CUDA_VISIBLE_DEVICES'] = '-1'
-os.environ['TF_ENABLE_ONEDNN_OPTS'] = '1'
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
-# NEW: Enable XLA (Accelerated Linear Algebra)
-os.environ['TF_XLA_FLAGS'] = '--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit'
-os.environ['TF_ENABLE_AUTO_MIXED_PRECISION'] = '1'
+os.environ['CUDA_VISIBLE_DEVICES'] = '-1'  # Force CPU only
+os.environ['TF_ENABLE_ONEDNN_OPTS'] = '1'  # Intel optimization
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'   # Reduce TF logging
 
 import gradio as gr
 import tensorflow as tf
@@ -23,13 +20,11 @@ from tokenizers import Tokenizer
 import numpy as np
 import time
 
-# Configure TF for maximum performance
+# Configure TF threading
 tf.config.threading.set_inter_op_parallelism_threads(NUM_CORES)
 tf.config.threading.set_intra_op_parallelism_threads(NUM_CORES)
-tf.config.optimizer.set_jit(True)  # Enable XLA JIT compilation
-tf.config.optimizer.set_experimental_options({'auto_mixed_precision': True})
 
-print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN + XLA enabled")
+print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN enabled")
 
 # ============================================================================
 # 🎊 FESTIVE MODE TOGGLE 🎊
@@ -62,8 +57,6 @@ class RotaryEmbedding(keras.layers.Layer):
 
     def build(self, input_shape):
         super().build(input_shape)
-        # Pre-build cache immediately
-        self._build_cache()
 
     def _build_cache(self):
         if not self.built_cache:
@@ -71,19 +64,17 @@ class RotaryEmbedding(keras.layers.Layer):
             t = tf.range(self.max_len, dtype=tf.float32)
             freqs = tf.einsum("i,j->ij", t, inv_freq)
             emb = tf.concat([freqs, freqs], axis=-1)
-            # Use tf.constant for faster access
             self.cos_cached = tf.constant(np.cos(emb.numpy()), dtype=tf.float32)
             self.sin_cached = tf.constant(np.sin(emb.numpy()), dtype=tf.float32)
             self.built_cache = True
 
-    @tf.function(jit_compile=True)
     def rotate_half(self, x):
         x1, x2 = tf.split(x, 2, axis=-1)
         return tf.concat([-x2, x1], axis=-1)
 
-    @tf.function(jit_compile=True)
     def call(self, q, k, offset=0):
         """Apply rotary embeddings with position offset for KV-cache."""
+        self._build_cache()
         seq_len = tf.shape(q)[2]
         dtype = q.dtype
 
@@ -111,7 +102,6 @@ class RMSNorm(keras.layers.Layer):
         self.scale = self.add_weight(name="scale", shape=(input_shape[-1],), initializer="ones")
         super().build(input_shape)
 
-    @tf.function(jit_compile=True)
     def call(self, x):
         variance = tf.reduce_mean(tf.square(x), axis=-1, keepdims=True)
         return x * tf.math.rsqrt(variance + self.epsilon) * self.scale
@@ -149,8 +139,15 @@ class TransformerBlock(keras.layers.Layer):
         self.dropout = keras.layers.Dropout(self.dropout_rate)
         super().build(input_shape)
 
-    # Removed @tf.function here to allow flexible past_kv handling
     def call(self, x, training=None, past_kv=None, use_cache=False):
+        """
+        Args:
+            x: input tensor [B, T, D] (T=1 during cached generation)
+            past_kv: tuple of (past_k, past_v) each [B, n_heads, past_len, head_dim]
+            use_cache: whether to return updated kv cache
+        Returns:
+            output, (new_k, new_v) if use_cache else output, None
+        """
         B = tf.shape(x)[0]
         T = tf.shape(x)[1]
         dtype = x.dtype
@@ -158,9 +155,9 @@ class TransformerBlock(keras.layers.Layer):
         res = x
         y = self.pre_attn_norm(x)
 
-        # Project Q, K, V
+        # Project Q, K, V for current input
         q = tf.reshape(self.q_proj(y), [B, T, self.n_heads, self.head_dim])
-        q = tf.transpose(q, [0, 2, 1, 3])
+        q = tf.transpose(q, [0, 2, 1, 3])  # [B, n_heads, T, head_dim]
 
         k = tf.reshape(self.k_proj(y), [B, T, self.n_heads, self.head_dim])
         k = tf.transpose(k, [0, 2, 1, 3])
@@ -168,13 +165,13 @@ class TransformerBlock(keras.layers.Layer):
         v = tf.reshape(self.v_proj(y), [B, T, self.n_heads, self.head_dim])
         v = tf.transpose(v, [0, 2, 1, 3])
 
-        # Position offset
+        # Determine position offset for RoPE
         if past_kv is not None:
             past_len = tf.shape(past_kv[0])[2]
         else:
             past_len = 0
 
-        # Apply RoPE
+        # Apply RoPE with position offset
         q, k = self.rope(q, k, offset=past_len)
 
         # Concatenate with past KV
@@ -200,13 +197,13 @@ class TransformerBlock(keras.layers.Layer):
         attn_out = tf.transpose(attn_out, [0, 2, 1, 3])
         attn_out = tf.reshape(attn_out, [B, T, self.d_model])
 
-        x = res + self.out_proj(attn_out)
+        x = res + self.dropout(self.out_proj(attn_out), training=training)
 
         # FFN
         res = x
         y = self.pre_ffn_norm(x)
         ffn = self.down_proj(keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
-        output = res + ffn
+        output = res + self.dropout(ffn, training=training)
 
         return output, new_kv
 
@@ -252,8 +249,15 @@ class SAM1Model(keras.Model):
         self.norm = RMSNorm(name="final_norm")
         self.lm_head = keras.layers.Dense(self.cfg['vocab_size'], use_bias=False, name="lm_head")
 
-    # Don't use @tf.function on the main call to allow flexible caching
     def call(self, input_ids, training=None, past_kv=None, use_cache=False):
+        """
+        Args:
+            input_ids: [B, T]
+            past_kv: list of (k, v) tuples, one per layer
+            use_cache: whether to return updated cache
+        Returns:
+            logits, new_past_kv (or None)
+        """
         x = self.embed(input_ids)
 
         new_past_kv = [] if use_cache else None
@@ -318,11 +322,12 @@ if use_checkpoint:
         'n_heads': config['num_attention_heads'],
         'ff_mult': config['intermediate_size'] / config['hidden_size'],
         'max_len': config['max_position_embeddings'],
-        'dropout': 0.0,  # Set to 0 for inference
+        'dropout': 0.1,
         'rope_theta': config['rope_theta']
     }
     model = SAM1Model(config=model_config)
     
+    # Build model with dummy input
     dummy_input = tf.zeros((1, 16), dtype=tf.int32)
     _ = model(dummy_input, training=False, use_cache=False)
     print(f"✅ Model architecture built: {model.count_params():,} parameters")
@@ -351,37 +356,11 @@ else:
 if model:
     print(f"✅ Model loaded: {config['num_hidden_layers']} layers, {config['vocab_size']} vocab")
 
-# ============================================================================
-# OPTIMIZATION: Create compiled inference functions
-# ============================================================================
-
-print("⚡ Compiling optimized inference functions...")
-
-# Create traced functions for prefill and decode
-@tf.function(
-    input_signature=[tf.TensorSpec(shape=[1, None], dtype=tf.int32)],
-    jit_compile=True
-)
-def prefill_fn(input_ids):
-    """Optimized prefill phase - processes entire prompt at once."""
-    return model(input_ids, training=False, use_cache=True)
-
-@tf.function(
-    input_signature=[tf.TensorSpec(shape=[1, 1], dtype=tf.int32)],
-    jit_compile=True  
-)
-def decode_fn_no_cache(input_ids):
-    """Fast decode for single token without cache (fallback)."""
-    logits, _ = model(input_ids, training=False, use_cache=False)
-    return logits
-
-# Warm up with compilation
-print("🔥 Warming up and compiling model...")
+# Warm up the model
+print("🔥 Warming up model...")
 warmup_input = tf.constant([[1, 2, 3, 4, 5]], dtype=tf.int32)
-_ = prefill_fn(warmup_input)
-warmup_single = tf.constant([[1]], dtype=tf.int32)
-_ = decode_fn_no_cache(warmup_single)
-print("✅ Model warmed up and compiled")
+_, _ = model(warmup_input, training=False, use_cache=True)
+print("✅ Model warmed up")
 
 # ============================================================================
 # Optimized Inference Logic with KV-Cache
@@ -389,50 +368,39 @@ print("✅ Model warmed up and compiled")
 
 stop_generation = False
 
-# Pre-compile sampling helper
-@tf.function(jit_compile=True)
-def apply_temperature_tf(logits, temperature):
-    """GPU-accelerated temperature scaling."""
-    return logits / temperature
 
 def sample_token(logits, temperature, top_k, top_p, token_freq, repetition_penalty):
-    """Optimized NumPy sampling."""
-    # Temperature
-    if temperature != 1.0:
-        logits = logits / temperature
-
-    # Repetition penalty (vectorized)
-    if repetition_penalty != 1.0 and token_freq:
-        penalty = np.ones_like(logits)
+    """Pure NumPy sampling for speed."""
+    # Temperature scaling
+    scaled_logits = logits / temperature
+
+    # Repetition penalty
+    if repetition_penalty != 1.0:
         for token_id, freq in token_freq.items():
-            if token_id < len(logits):
-                penalty[token_id] = repetition_penalty ** freq
-        logits = logits / penalty
-
-    # Top-K
-    if top_k > 0 and top_k < len(logits):
-        top_k_indices = np.argpartition(logits, -top_k)[-top_k:]
-        top_k_logits = logits[top_k_indices]
+            if token_id < len(scaled_logits):
+                scaled_logits[token_id] /= (repetition_penalty ** freq)
+
+    # Top-K filtering
+    if top_k > 0 and top_k < len(scaled_logits):
+        top_k_indices = np.argpartition(scaled_logits, -top_k)[-top_k:]
+        top_k_logits = scaled_logits[top_k_indices]
     else:
-        top_k_indices = np.arange(len(logits))
-        top_k_logits = logits
+        top_k_indices = np.arange(len(scaled_logits))
+        top_k_logits = scaled_logits
 
-    # Softmax (stable)
-    max_logit = np.max(top_k_logits)
-    top_k_logits = top_k_logits - max_logit
+    # Softmax (numerically stable)
+    top_k_logits = top_k_logits - np.max(top_k_logits)
     top_k_probs = np.exp(top_k_logits)
-    sum_probs = np.sum(top_k_probs)
-    top_k_probs = top_k_probs / sum_probs
+    top_k_probs /= top_k_probs.sum()
 
-    # Top-P
+    # Top-P (nucleus) filtering
     if top_p < 1.0:
         sorted_idx = np.argsort(top_k_probs)[::-1]
         cumsum = np.cumsum(top_k_probs[sorted_idx])
         cutoff = np.searchsorted(cumsum, top_p) + 1
-        cutoff = min(cutoff, len(sorted_idx))
         nucleus_idx = sorted_idx[:cutoff]
         nucleus_probs = top_k_probs[nucleus_idx]
-        nucleus_probs = nucleus_probs / np.sum(nucleus_probs)
+        nucleus_probs /= nucleus_probs.sum()
         sampled = np.random.choice(len(nucleus_probs), p=nucleus_probs)
         return int(top_k_indices[nucleus_idx[sampled]])
     else:
@@ -448,23 +416,23 @@ def generate_stream(
     top_p: float = 0.9,
     repetition_penalty: float = 1.1
 ):
-    """Optimized generation with compiled functions."""
+    """Generate text with KV-cache for fast CPU inference."""
     global stop_generation
     stop_generation = False
 
-    # Tokenize
+    # Tokenize prompt
     prompt_ids = tokenizer.encode(prompt).ids
     input_ids = [i for i in prompt_ids if i != eos_token_id]
 
     if len(input_ids) == 0:
-        yield "Error: Empty prompt"
+        yield "Error: Empty prompt after tokenization"
         return
 
     generated_text = ""
     token_count = 0
     token_freq = {}
 
-    # Stop tokens
+    # Get special token IDs
     im_end_id = tokenizer.token_to_id("<|im_end|>")
     model_end_id = tokenizer.token_to_id("<im end for model tun>")
     stop_ids = {eos_token_id, im_end_id, model_end_id}
@@ -474,81 +442,81 @@ def generate_stream(
 
     start_time = time.time()
 
-    # === PREFILL (optimized) ===
+    # === PREFILL PHASE ===
+    # Truncate if prompt is too long
     if len(input_ids) > max_context - max_tokens:
         input_ids = input_ids[-(max_context - max_tokens):]
 
     input_tensor = tf.constant([input_ids], dtype=tf.int32)
     
     try:
-        logits, past_kv = prefill_fn(input_tensor)
-        logits_np = logits[0, -1, :].numpy()
+        logits, past_kv = model(input_tensor, training=False, use_cache=True)
     except Exception as e:
-        yield f"Error: {e}"
+        yield f"Error during prefill: {e}"
         return
 
+    # Get logits for last position
+    next_token_logits = logits[0, -1, :].numpy()
+
     prefill_time = time.time() - start_time
-    print(f"⚡ Prefill: {len(input_ids)} tokens in {prefill_time:.3f}s ({len(input_ids)/prefill_time:.1f} tok/s)")
+    print(f"⚡ Prefill: {len(input_ids)} tokens in {prefill_time:.2f}s")
 
-    # === DECODE LOOP ===
+    # === GENERATION LOOP ===
     decode_start = time.time()
-    decode_times = []
     
     for step in range(max_tokens):
         if stop_generation:
-            yield generated_text + "\n\n*[Stopped]*"
+            yield generated_text + "\n\n*[Generation stopped]*"
             return
 
-        # Sample
-        token_start = time.time()
-        next_token_id = sample_token(logits_np, temperature, top_k, top_p, token_freq, repetition_penalty)
-        
+        # Sample next token
+        next_token_id = sample_token(
+            next_token_logits, temperature, top_k, top_p, token_freq, repetition_penalty
+        )
+
+        # Stop conditions
         if next_token_id in stop_ids:
             break
 
+        # Update frequency tracking
         token_freq[next_token_id] = token_freq.get(next_token_id, 0) + 1
 
-        # Decode
+        # Decode and yield
         token_text = tokenizer.decode([next_token_id])
         generated_text += token_text
         token_count += 1
-        
-        # Yield every token for streaming
-        if token_count % 1 == 0:  # Stream every token
-            yield generated_text
+        yield generated_text
 
-        # === OPTIMIZED DECODE ===
+        # === DECODE PHASE (single token, reuse cache) ===
         next_input = tf.constant([[next_token_id]], dtype=tf.int32)
         
         try:
             logits, past_kv = model(next_input, training=False, past_kv=past_kv, use_cache=True)
-            logits_np = logits[0, -1, :].numpy()
         except Exception as e:
-            yield generated_text + f"\n\n*[Error: {e}]*"
+            yield generated_text + f"\n\n*[Error during generation: {e}]*"
             return
-        
-        decode_times.append(time.time() - token_start)
-
-        # Cache management
-        if past_kv and len(past_kv) > 0:
-            current_len = past_kv[0][0].shape[2]
-            if current_len > max_context:
-                trim = current_len - max_context + 100
-                past_kv = [(k[:, :, trim:, :], v[:, :, trim:, :]) for k, v in past_kv]
+            
+        next_token_logits = logits[0, -1, :].numpy()
+
+        # Truncate cache if too long
+        current_len = past_kv[0][0].shape[2] if past_kv and past_kv[0] is not None else 0
+        if current_len > max_context:
+            trim_amount = current_len - max_context + 100  # Keep some buffer
+            past_kv = [
+                (k[:, :, trim_amount:, :], v[:, :, trim_amount:, :])
+                for k, v in past_kv
+            ]
 
-    # Stats
     decode_time = time.time() - decode_start
     total_time = time.time() - start_time
     
     if token_count > 0:
-        avg_decode_time = np.mean(decode_times) if decode_times else 0
         decode_tps = token_count / decode_time if decode_time > 0 else 0
+        total_tps = token_count / total_time if total_time > 0 else 0
         
         stats = (
-            f"\n\n*[{token_count} tokens in {total_time:.1f}s | "
-            f"Prefill: {len(input_ids)/prefill_time:.1f} t/s | "
-            f"Decode: {decode_tps:.1f} t/s | "
-            f"Avg: {1/avg_decode_time:.1f} t/s]*"
+            f"\n\n*[Generated {token_count} tokens in {total_time:.1f}s "
+            f"(prefill: {prefill_time:.1f}s, decode: {decode_tps:.1f} tok/s)]*"
         )
         
         if not stop_generation:
@@ -558,7 +526,7 @@ def generate_stream(
 
 
 # ============================================================================
-# Chat Interface Logic (unchanged from original)
+# Chat Interface Logic
 # ============================================================================
 
 def format_chat_prompt(message: str, history: list, reasoning_enabled: bool) -> str:
@@ -567,6 +535,7 @@ def format_chat_prompt(message: str, history: list, reasoning_enabled: bool) ->
     for user_msg, assistant_msg in history:
         prompt += f"<|im_start|>user\n{user_msg}<|im_end|>\n"
         if assistant_msg:
+            # Clean up any stats from previous messages
             clean_msg = assistant_msg.split("\n\n*[")[0]
             prompt += f"<|im_start|>assistant\n{clean_msg}<|im_end|>\n"
 
@@ -600,6 +569,7 @@ def chat_stream(
     ):
         partial_response = generated
 
+        # Robust end-of-turn detection
         stop_tags = ["<|im_end|>", "<im end for model tun>"]
         earliest_stop = len(partial_response)
         should_stop = False
@@ -613,12 +583,14 @@ def chat_stream(
 
         display_response = partial_response
         if should_stop:
+            # Keep the stats portion if present
             stats_start = partial_response.find("\n\n*[")
             if stats_start > earliest_stop:
                 display_response = partial_response[:earliest_stop] + partial_response[stats_start:]
             else:
                 display_response = partial_response[:earliest_stop]
 
+        # Post-process reasoning tags for display
         if reasoning_enabled:
             if '<think>' in display_response and '</think>' in display_response:
                 start_idx = display_response.find('<think>')
@@ -650,7 +622,7 @@ def stop_gen():
 
 
 # ============================================================================
-# Gradio UI (same as original, just using optimized backend)
+# Gradio UI
 # ============================================================================
 
 custom_css = """
@@ -725,15 +697,15 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
         <div id="welcome-modal" class="modal-overlay" style="display:none;">
             <div class="modal-content">
                 <h2>🧠 Welcome to Sam-large-2: Dual-Mode Reasoning Demo</h2>
-                <p>⚡ <strong>OPTIMIZED VERSION</strong> with XLA JIT compilation for 5-10x faster CPU inference!</p>
+                <p>Our latest model features <strong>Chain-of-Thought (CoT)</strong> functionality and <strong>KV-Cache optimization</strong> for fast CPU inference!</p>
                 <div class="comparison-box">
                     <div class="comparison-mode mode-reasoning">
                         <h3>💡 Reasoning Mode (ON)</h3>
-                        <p>Chain-of-Thought with <code>&lt;think>...&lt;/think></code> tags.</p>
+                        <p>The model performs a <strong>CoT step</strong> first. The internal thought process is contained within <code>&lt;think>...&lt;/think></code> tags.</p>
                     </div>
                     <div class="comparison-mode mode-direct">
                         <h3>⚪ Direct Mode (OFF)</h3>
-                        <p>Immediate answer generation for maximum speed.</p>
+                        <p>The model generates the final answer immediately, maximizing speed.</p>
                     </div>
                 </div>
                 <button class="close-btn" onclick="document.getElementById('welcome-modal').style.display='none'">Got it! Start Chatting</button>
@@ -749,13 +721,13 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                 <img src="https://cdn-uploads.huggingface.co/production/uploads/64e3486b82fb6ae7a06c749c/yBUDdaTze1L84NaDSpZGf.jpeg"
                     alt="Sam-large-2" style="max-width: 400px; border-radius: 12px; margin: 1rem auto; display: block; box-shadow: 0 8px 32px rgba(240, 147, 251, 0.3);">
                 <h1>🤖 Sam-large-2 Chat 🤖</h1>
-                <p><strong>XLA-OPTIMIZED!</strong> Our <strong>FASTEST Reasoning Model</strong> <span class="speed-indicator">⚡ 10-30 tok/s on CPU</span></p>
-                <div class="twin-badge">Reasoning Model + XLA JIT</div>
+                <p><strong>LATEST RELEASE!</strong> Our <strong>BEST Reasoning Model</strong> - Now with KV-Cache! <span class="speed-indicator">⚡ 5-20x Faster</span></p>
+                <div class="twin-badge">Reasoning Model</div>
                 <div class="celebration">🚀 💫 🎯 ⚡ 🔥</div>
             </div>
         """)
     else:
-        gr.HTML("""<div class="header"><h1>🤖 Sam-large-2 Chat</h1><p>XLA-Optimized Reasoning Model</p></div>""")
+        gr.HTML("""<div class="header"><h1>🤖 Sam-large-2 Chat</h1><p>Advanced Reasoning Model with KV-Cache</p></div>""")
 
     with gr.Row():
         with gr.Column(scale=4):
@@ -793,12 +765,11 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
             repetition_penalty = gr.Slider(minimum=1.0, maximum=2.0, value=1.1, step=0.1, label="Repetition Penalty")
             gr.Markdown("---")
             gr.Markdown(f"""### 🎊 Sam-large-2 Model Info
-**Type:** XLA-Optimized CoT Model
+**Type:** Chain-of-Thought Reasoning Model
 **Vocab:** {config['vocab_size']:,}
 **Layers:** {config['num_hidden_layers']}
 **Context:** {config['max_position_embeddings']:,} tokens
-**Optimization:** KV-Cache + XLA JIT ⚡
-**Expected Speed:** 10-30 tok/s (CPU)
+**Optimization:** KV-Cache enabled ⚡
 """)
 
     gr.Examples(
@@ -813,8 +784,8 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
 
     gr.HTML("""
         <footer>
-            <p><strong>⚡ Sam-large-2 - XLA-OPTIMIZED RELEASE! ⚡</strong></p>
-            <p style="font-size: 0.9rem; color: #999;">5-10x faster with XLA JIT compilation • KV-Cache enabled • Built by Smily studios</p>
+            <p><strong>🎉 Sam-large-2 - LATEST RELEASE with KV-Cache! 🎉</strong></p>
+            <p style="font-size: 0.9rem; color: #999;">Trained from scratch on TPU v5e-8 • Built by Smily studios with TensorFlow & Gradio</p>
         </footer>
     """)
 
@@ -874,7 +845,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
 
 if __name__ == "__main__":
     print("\n" + "=" * 60)
-    print("⚡ Starting Sam-large-2 Chat (XLA-OPTIMIZED)")
+    print("🚀 Starting Sam-large-2 Chat with KV-Cache Optimization")
     print("=" * 60 + "\n")
     demo.queue(max_size=20)
     demo.launch(server_name="0.0.0.0", server_port=7860, share=False, show_error=True)