import gradio as gr
import spaces
import torch
import os

from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    BitsAndBytesConfig,
)

MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"

HF_TOKEN = os.environ["HF_TOKEN"]

tokenizer = None
model = None


def load_model():
    global tokenizer, model

    if model is not None:
        return

    quant_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.float16,
        bnb_4bit_use_double_quant=True,
    )

    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_ID,
        token=HF_TOKEN,
    )

    model = AutoModelForCausalLM.from_pretrained(
        MODEL_ID,
        token=HF_TOKEN,
        quantization_config=quant_config,
        device_map="cuda",
        torch_dtype=torch.float16,
    )

    print("Model loaded on:", model.device)


@spaces.GPU
def greet(message):
    load_model()

    messages = [
        {
            "role": "user",
            "content": message,
        }
    ]

    prompt = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    inputs = tokenizer(
        prompt,
        return_tensors="pt",
    ).to("cuda")

    with torch.inference_mode():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
        )

    generated_tokens = outputs[0][inputs["input_ids"].shape[-1]:]

    return tokenizer.decode(
        generated_tokens,
        skip_special_tokens=True,
    )


demo = gr.Interface(
    fn=greet,
    inputs=gr.Textbox(
        label="Message",
        placeholder="Ask Llama something...",
    ),
    outputs=gr.Textbox(
        label="Llama 3.1 8B Response",
    ),
    title="Llama 3.1 8B Instruct",
)

demo.launch()