Update app.py

app.py

@@ -7,19 +7,21 @@ 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 using verified local configs...")
-# EXACT CONFIG FROM YOUR SUCCESSFUL scripts.chat_web LOG
+print("Creating GPT model skeleton with exact local config...")
+
+# Exact config from your local chat_web.py load
 config = GPTConfig(
     vocab_size=32768,
     n_layer=12,
-    n_head=6,            # Verified: 6 heads
-    n_kv_head=6,         # Verified: 6 KV heads
-    n_embd=768,          # Verified: 768 embedding dim
-    sequence_len=2048,   # Verified: 2048 context
-    # Note: If your GPT class accepts window_pattern, uncomment below:
-    # window_pattern='L', 
+    n_head=6,             # ← Critical: matches local load (n_head=6)
+    n_kv_head=6,          # ← Matches n_head (no GQA here)
+    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
 )
 
 model = GPT(config)
@@ -27,14 +29,14 @@ model = GPT(config)
 print("Loading weights...")
 state_dict = torch.load(MODEL_PATH, map_location="cpu", weights_only=False)
 
-# Remove the DDP wrapper prefix if present
 unwanted_prefix = '_orig_mod.'
 for k in list(state_dict.keys()):
     if k.startswith(unwanted_prefix):
         state_dict[k[len(unwanted_prefix):]] = state_dict.pop(k)
 
-# Loading weights - should be 100% match now
-model.load_state_dict(state_dict, strict=False)
+missing, unexpected = model.load_state_dict(state_dict, strict=False)
+print(f"Load info: {len(missing)} missing keys, {len(unexpected)} unexpected keys")
+
 model.to("cpu")
 model.eval()
 print("Model ready!")
@@ -43,17 +45,15 @@ engine = Engine(model=model, tokenizer=tokenizer)
 
 def chat_fn(message, history):
     try:
-        # Build prompt using the standard nanochat special tokens
+        # Build prompt like local chat_web.py (adjust tags if needed)
         prompt_tokens = []
-        for user_msg, assistant_msg in history:
-            prompt_tokens += tokenizer.encode(f"<|user|>{user_msg}<|end|>")
-            if assistant_msg:
-                prompt_tokens += tokenizer.encode(f"<|assistant|>{assistant_msg}<|end|>")
-        
-        # Current turn
+        for user, assistant in history:
+            prompt_tokens += tokenizer.encode(f"<|user|>{user}<|end|>")
+            if assistant:
+                prompt_tokens += tokenizer.encode(f"<|assistant|>{assistant}<|end|>")
+
         prompt_tokens += tokenizer.encode(f"<|user|>{message}<|end|><|assistant|>")
 
-        # Blocking generation (NanoChat Engine standard)
         new_tokens = engine.generate(
             prompt_tokens,
             max_tokens=512,
@@ -62,25 +62,20 @@ def chat_fn(message, history):
         )
 
         response = tokenizer.decode(new_tokens).strip()
+        if "<|end|>" in response:
+            response = response.split("<|end|>")[0].strip()
 
-        # Clean up tags
-        for stop_tag in ["<|end|>", "<|assistant_end|>", "<|EOS|>"]:
-            if stop_tag in response:
-                response = response.split(stop_tag)[0].strip()
-        
-        return response or "The toddler is lost in thought... 😅"
+        return response or "Toddler says: ... 😅"
 
     except Exception as e:
         return f"Toddler tantrum: {str(e)}"
 
-# Gradio UI
 with gr.Blocks(theme=gr.themes.Soft(primary_hue="blue")) as demo:
-    gr.Markdown("# 🧸 NanoChat-ClimbMix-D12")
-    gr.Markdown("Architecture: 12 Layers, 6 Heads, GQA. Optimized for Fintech testing.")
-    
+    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.ChatInterface(
         fn=chat_fn,
-        examples=["Tell me a joke", "What is NPCI FiMI?", "Explain UPI transaction flow"],
+        examples=["Tell me a joke", "What is UPI?", "Write hello world Python"],
         title="Chat with the Toddler"
     )