Update app.py

app.py

@@ -7,21 +7,18 @@ from nanochat.gpt import GPT, GPTConfig
 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 config from your local chat_web.py load
+# 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}
 config = GPTConfig(
     vocab_size=32768,
     n_layer=12,
-    n_head=6,             # ← Critical: matches local load (n_head=6)
-    n_kv_head=6,          # ← Matches n_head (no GQA here)
+    n_head=6,            # Critical fix
+    n_kv_head=6,         # Critical fix
     n_embd=768,
-    sequence_len=2048,    # ← From local log (2048, not 1024)
-    # dropout=0.1,        # omit if TypeError
-    # window_pattern='L'  # if GPTConfig has this field — add if needed
+    sequence_len=2048,   # Updated from 1024
 )
 
 model = GPT(config)
@@ -34,9 +31,7 @@ for k in list(state_dict.keys()):
     if k.startswith(unwanted_prefix):
         state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
 
-missing, unexpected = model.load_state_dict(state_dict, strict=False)
-print(f"Load info: {len(missing)} missing keys, {len(unexpected)} unexpected keys")
-
+model.load_state_dict(state_dict, strict=False)
 model.to("cpu")
 model.eval()
 print("Model ready!")
@@ -45,46 +40,53 @@ engine = Engine(model=model, tokenizer=tokenizer)
 
 def chat_fn(message, history):
     try:
+        # 1. Build prompt and FORCE everything to be a flat list of integers
         prompt_tokens = []
-
-        # Force list conversion on encode output
-        for user, assistant in history:
-            prompt_tokens.extend(list(tokenizer.encode(f"<|user|>{user}<|end|>")))
-            if assistant:
-                prompt_tokens.extend(list(tokenizer.encode(f"<|assistant|>{assistant}<|end|>")))
-
+        
+        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|>")))
+            if assistant_msg:
+                prompt_tokens.extend(list(tokenizer.encode(f"<|assistant|>{assistant_msg}<|end|>")))
+        
+        # Current message
         prompt_tokens.extend(list(tokenizer.encode(f"<|user|>{message}<|end|><|assistant|>")))
 
-        # Optional: add BOS if your model needs it
-        # bos = tokenizer.bos_token_id or list(tokenizer.encode("<|BOS|>"))[0]
-        # prompt_tokens = [bos] + prompt_tokens
-
+        # 2. Generate (Blocking)
+        # Note: We pass prompt_tokens which is now guaranteed to be a List[int]
         new_tokens = engine.generate(
-            prompt_tokens,  # now definitely list[int]
+            prompt_tokens,
             max_tokens=512,
-            temperature=0.85,
+            temperature=0.8,
             top_k=50,
         )
 
-        response = tokenizer.decode(new_tokens).strip()
-
-        # Clean end token
-        for end in ["<|end|>", "<|assistant_end|>", "<|EOS|>"]:
-            if end in response:
-                response = response.split(end)[0].strip()
-                break
+        # 3. Handle the output (ensure it's a list for decoding)
+        if hasattr(new_tokens, 'tolist'):
+            new_tokens = new_tokens.tolist()
+        
+        response = tokenizer.decode(list(new_tokens)).strip()
 
-        return response or "Toddler says: ... 😅"
+        # 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()
+        
+        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("Loaded with exact config from local chat_web.py — n_head=6, seq_len=2048")
+    gr.Markdown("Architecture: 12 Layers, 6 Heads (Verified). Fixed Sequence/Generator error.")
+    
     gr.ChatInterface(
         fn=chat_fn,
-        examples=["Tell me a joke", "What is UPI?", "Write hello world Python"],
+        examples=["Tell me a joke", "What is UPI?", "Explain QA automation"],
         title="Chat with the Toddler"
     )