Update app.py

app.py

@@ -225,8 +225,8 @@ from transformers import AutoTokenizer
 
 hf_tokenizer = AutoTokenizer.from_pretrained("gpt2")
 
-# Add custom tokens
-custom_tokens = ["<|im_start|>", "<|im_end|>"]
+# Add custom tokens to match model's vocab size
+custom_tokens = ["<|im_start|>", "<|im_end|>", "<think>", "<think/>"]
 hf_tokenizer.add_special_tokens({"additional_special_tokens": custom_tokens})
 
 # Save and reload as tokenizers format
@@ -235,7 +235,13 @@ hf_tokenizer.save_pretrained("./temp_tokenizer")
 tokenizer = Tokenizer.from_file("./temp_tokenizer/tokenizer.json")
 
 print(f"✅ Tokenizer created with vocab size: {tokenizer.get_vocab_size()}")
-print(f"   Custom tokens: {custom_tokens}")
+print(f"   Custom tokens added: {custom_tokens}")
+print(f"   Model vocab size: {config.get('vocab_size', 'unknown')}")
+
+# Verify vocab sizes match
+if tokenizer.get_vocab_size() != config.get('vocab_size'):
+    print(f"⚠️  WARNING: Tokenizer vocab ({tokenizer.get_vocab_size()}) != Model vocab ({config.get('vocab_size')})")
+    print(f"   Model was trained with these tokens, but SAM-Z-1 doesn't use <think> tags in generation")
 
 eos_token_id = config.get('eos_token_id', 50256)
 
@@ -681,8 +687,10 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                     **Speed:** ⚡ Optimized with TF Functions
                     
                     **Twin Model:**
-                    - **SAM-X-1**: Reasoning model (with thinking)
-                    - **SAM-Z-1**: Fast model (YOU ARE HERE! 🎉)
+                    - **SAM-X-1**: Reasoning model (uses `<think>` tags)
+                    - **SAM-Z-1**: Fast model (no thinking, direct answers! 🎉)
+                    
+                    **Note:** Model includes `<think>` tokens in vocab but doesn't use them. Training used same tokenizer as SAM-X-1.
                     
                     **Architecture:**
                     - RoPE positional encoding
@@ -706,8 +714,10 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
                     **Vocab:** {config['vocab_size']}
                     
                     **Twin Models:**
-                    - SAM-X-1: Reasoning model
-                    - SAM-Z-1: Direct response model
+                    - SAM-X-1: Reasoning model (uses `<think>` tags)
+                    - SAM-Z-1: Direct response model (no thinking)
+                    
+                    **Note:** Vocab includes `<think>` tokens but model doesn't use them in generation.
                     
                     **Features:**
                     - RoPE positional encoding