Update app.py

app.py

@@ -9,23 +9,20 @@ MODEL_PATH = "model_000971.pt"
 print("Waking up the toddler (NanoChat-ClimbMix-D12)...")
 tokenizer = get_tokenizer()
 
-print("Creating GPT model skeleton with exact local config...")
-# EXACT values from your local scripts.chat_web output:
-# {'sequence_len': 2048, 'vocab_size': 32768, 'n_layer': 12, 'n_head': 6, 'n_kv_head': 6, 'n_embd': 768}
+# EXACT values from your local scripts.chat_web output
 config = GPTConfig(
     vocab_size=32768,
     n_layer=12,
-    n_head=6,            # Critical fix
-    n_kv_head=6,         # Critical fix
+    n_head=6,            
+    n_kv_head=6,         
     n_embd=768,
-    sequence_len=2048,   # Updated from 1024
+    sequence_len=2048,   
 )
 
 model = GPT(config)
 
 print("Loading weights...")
 state_dict = torch.load(MODEL_PATH, map_location="cpu", weights_only=False)
-
 unwanted_prefix = '_orig_mod.'
 for k in list(state_dict.keys()):
     if k.startswith(unwanted_prefix):
@@ -38,22 +35,26 @@ print("Model ready!")
 
 engine = Engine(model=model, tokenizer=tokenizer)
 
+def safe_encode(text):
+    """Helper to ensure we only get the list of token IDs."""
+    encoded = tokenizer.encode(text)
+    # If it's a tuple (ids, mask), just take the ids
+    if isinstance(encoded, tuple):
+        return list(encoded[0])
+    return list(encoded)
+
 def chat_fn(message, history):
     try:
-        # 1. Build prompt and FORCE everything to be a flat list of integers
         prompt_tokens = []
         
         for user_msg, assistant_msg in history:
-            # We wrap in list() to kill any generator objects returned by tokenizer
-            prompt_tokens.extend(list(tokenizer.encode(f"<|user|>{user_msg}<|end|>")))
+            prompt_tokens.extend(safe_encode(f"<|user|>{user_msg}<|end|>"))
             if assistant_msg:
-                prompt_tokens.extend(list(tokenizer.encode(f"<|assistant|>{assistant_msg}<|end|>")))
+                prompt_tokens.extend(safe_encode(f"<|assistant|>{assistant_msg}<|end|>"))
         
-        # Current message
-        prompt_tokens.extend(list(tokenizer.encode(f"<|user|>{message}<|end|><|assistant|>")))
+        prompt_tokens.extend(safe_encode(f"<|user|>{message}<|end|><|assistant|>"))
 
-        # 2. Generate (Blocking)
-        # Note: We pass prompt_tokens which is now guaranteed to be a List[int]
+        # Blocking generation
         new_tokens = engine.generate(
             prompt_tokens,
             max_tokens=512,
@@ -61,13 +62,14 @@ def chat_fn(message, history):
             top_k=50,
         )
 
-        # 3. Handle the output (ensure it's a list for decoding)
+        # Convert output to list if it's a tensor or tuple
         if hasattr(new_tokens, 'tolist'):
             new_tokens = new_tokens.tolist()
-        
+        elif isinstance(new_tokens, tuple):
+            new_tokens = new_tokens[0]
+            
         response = tokenizer.decode(list(new_tokens)).strip()
 
-        # 4. Clean up the tags
         for stop_tag in ["<|end|>", "<|assistant_end|>", "<|EOS|>"]:
             if stop_tag in response:
                 response = response.split(stop_tag)[0].strip()
@@ -75,18 +77,14 @@ def chat_fn(message, history):
         return response or "Toddler is thinking... 😅"
 
     except Exception as e:
-        # If it still fails, this will tell us exactly which line crashed
-        import traceback
-        print(traceback.format_exc())
         return f"Toddler tantrum: {str(e)}"
 
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
-    gr.Markdown("# 🧸 NanoChat-ClimbMix-D12 – Confident Toddler")
-    gr.Markdown("Architecture: 12 Layers, 6 Heads (Verified). Fixed Sequence/Generator error.")
-    
+    gr.Markdown("# 🧸 NanoChat-ClimbMix-D12")
+    gr.Markdown("Architecture verified. Tuple/Generator issues handled.")
     gr.ChatInterface(
         fn=chat_fn,
-        examples=["Tell me a joke", "What is UPI?", "Explain QA automation"],
+        examples=["Tell me a joke", "What is UPI?"],
         title="Chat with the Toddler"
     )