Add PEFT LoRA support

app.py

@@ -1,32 +1,35 @@
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
 
-# Load the PEFT model and tokenizer from Hugging Face Hub
-model_name = "JamieAi33/Phi-2_PEFT"
-model = AutoModelForCausalLM.from_pretrained(model_name)
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Define model details
+base_model_name = "microsoft/phi-2"
+adapter_name = "JamieAi33/Phi-2-QLora"
 
-# Define the prediction function
-def generate_text(prompt, max_length=100):
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=max_length)
+# Load base model
+print("Loading base model...")
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+
+# Apply LoRA adapter
+print("Loading LoRA adapter...")
+model = PeftModel.from_pretrained(base_model, adapter_name)
+
+# Function to generate text
+def generate_text(prompt, max_tokens):
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
+    outputs = model.generate(**inputs, max_new_tokens=max_tokens)
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-# Create the Gradio interface
+# Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# PEFT LLM Demo")
-    gr.Markdown("Generate text using the Phi-2 PEFT model.")
+    gr.Markdown("# PEFT LoRA Model")
     with gr.Row():
-        prompt_input = gr.Textbox(label="Input Prompt", placeholder="Enter a prompt here...")
-        max_tokens_input = gr.Slider(label="Max Tokens", minimum=10, maximum=200, value=100, step=10)
-    generate_button = gr.Button("Generate")
-    output_text = gr.Textbox(label="Generated Text", placeholder="Generated text will appear here.")
+        prompt = gr.Textbox(label="Prompt", lines=4)
+        max_tokens = gr.Slider(label="Max Tokens", minimum=10, maximum=200, value=50)
+    output = gr.Textbox(label="Generated Text", lines=6)
 
-    generate_button.click(
-        fn=generate_text,
-        inputs=[prompt_input, max_tokens_input],
-        outputs=output_text
-    )
+    generate_button = gr.Button("Generate")
+    generate_button.click(generate_text, inputs=[prompt, max_tokens], outputs=output)
 
-# Launch the app
-demo.launch()
+demo.launch()

requirements.txt

@@ -1,3 +1,4 @@
-gradio
-transformers
 torch
+transformers
+peft
+gradio