Update app.py

app.py

@@ -461,8 +461,9 @@ class ModelWrapper:
         
         print(f"✅ Model loaded: {self.d_model}d × {self.n_layers}L × {self.n_heads}H")
     
-    def generate(self, prompt: str, max_new_tokens: int = 200, 
-                 temperature: float = 0.8, top_k: int = 50, top_p: float = 0.9):
+    def generate_stream(self, prompt: str, max_new_tokens: int = 200, 
+                        temperature: float = 0.8, top_k: int = 50, top_p: float = 0.9):
+        """Generator that yields tokens one at a time for streaming"""
         # Format prompt correctly (NO newline between User: and Sam:)
         if not prompt.startswith("User:"):
             prompt = f"User: {prompt} Sam:"
@@ -479,6 +480,7 @@ class ModelWrapper:
         
         rng = random.PRNGKey(42)
         generated_ids = input_ids
+        response_text = ""
         
         # Generate tokens
         for _ in range(max_new_tokens):
@@ -504,21 +506,17 @@ class ModelWrapper:
             
             generated_ids = jnp.concatenate([generated_ids, next_token], axis=1)
             
+            # Decode the new token
+            token_id = int(next_token[0, 0])
+            
             # Stop on EOS
-            if next_token[0, 0] == self.tokenizer.token_to_id("<|endoftext|>"):
+            if token_id == self.tokenizer.token_to_id("<|endoftext|>"):
                 break
-        
-        generated_text = self.tokenizer.decode(generated_ids[0].tolist())
-        
-        # Extract response after "Sam:"
-        if "Sam:" in generated_text:
-            response = generated_text.split("Sam:")[-1].strip()
-            # Clean up
-            if "<|endoftext|>" in response:
-                response = response.split("<|endoftext|>")[0].strip()
-            return response
-        else:
-            return generated_text
+            
+            # Decode and yield the token
+            token_text = self.tokenizer.decode([token_id])
+            response_text += token_text
+            yield response_text
 
 # ==============================================================================
 # GRADIO INTERFACE