Update app.py

app.py

@@ -24,12 +24,7 @@ import time
 tf.config.threading.set_inter_op_parallelism_threads(NUM_CORES)
 tf.config.threading.set_intra_op_parallelism_threads(NUM_CORES)
 
-# Enable XLA JIT compilation for CPU
-try:
-    tf.config.optimizer.set_jit(True)
-    print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN enabled, XLA JIT enabled")
-except:
-    print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN enabled")
+print(f"✅ CPU optimized: {NUM_CORES} threads, oneDNN enabled")
 
 # ============================================================================
 # 🎊 FESTIVE MODE TOGGLE 🎊
@@ -46,48 +41,48 @@ MODEL_REPO = "Smilyai-labs/Sam-large-2"
 CACHE_DIR = "./model_cache"
 
 # ============================================================================
-# Model Architecture - MATCHES CHECKPOINT STRUCTURE
+# Model Architecture Definitions (Optimized with KV-Cache)
 # ============================================================================
 
 @keras.saving.register_keras_serializable()
 class RotaryEmbedding(keras.layers.Layer):
-    """RoPE with cache built during layer build phase."""
-    
     def __init__(self, dim, max_len=2048, theta=10000, **kwargs):
         super().__init__(**kwargs)
         self.dim = dim
         self.max_len = max_len
         self.theta = theta
+        self.built_cache = False
+        self.cos_cached = None
+        self.sin_cached = None
 
     def build(self, input_shape):
-        # Pre-compute RoPE cache as numpy arrays during build
-        inv_freq = 1.0 / (self.theta ** (np.arange(0, self.dim, 2, dtype=np.float32) / self.dim))
-        t = np.arange(self.max_len, dtype=np.float32)
-        freqs = np.outer(t, inv_freq)
-        emb = np.concatenate([freqs, freqs], axis=-1)
-        
-        # Store as numpy arrays - will be converted to tensors in call()
-        self._cos_cached = np.cos(emb).astype(np.float32)
-        self._sin_cached = np.sin(emb).astype(np.float32)
-        
         super().build(input_shape)
 
+    def _build_cache(self):
+        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)
+            freqs = tf.einsum("i,j->ij", t, inv_freq)
+            emb = tf.concat([freqs, freqs], axis=-1)
+            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
+
+    def rotate_half(self, x):
+        x1, x2 = tf.split(x, 2, axis=-1)
+        return tf.concat([-x2, x1], axis=-1)
+
     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
 
-        # Slice the pre-computed values
-        cos = tf.cast(self._cos_cached[offset:offset + seq_len, :], dtype)[None, None, :, :]
-        sin = tf.cast(self._sin_cached[offset:offset + seq_len, :], dtype)[None, None, :, :]
+        cos = tf.cast(self.cos_cached[offset:offset + seq_len, :], dtype)[None, None, :, :]
+        sin = tf.cast(self.sin_cached[offset:offset + seq_len, :], dtype)[None, None, :, :]
 
-        # Fused rotate_half
-        x1_q, x2_q = tf.split(q, 2, axis=-1)
-        x1_k, x2_k = tf.split(k, 2, axis=-1)
-        
-        q_embed = (q * cos) + (tf.concat([-x2_q, x1_q], axis=-1) * sin)
-        k_embed = (k * cos) + (tf.concat([-x2_k, x1_k], axis=-1) * sin)
-        
+        q_embed = (q * cos) + (self.rotate_half(q) * sin)
+        k_embed = (k * cos) + (self.rotate_half(k) * sin)
         return q_embed, k_embed
 
     def get_config(self):
@@ -119,8 +114,6 @@ class RMSNorm(keras.layers.Layer):
 
 @keras.saving.register_keras_serializable()
 class TransformerBlock(keras.layers.Layer):
-    """Transformer block - MATCHES ORIGINAL CHECKPOINT STRUCTURE."""
-    
     def __init__(self, d_model, n_heads, ff_dim, dropout, max_len, rope_theta, layer_idx=0, **kwargs):
         super().__init__(**kwargs)
         self.d_model = d_model
@@ -131,37 +124,38 @@ class TransformerBlock(keras.layers.Layer):
         self.rope_theta = rope_theta
         self.head_dim = d_model // n_heads
         self.layer_idx = layer_idx
-        self.scale = 1.0 / np.sqrt(self.head_dim)
 
     def build(self, input_shape):
-        # MUST use same layer names as checkpoint
         self.pre_attn_norm = RMSNorm(name="pre_attn_norm")
         self.pre_ffn_norm = RMSNorm(name="pre_ffn_norm")
-        
-        # Separate Q, K, V projections (matches checkpoint)
         self.q_proj = keras.layers.Dense(self.d_model, use_bias=False, name="q_proj")
         self.k_proj = keras.layers.Dense(self.d_model, use_bias=False, name="k_proj")
         self.v_proj = keras.layers.Dense(self.d_model, use_bias=False, name="v_proj")
         self.out_proj = keras.layers.Dense(self.d_model, use_bias=False, name="o_proj")
-        
         self.rope = RotaryEmbedding(self.head_dim, max_len=self.max_len, theta=self.rope_theta)
-        
-        # Separate gate, up, down projections (matches checkpoint)
         self.gate_proj = keras.layers.Dense(self.ff_dim, use_bias=False, name="gate_proj")
         self.up_proj = keras.layers.Dense(self.ff_dim, use_bias=False, name="up_proj")
         self.down_proj = keras.layers.Dense(self.d_model, use_bias=False, name="down_proj")
-        
         self.dropout = keras.layers.Dropout(self.dropout_rate)
         super().build(input_shape)
 
     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
 
         res = x
         y = self.pre_attn_norm(x)
 
-        # Separate Q, K, V projections
+        # 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])  # [B, n_heads, T, head_dim]
 
@@ -187,14 +181,15 @@ class TransformerBlock(keras.layers.Layer):
 
         new_kv = (k, v) if use_cache else None
 
-        # Scaled dot-product attention
+        # Attention
         full_len = tf.shape(k)[2]
-        scores = tf.matmul(q, k, transpose_b=True) * self.scale
+        scores = tf.matmul(q, k, transpose_b=True) / tf.sqrt(tf.cast(self.head_dim, dtype))
 
         # Causal mask
         q_positions = tf.range(past_len, past_len + T)
         k_positions = tf.range(full_len)
-        mask = tf.cast(q_positions[:, None] < k_positions[None, :], scores.dtype) * -1e9
+        mask = tf.cast(q_positions[:, None] >= k_positions[None, :], dtype)
+        mask = tf.where(mask == 0, tf.constant(-1e9, dtype=dtype), tf.constant(0.0, dtype=dtype))
         scores = scores + mask[None, None, :, :]
 
         attn = tf.nn.softmax(scores, axis=-1)
@@ -204,10 +199,10 @@ class TransformerBlock(keras.layers.Layer):
 
         x = res + self.dropout(self.out_proj(attn_out), training=training)
 
-        # FFN with SwiGLU
+        # FFN
         res = x
         y = self.pre_ffn_norm(x)
-        ffn = self.down_proj(tf.nn.silu(self.gate_proj(y)) * self.up_proj(y))
+        ffn = self.down_proj(keras.activations.silu(self.gate_proj(y)) * self.up_proj(y))
         output = res + self.dropout(ffn, training=training)
 
         return output, new_kv
@@ -255,6 +250,14 @@ class SAM1Model(keras.Model):
         self.lm_head = keras.layers.Dense(self.cfg['vocab_size'], use_bias=False, name="lm_head")
 
     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
@@ -274,62 +277,6 @@ class SAM1Model(keras.Model):
         return base_config
 
 
-# ============================================================================
-# Optimized Sampling
-# ============================================================================
-
-class FastSampler:
-    """Vectorized sampler for faster token selection."""
-    
-    def __init__(self, vocab_size):
-        self.vocab_size = vocab_size
-        self.rng = np.random.default_rng()
-    
-    def sample(self, logits, temperature, top_k, top_p, token_freq, repetition_penalty):
-        """Optimized sampling with vectorized operations."""
-        logits = logits.copy()
-        
-        if temperature != 1.0:
-            logits = logits / temperature
-
-        # Vectorized repetition penalty
-        if repetition_penalty != 1.0 and token_freq:
-            freq_tokens = np.array(list(token_freq.keys()), dtype=np.int32)
-            freq_values = np.array(list(token_freq.values()), dtype=np.float32)
-            valid_mask = freq_tokens < len(logits)
-            freq_tokens = freq_tokens[valid_mask]
-            freq_values = freq_values[valid_mask]
-            if len(freq_tokens) > 0:
-                logits[freq_tokens] /= np.power(repetition_penalty, freq_values)
-
-        # Top-K filtering with partial sort
-        if 0 < top_k < len(logits):
-            top_k_indices = np.argpartition(logits, -top_k)[-top_k:]
-            top_k_logits = logits[top_k_indices]
-        else:
-            top_k_indices = np.arange(len(logits))
-            top_k_logits = logits
-
-        # Stable softmax
-        top_k_logits = top_k_logits - np.max(top_k_logits)
-        exp_logits = np.exp(top_k_logits)
-        top_k_probs = exp_logits / exp_logits.sum()
-
-        # 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
-            nucleus_idx = sorted_idx[:cutoff]
-            nucleus_probs = top_k_probs[nucleus_idx]
-            nucleus_probs /= nucleus_probs.sum()
-            sampled = self.rng.choice(len(nucleus_probs), p=nucleus_probs)
-            return int(top_k_indices[nucleus_idx[sampled]])
-        else:
-            sampled = self.rng.choice(len(top_k_probs), p=top_k_probs)
-            return int(top_k_indices[sampled])
-
-
 # --- Model and Tokenizer Loading ---
 
 config_path = hf_hub_download(MODEL_REPO, "config.json", cache_dir=CACHE_DIR)
@@ -375,7 +322,7 @@ 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,  # Disable dropout for inference
+        'dropout': 0.1,
         'rope_theta': config['rope_theta']
     }
     model = SAM1Model(config=model_config)
@@ -409,77 +356,56 @@ else:
 if model:
     print(f"✅ Model loaded: {config['num_hidden_layers']} layers, {config['vocab_size']} vocab")
 
-# Initialize fast sampler
-sampler = FastSampler(config['vocab_size'])
-
-# Warm up the model (without tf.function first)
+# Warm up the model
 print("🔥 Warming up model...")
 warmup_input = tf.constant([[1, 2, 3, 4, 5]], dtype=tf.int32)
-
-# Initial warmup to build all internal caches
-for _ in range(2):
-    logits, past_kv = model(warmup_input, training=False, use_cache=True)
-
-# Warmup decode step
-single_token = tf.constant([[1]], dtype=tf.int32)
-for _ in range(2):
-    logits, past_kv = model(single_token, training=False, past_kv=past_kv, use_cache=True)
-
+_, _ = model(warmup_input, training=False, use_cache=True)
 print("✅ Model warmed up")
 
-
 # ============================================================================
-# Inference wrapper class for clean tf.function usage
+# Optimized Inference Logic with KV-Cache
 # ============================================================================
 
-class InferenceEngine:
-    """Wrapper for compiled inference functions."""
-    
-    def __init__(self, model):
-        self.model = model
-        self._prefill_fn = None
-        self._decode_fn = None
-        
-    def prefill(self, input_ids):
-        """Run prefill (first call builds trace)."""
-        if self._prefill_fn is None:
-            # First call - run eagerly to ensure all caches are built
-            return self.model(input_ids, training=False, use_cache=True)
-        return self._prefill_fn(input_ids)
-    
-    def decode(self, input_ids, past_kv):
-        """Run single-token decode."""
-        return self.model(input_ids, training=False, past_kv=past_kv, use_cache=True)
-    
-    def compile_traces(self):
-        """Compile tf.function traces after warmup."""
-        print("🔥 Compiling optimized traces...")
-        
-        @tf.function(reduce_retracing=True)
-        def prefill_fn(input_ids):
-            return self.model(input_ids, training=False, use_cache=True)
-        
-        self._prefill_fn = prefill_fn
-        
-        # Trace with sample inputs
-        sample_input = tf.constant([[1, 2, 3, 4, 5]], dtype=tf.int32)
-        _ = self._prefill_fn(sample_input)
-        
-        print("✅ Traces compiled")
-
-
-# Create inference engine
-engine = InferenceEngine(model)
+stop_generation = False
 
-# Compile traces after warmup
-engine.compile_traces()
 
+def sample_token(logits, temperature, top_k, top_p, token_freq, repetition_penalty):
+    """Pure NumPy sampling for speed."""
+    # Temperature scaling
+    scaled_logits = logits / temperature
 
-# ============================================================================
-# Optimized Inference Logic with KV-Cache
-# ============================================================================
+    # Repetition penalty
+    if repetition_penalty != 1.0:
+        for token_id, freq in token_freq.items():
+            if token_id < len(scaled_logits):
+                scaled_logits[token_id] /= (repetition_penalty ** freq)
 
-stop_generation = False
+    # 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(scaled_logits))
+        top_k_logits = scaled_logits
+
+    # Softmax (numerically stable)
+    top_k_logits = top_k_logits - np.max(top_k_logits)
+    top_k_probs = np.exp(top_k_logits)
+    top_k_probs /= top_k_probs.sum()
+
+    # 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
+        nucleus_idx = sorted_idx[:cutoff]
+        nucleus_probs = top_k_probs[nucleus_idx]
+        nucleus_probs /= nucleus_probs.sum()
+        sampled = np.random.choice(len(nucleus_probs), p=nucleus_probs)
+        return int(top_k_indices[nucleus_idx[sampled]])
+    else:
+        sampled = np.random.choice(len(top_k_probs), p=top_k_probs)
+        return int(top_k_indices[sampled])
 
 
 def generate_stream(
@@ -514,16 +440,17 @@ def generate_stream(
 
     max_context = config['max_position_embeddings']
 
-    start_time = time.perf_counter()
+    start_time = time.time()
 
     # === 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 = engine.prefill(input_tensor)
+        logits, past_kv = model(input_tensor, training=False, use_cache=True)
     except Exception as e:
         yield f"Error during prefill: {e}"
         return
@@ -531,12 +458,11 @@ def generate_stream(
     # Get logits for last position
     next_token_logits = logits[0, -1, :].numpy()
 
-    prefill_time = time.perf_counter() - start_time
-    prefill_tps = len(input_ids) / prefill_time if prefill_time > 0 else 0
-    print(f"⚡ Prefill: {len(input_ids)} tokens in {prefill_time:.3f}s ({prefill_tps:.1f} tok/s)")
+    prefill_time = time.time() - start_time
+    print(f"⚡ Prefill: {len(input_ids)} tokens in {prefill_time:.2f}s")
 
     # === GENERATION LOOP ===
-    decode_start = time.perf_counter()
+    decode_start = time.time()
     
     for step in range(max_tokens):
         if stop_generation:
@@ -544,7 +470,7 @@ def generate_stream(
             return
 
         # Sample next token
-        next_token_id = sampler.sample(
+        next_token_id = sample_token(
             next_token_logits, temperature, top_k, top_p, token_freq, repetition_penalty
         )
 
@@ -561,11 +487,11 @@ def generate_stream(
         token_count += 1
         yield generated_text
 
-        # === DECODE PHASE ===
+        # === DECODE PHASE (single token, reuse cache) ===
         next_input = tf.constant([[next_token_id]], dtype=tf.int32)
         
         try:
-            logits, past_kv = engine.decode(next_input, past_kv)
+            logits, past_kv = model(next_input, training=False, past_kv=past_kv, use_cache=True)
         except Exception as e:
             yield generated_text + f"\n\n*[Error during generation: {e}]*"
             return
@@ -573,24 +499,24 @@ def generate_stream(
         next_token_logits = logits[0, -1, :].numpy()
 
         # Truncate cache if too long
-        if step % 100 == 99:
-            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
-                past_kv = [
-                    (k[:, :, trim_amount:, :], v[:, :, trim_amount:, :])
-                    for k, v in past_kv
-                ]
-
-    decode_time = time.perf_counter() - decode_start
-    total_time = time.perf_counter() - start_time
+        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
+            ]
+
+    decode_time = time.time() - decode_start
+    total_time = time.time() - start_time
     
     if token_count > 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*[Generated {token_count} tokens in {total_time:.1f}s "
-            f"(prefill: {prefill_time:.2f}s, decode: {decode_tps:.1f} tok/s)]*"
+            f"(prefill: {prefill_time:.1f}s, decode: {decode_tps:.1f} tok/s)]*"
         )
         
         if not stop_generation:
@@ -605,21 +531,20 @@ def generate_stream(
 
 def format_chat_prompt(message: str, history: list, reasoning_enabled: bool) -> str:
     """Format message history and seed <think> if enabled."""
-    prompt_parts = []
-    
+    prompt = ""
     for user_msg, assistant_msg in history:
-        prompt_parts.append(f"<|im_start|>user\n{user_msg}<|im_end|>")
+        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_parts.append(f"<|im_start|>assistant\n{clean_msg}<|im_end|>")
+            prompt += f"<|im_start|>assistant\n{clean_msg}<|im_end|>\n"
+
+    prompt += f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
 
-    prompt_parts.append(f"<|im_start|>user\n{message}<|im_end|>")
-    prompt_parts.append("<|im_start|>assistant")
-    
     if reasoning_enabled:
-        prompt_parts.append("<think>")
+        prompt += "<think>"
 
-    return "\n".join(prompt_parts)
+    return prompt
 
 
 def chat_stream(
@@ -658,6 +583,7 @@ 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:]
@@ -843,7 +769,7 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
 **Vocab:** {config['vocab_size']:,}
 **Layers:** {config['num_hidden_layers']}
 **Context:** {config['max_position_embeddings']:,} tokens
-**Optimization:** KV-Cache + XLA ⚡
+**Optimization:** KV-Cache enabled ⚡
 """)
 
     gr.Examples(
@@ -919,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 with Optimized Inference")
+    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)