Update app.py

app.py

@@ -3,9 +3,9 @@ import torch
 from transformers import GPT2Tokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
-# 1️⃣ Load tokenizer (fallback)
+# 1️⃣ Load fallback tokenizer (GPT2)
 tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
-tokenizer.pad_token = tokenizer.eos_token  # Important for causal LM
+tokenizer.pad_token = tokenizer.eos_token  # Required for causal LM
 
 # 2️⃣ Load base model
 base_model_name = "TRM-coding/PythonCopilot"
@@ -13,18 +13,19 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 base_model = AutoModelForCausalLM.from_pretrained(
     base_model_name,
-    torch_dtype=torch.float16 if device=="cuda" else torch.float32
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
 ).to(device)
 
-# 3️⃣ Load PEFT/LoRA weights
+# 3️⃣ Resize embeddings to match PEFT checkpoint vocab
+checkpoint_vocab_size = 50257  # From DSUDUDe/funfox PEFT model
+base_model.resize_token_embeddings(checkpoint_vocab_size)
+
+# 4️⃣ Load PEFT/LoRA adapter
 peft_model_name = "DSDUDEd/funfox"
 model = PeftModel.from_pretrained(base_model, peft_model_name)
 model.eval()
 
-# 4️⃣ Resize embeddings to match tokenizer
-model.resize_token_embeddings(len(tokenizer))
-
-# 5️⃣ Text generation function
+# 5️⃣ Define generation function
 def generate_text(prompt, max_tokens=50):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
     outputs = model.generate(
@@ -36,7 +37,7 @@ def generate_text(prompt, max_tokens=50):
     )
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# 6️⃣ Gradio interface
+# 6️⃣ Build Gradio interface
 iface = gr.Interface(
     fn=generate_text,
     inputs=[