Update space

app.py

@@ -27,8 +27,18 @@ def load_model():
         low_cpu_mem_usage=True
     )
 
+    # Resize token embeddings to match the adapter
+    if base_model.get_input_embeddings().weight.shape[0] != 128257:
+        base_model.resize_token_embeddings(128257)
+        print(f"Resized embeddings from {base_model.get_input_embeddings().weight.shape[0]-1} to 128257")
+
     # Load tokenizer
     tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+    
+    # Ensure tokenizer has the same vocab size as the model
+    if len(tokenizer) != base_model.get_input_embeddings().weight.shape[0]:
+        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+        print(f"Added padding token to tokenizer. New size: {len(tokenizer)}")
 
     # Load LoRA adapter
     print(f"Loading LoRA adapter: {ADAPTER_MODEL}")
@@ -79,4 +89,4 @@ demo = gr.Interface(
     description=f"Running on {'CUDA' if torch.cuda.is_available() else 'CPU'} - May be slow on CPU."
 )
 
-demo.launch()
+demo.launch()