app.py

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

# Set the HF repo and LoRA model location
base_model_id = "unsloth/gemma-2-9b"
lora_model_id = "Futuresony/gemma2-9b-lora-alpaca"

# Load base model on CPU
base_model = AutoModelForCausalLM.from_pretrained(
    base_model_id,
    device_map="cpu",
    torch_dtype=torch.float32,
)

# Load tokenizer from base model
tokenizer = AutoTokenizer.from_pretrained(base_model_id)

# Load LoRA adapter
model = PeftModel.from_pretrained(base_model, lora_model_id)
model.eval()

# === Alpaca-style formatter ===
def format_alpaca_prompt(user_input, system_prompt, history):
    history_str = "\n".join([f"### Instruction:\n{h[0]}\n### Response:\n{h[1]}" for h in history])
    prompt = f"""{system_prompt}
{history_str}

### Instruction:
{user_input}

### Response:"""
    return prompt

# === Chat logic ===
def respond(message, history, system_message, max_tokens, temperature, top_p):
    prompt = format_alpaca_prompt(message, system_message, history)
    inputs = tokenizer(prompt, return_tensors="pt").to("cpu")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            do_sample=True,
            pad_token_id=tokenizer.eos_token_id,
        )

    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    # Only return the part after "### Response:"
    if "### Response:" in response_text:
        final_output = response_text.split("### Response:")[-1].strip()
    else:
        final_output = response_text.strip()

    history.append((message, final_output))
    yield final_output

# === Gradio Interface ===
demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=1024, value=256, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.01, label="Top-p"),
    ],
    title="Offline Gemma-2B Alpaca Chatbot (LoRA)",
)

if __name__ == "__main__":
    demo.launch()