from peft import PeftModel
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Define model details
base_model_name = "microsoft/phi-2"
adapter_name = "JamieAi33/Phi-2-QLora"

# 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)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# PEFT LoRA Model")
    with gr.Row():
        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 = gr.Button("Generate")
    generate_button.click(generate_text, inputs=[prompt, max_tokens], outputs=output)

demo.launch()