app.py

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

# Base model and adapter repo
BASE_MODEL_NAME = "microsoft/phi-2"
ADAPTER_REPO = "Shriti09/Microsoft-Phi-QLora" 

# Load the tokenizer
print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

# Load the base model
print("Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(BASE_MODEL_NAME, device_map="auto")

# Load adapter weights
print("Loading LoRA adapter...")
model = PeftModel.from_pretrained(base_model, ADAPTER_REPO)

# Merge adapter into base model (optional, makes inference simpler)
model = model.merge_and_unload()

# Put model in eval mode
model.eval()

# Function to generate response from prompt
def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_length=256,
        do_sample=True,
        top_p=0.95,
        temperature=0.7
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Gradio UI
gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Ask me something..."),
    outputs="text",
    title="Phi-2 QLoRA Chatbot",
    description="Chat with Phi-2 fine-tuned with QLoRA adapters!"
).launch()