Update app.py

app.py

@@ -1,26 +1,27 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel
 
-# 1️⃣ Load model and tokenizer
-model_name = "DSDUDEd/funfox"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
+# 1️⃣ 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
 
-# 2️⃣ Setup text generation pipeline
-generator = pipeline(
-    task="text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device=0 if torch.cuda.is_available() else -1
-)
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name, torch_dtype=torch.float16).to("cuda" if torch.cuda.is_available() else "cpu")
+
+# 2️⃣ Load PEFT model
+peft_model_name = "DSDUDEd/funfox"
+model = PeftModel.from_pretrained(base_model, peft_model_name)
+model.eval()
 
-# 3️⃣ Define function for Gradio interface
+# 3️⃣ Define text generation function
 def generate_text(prompt, max_tokens=50):
-    output = generator(prompt, max_new_tokens=max_tokens)
-    return output[0]["generated_text"]
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(**inputs, max_new_tokens=max_tokens)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# 4️⃣ Create Gradio interface
+# 4️⃣ Build Gradio interface
 iface = gr.Interface(
     fn=generate_text,
     inputs=[
@@ -28,9 +29,9 @@ iface = gr.Interface(
         gr.Slider(minimum=10, maximum=200, step=10, label="Max Tokens")
     ],
     outputs=gr.Textbox(label="Generated Text"),
-    title="FunFox Text Generation",
-    description="Enter a prompt and generate text with the FunFox model."
+    title="FunFox PEFT Model",
+    description="FunFox LoRA model fine-tuned on PythonCopilot base."
 )
 
-# 5️⃣ Launch the interface
+# 5️⃣ Launch
 iface.launch(share=True)