Update app.py

app.py

@@ -1,104 +1,111 @@
+import os
 import torch
 import gradio as gr
-from nanochat.engine import Engine
-from nanochat.tokenizer import get_tokenizer
 from nanochat.gpt import GPT, GPTConfig
+from nanochat.tokenizer import RustBPETokenizer
 
+# Configuration
 MODEL_PATH = "model_000971.pt"
+# The Dockerfile moves files to this specific cache location
+CACHE_DIR = os.path.expanduser("~/.cache/nanochat/tokenizer/")
+# Fallback to current directory if cache doesn't exist (local testing)
+TOKENIZER_DIR = CACHE_DIR if os.path.exists(CACHE_DIR) else "."
 
-print("Waking up the toddler (NanoChat-ClimbMix-D12)...")
-tokenizer = get_tokenizer()
+print(f"--- Waking up the Toddler ---")
+print(f"Loading tokenizer from: {TOKENIZER_DIR}")
 
-# SET SPECIAL TOKENS (Aligned with Saint Iberis working space)
-# We use .get() or try/except to handle different tokenizer versions
-try:
-    bos_id = tokenizer.encode("<|bos|>")[0]
-    user_start_id = tokenizer.encode("<|user_start|>")[0]
-    user_end_id = tokenizer.encode("<|user_end|>")[0]
-    assistant_start_id = tokenizer.encode("<|assistant_start|>")[0]
-    assistant_end_id = tokenizer.encode("<|assistant_end|>")[0]
-except:
-    # Fallback to standard tags if the specific Saint Iberis ones aren't in your vocab
-    bos_id = tokenizer.encode("<|endoftext|>")[0]
-    user_start_id = tokenizer.encode("<|user|>")[0]
-    user_end_id = tokenizer.encode("<|end|>")[0]
-    assistant_start_id = tokenizer.encode("<|assistant|>")[0]
-    assistant_end_id = tokenizer.encode("<|end|>")[0]
+# 1. Load Tokenizer & Map Special Tokens
+tokenizer = RustBPETokenizer.from_directory(TOKENIZER_DIR)
 
-print("Creating GPT model skeleton (6-head, 2048 seq)...")
+# These must match your training vocab
+tokenizer.bos_token_id = tokenizer.enc.encode_single_token("<|bos|>")
+tokenizer.user_start_id = tokenizer.enc.encode_single_token("<|user_start|>")
+tokenizer.user_end_id = tokenizer.enc.encode_single_token("<|user_end|>")
+tokenizer.assistant_start_id = tokenizer.enc.encode_single_token("<|assistant_start|>")
+tokenizer.assistant_end_id = tokenizer.enc.encode_single_token("<|assistant_end|>")
+
+# 2. Build Model Architecture
 config = GPTConfig(
     vocab_size=32768,
     n_layer=12,
-    n_head=6,            
-    n_kv_head=6,         
+    n_head=6,
+    n_kv_head=6,
     n_embd=768,
-    sequence_len=2048,   
+    sequence_len=2048,
 )
 
 model = GPT(config)
 
+# 3. Load Weights (with _orig_mod. cleaning)
 print("Loading weights...")
 state_dict = torch.load(MODEL_PATH, map_location="cpu", weights_only=False)
 state_dict = {k.replace("_orig_mod.", ""): v for k, v in state_dict.items()}
-
 model.load_state_dict(state_dict, strict=False)
 model.to("cpu")
 model.eval()
-print("Model ready!")
+print("Toddler is awake and ready!")
 
-# We use the model directly to avoid 'Engine' type-hinting issues
 def chat_fn(message, history):
     try:
-        # 1. Build token list
-        tokens = [bos_id]
+        # Build Chat History
+        tokens = [tokenizer.bos_token_id]
         for user_msg, assistant_msg in history:
-            tokens.extend([user_start_id] + list(tokenizer.encode(user_msg)) + [user_end_id])
+            tokens.extend([tokenizer.user_start_id] + tokenizer.encode(user_msg) + [tokenizer.user_end_id])
             if assistant_msg:
-                tokens.extend([assistant_start_id] + list(tokenizer.encode(assistant_msg)) + [assistant_end_id])
+                tokens.extend([tokenizer.assistant_start_id] + tokenizer.encode(assistant_msg) + [tokenizer.assistant_end_id])
         
-        # Current turn
-        tokens.extend([user_start_id] + list(tokenizer.encode(message)) + [user_end_id])
-        tokens.append(assistant_start_id)
+        # Add current user prompt
+        tokens.extend([tokenizer.user_start_id] + tokenizer.encode(message) + [tokenizer.user_end_id])
+        tokens.append(tokenizer.assistant_start_id)
 
-        # 2. THE FIX: Convert to Tensor before generating
-        input_ids = torch.tensor([tokens], dtype=torch.long).to("cpu")
+        input_ids = torch.tensor([tokens], dtype=torch.long)
 
-        # 3. Generate
-        # If your model.generate is a generator (streaming), we'll take the result
+        # 4. Generate with Streaming Logic
+        # Note: In nanochat.gpt, generate is typically an autoregressive loop. 
+        # If your version returns a generator, we iterate. If a tensor, we slice.
         with torch.no_grad():
             output_ids = model.generate(
                 input_ids,
                 max_tokens=512,
                 temperature=0.8,
-                top_k=50,
+                top_k=40
             )
 
-        # 4. Decode (handling both streaming and blocking outputs)
-        # If output_ids is a generator, we collect it; if it's a tensor, we decode it.
+        # Handle Tensor vs Generator output
         if isinstance(output_ids, torch.Tensor):
-            # Take only the newly generated tokens
+            # Just take the new parts
             new_tokens = output_ids[0][input_ids.shape[1]:]
-            response = tokenizer.decode(new_tokens.tolist()).strip()
+            response = tokenizer.decode(new_tokens.tolist())
         else:
-            # It's a generator (streaming)
-            full_response = ""
+            # It's a generator yielding token by token
+            response = ""
             for token in output_ids:
-                full_response += tokenizer.decode([token])
-            response = full_response.strip()
+                decoded = tokenizer.decode([token])
+                if "<|assistant_end|>" in decoded:
+                    break
+                response += decoded
+                yield response # Yield for streaming UI effect
 
-        # Clean up stop tags
+        # Final cleanup for non-streaming return
         for tag in ["<|assistant_end|>", "<|end|>", "<|user_start|>"]:
-            if tag in response:
-                response = response.split(tag)[0].strip()
-
-        return response or "Toddler is thinking... 😅"
+            response = response.split(tag)[0]
+        
+        return response.strip()
 
     except Exception as e:
+        # Crucial for QA: see the actual error in Space logs
+        print(f"ERROR: {e}")
         return f"Toddler tantrum: {str(e)}"
 
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
+# 5. Launch UI
+with gr.Blocks(theme=gr.themes.Default(primary_hue="orange")) as demo:
     gr.Markdown("# 🧸 NanoChat-ClimbMix-D12")
-    gr.ChatInterface(fn=chat_fn, title="Toddler Chat")
+    gr.Markdown("A custom-trained small language model running on your CPU.")
+    gr.ChatInterface(
+        fn=chat_fn, 
+        type="messages", # Updated for latest Gradio versions
+        examples=["Hi Toddler!", "How does UPI work?", "Tell me a story."]
+    )
 
 if __name__ == "__main__":
     demo.launch(server_name="0.0.0.0", server_port=7860)