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

BASE_MODEL =  "unsloth/Llama-3.2-1B-Instruct" # change to 1B to use smaller model
LORA_REPO = "./1B/" # Change this to 1B to use smaller model
device = "cpu"

print('loading tokenizer')
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print('loading base model')
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL, 
    trust_remote_code=True)

print('loading LoRA adapter')
model = PeftModel.from_pretrained(base_model, LORA_REPO)
model.to(device)
model.eval()

def respond(message, history):
    messages = [{"role": "system", "content": "You are a helpful assistant."}]
    for t in history:
        messages.append(t)
    messages.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(device)

    with torch.no_grad():
        out = model.generate(
            input_ids=input_ids,
            max_new_tokens=256,
            do_sample=False,
            temperature=0.7,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id
        )

    output = tokenizer.decode(out[0, input_ids.shape[1]:], skip_special_tokens=True)
    return output

"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
chatbot = gr.ChatInterface(
    respond,
    type="messages",
)

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

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