app.py

app.py

@@ -1,27 +1,42 @@
 import gradio as gr
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import GPT2Tokenizer, AutoModelForCausalLM
 from peft import PeftModel
 
-# 1️⃣ Load base model
+# 1️⃣ Load tokenizer (fallback)
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+tokenizer.pad_token = tokenizer.eos_token  # Important for causal LM
+
+# 2️⃣ Load base model
 base_model_name = "TRM-coding/PythonCopilot"
-tokenizer = AutoTokenizer.from_pretrained(base_model_name)
-tokenizer.pad_token = tokenizer.eos_token  # important for causal LM
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float16).to("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
+).to(device)
 
-# 2️⃣ Load PEFT model
+# 3️⃣ Load PEFT/LoRA weights
 peft_model_name = "DSDUDEd/funfox"
 model = PeftModel.from_pretrained(base_model, peft_model_name)
 model.eval()
 
-# 3️⃣ Define text generation function
+# 4️⃣ Resize embeddings to match tokenizer
+model.resize_token_embeddings(len(tokenizer))
+
+# 5️⃣ Text generation function
 def generate_text(prompt, max_tokens=50):
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=max_tokens)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=max_tokens,
+        do_sample=True,
+        top_p=0.9,
+        temperature=0.8
+    )
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# 4️⃣ Build Gradio interface
+# 6️⃣ Gradio interface
 iface = gr.Interface(
     fn=generate_text,
     inputs=[
@@ -33,5 +48,5 @@ iface = gr.Interface(
     description="FunFox LoRA model fine-tuned on PythonCopilot base."
 )
 
-# 5️⃣ Launch
+# 7️⃣ Launch
 iface.launch(share=True)