Update app.py

app.py

@@ -1,15 +1,24 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel, PeftConfig
 
-# Model and tokenizer names
-model_name = "djmax13/qween7.5-arabic-story-teller-bnb-4bit"
+# Base model and adapter model names
+base_model_name = "unsloth/Qwen2.5-7B-Instruct-bnb-4bit"
+adapter_model_name = "djmax13/qween7.5-arabic-story-teller-bnb-4bit"
 
-# Load model and tokenizer
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Load base model
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name)
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+
+# Load LoRA configuration
+config = PeftConfig.from_pretrained(adapter_model_name)
+
+# Load LoRA adapter and merge it with the base model
+model = PeftModel.from_pretrained(base_model, adapter_model_name)
+model = model.merge_and_unload() # Optional: Merge adapter weights into base model for potential speedup
 
 def generate_text(prompt):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
+    input_ids = tokenizer.encode(prompt, return_tensors="pt").to(model.device) # Move input to model's device
 
     # Generate text (you might need to adjust generation parameters)
     output = model.generate(
@@ -30,7 +39,7 @@ iface = gr.Interface(
     inputs=gr.Textbox(lines=5, placeholder="Enter your story prompt here..."),
     outputs=gr.Textbox(),
     title="Arabic Story Teller",
-    description="A Qwen2.5-7B model finetuned for Arabic story generation.",
+    description="A Qwen2.5-7B model finetuned for Arabic story generation using LoRA.",
 )
 
 iface.launch()