Update app.py

app.py

@@ -44,7 +44,7 @@ print("Toddler is awake and ready!")
 
 def chat_fn(message, history):
     try:
-        # Build Chat History (Handling standard Gradio list-of-lists format)
+        # 1. Build Token List
         tokens = [tokenizer.bos_token_id]
         for user_msg, assistant_msg in history:
             if user_msg:
@@ -52,14 +52,14 @@ def chat_fn(message, history):
             if assistant_msg:
                 tokens.extend([tokenizer.assistant_start_id] + tokenizer.encode(assistant_msg) + [tokenizer.assistant_end_id])
         
-        # Add current user prompt
         tokens.extend([tokenizer.user_start_id] + tokenizer.encode(message) + [tokenizer.user_end_id])
         tokens.append(tokenizer.assistant_start_id)
 
         input_ids = torch.tensor([tokens], dtype=torch.long)
 
-        # 4. Generate
+        # 2. Generate (Non-streaming for stability)
         with torch.no_grad():
+            # In nanochat, generate usually returns the full sequence tensor
             output_ids = model.generate(
                 input_ids,
                 max_tokens=512,
@@ -67,28 +67,24 @@ def chat_fn(message, history):
                 top_k=40
             )
 
-        # Handle output
+        # 3. Process Output
         if isinstance(output_ids, torch.Tensor):
+            # Slicing to get only new tokens
             new_tokens = output_ids[0][input_ids.shape[1]:]
             response = tokenizer.decode(new_tokens.tolist())
         else:
-            # Generator logic
-            response = ""
-            for token in output_ids:
-                decoded = tokenizer.decode([token])
-                if "<|assistant_end|>" in decoded:
-                    break
-                response += decoded
-                yield response 
+            # If it's a generator, collect it all into one string
+            response = "".join([tokenizer.decode([t]) for t in output_ids])
 
-        # Final cleanup
-        for tag in ["<|assistant_end|>", "<|end|>", "<|user_start|>"]:
+        # 4. Clean up tags
+        for tag in ["<|assistant_end|>", "<|end|>", "<|user_start|>", "<|bos|>"]:
             response = response.split(tag)[0]
         
-        return response.strip()
+        final_text = response.strip()
+        return final_text if final_text else "..."
 
     except Exception as e:
-        print(f"ERROR: {e}")
+        print(f"CRITICAL ERROR: {e}")
         return f"Toddler tantrum: {str(e)}"
 
 # 5. Launch UI (Cleaned for Gradio 6.0 compatibility)