app.py

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

# -----------------------------
# Load 4-bit Qwen model locally
# -----------------------------
MODEL_NAME = "unsloth/qwen2.5-math-1.5b-bnb-4bit"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)

print("Loading model in 4-bit...")
base_model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    device_map="auto",
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True
)

# Check if LoRA adapter exists
try:
    model = PeftModel.from_pretrained(base_model, MODEL_NAME, device_map="auto")
except:
    model = base_model

model.eval()

# -----------------------------
# Respond function
# -----------------------------
def respond(message, history, system_message, max_tokens, temperature, top_p):
    # Build chat prompt
    prompt = system_message + "\n"
    for h in history:
        prompt += f"User: {h['content']}\n"
    prompt += f"User: {message}\nBot:"

    # Tokenize
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # Generation config
    gen_config = GenerationConfig(
        max_new_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        do_sample=True
    )

    # Generate output
    with torch.no_grad():
        output_ids = model.generate(**inputs, **gen_config.to_dict())
    output = tokenizer.decode(output_ids[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
    return output

# -----------------------------
# Gradio Chat Interface
# -----------------------------
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, 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.05, label="Top-p (nucleus sampling)"),
    ],
)

with gr.Blocks() as demo:
    chatbot.render()

# -----------------------------
# Launch Gradio app
# -----------------------------
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)