"""
app.py - Hugging Face Spaces (CPU) deployment for the fine-tuned Llama 3.2 1B LoRA.

This runs WITHOUT unsloth or bitsandbytes, because both require a GPU/CUDA.
Instead it loads a full-precision base model on CPU with `transformers` and
applies your LoRA adapter on top with `peft`, then merges the two so CPU
generation is a little faster.

The adapter repo (mthuy/dssp-llama-1b-lora) is public, so no token is needed.
"""

import os
import torch
import gradio as gr

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

# ------------------------------------------------------------
# Configuration
# ------------------------------------------------------------

# Your fine-tuned LoRA adapter (also holds the tokenizer + chat template).
ADAPTER_REPO_ID = "mthuy/dssp-llama-1b-lora"

# Full-precision base model to apply the adapter on top of.
# The adapter was trained on "unsloth/Llama-3.2-1B-Instruct-bnb-4bit" (4-bit),
# which can't load on CPU. This is the same model, un-quantized and ungated.
BASE_MODEL_ID = "unsloth/Llama-3.2-1B-Instruct"

# Keep CPU generation responsive.
torch.set_num_threads(os.cpu_count() or 2)

# ------------------------------------------------------------
# Load model + tokenizer once at startup
# ------------------------------------------------------------

print("Loading tokenizer from adapter repo...")
tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO_ID)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

print("Loading base model on CPU (float32)...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    torch_dtype=torch.float32,   # float32 is the safe choice on CPU
    low_cpu_mem_usage=True,
)

print("Applying LoRA adapter...")
model = PeftModel.from_pretrained(base_model, ADAPTER_REPO_ID)

# Merge LoRA weights into the base for faster CPU inference, then drop PEFT wrappers.
model = model.merge_and_unload()
model.eval()

print("Model ready.")

# ------------------------------------------------------------
# Inference
# ------------------------------------------------------------

def ask_model(question, max_new_tokens=60, temperature=0.0):
    if not question or not question.strip():
        return ""

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

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )
    # Everything stays on CPU.

    generation_args = {
        "max_new_tokens": int(max_new_tokens),
        "pad_token_id": tokenizer.eos_token_id,
        "eos_token_id": tokenizer.eos_token_id,
        "use_cache": True,
    }

    if temperature and temperature > 0:
        generation_args["do_sample"] = True
        generation_args["temperature"] = float(temperature)
        generation_args["top_p"] = 0.9
    else:
        generation_args["do_sample"] = False

    with torch.inference_mode():
        outputs = model.generate(**inputs, **generation_args)

    input_length = inputs["input_ids"].shape[-1]
    new_tokens = outputs[0][input_length:]
    response = tokenizer.decode(new_tokens, skip_special_tokens=True)

    return response.strip()

# ------------------------------------------------------------
# Gradio UI
# ------------------------------------------------------------

with gr.Blocks(title="DSSP Fine-Tuned Llama 1B Demo") as demo:

    gr.Markdown(
        """
        # DSSP Fine-Tuned Llama 1B Demo (Dog Food Edition)

        This Space runs a Llama 3.2 1B model fine-tuned with a LoRA adapter on a
        handful of dog food / canine nutrition facts. It runs on **CPU**, so the
        first answer after a cold start can take a little while.

        Adapter: `mthuy/dssp-llama-1b-lora`
        """
    )

    question = gr.Textbox(
        label="Question",
        lines=3,
        placeholder="Example: Is chocolate safe for dogs?"
    )

    with gr.Row():
        max_new_tokens = gr.Slider(
            minimum=20,
            maximum=100,
            value=60,
            step=10,
            label="Max new tokens"
        )

        temperature = gr.Slider(
            minimum=0.0,
            maximum=1.0,
            value=0.0,
            step=0.1,
            label="Temperature"
        )

    answer = gr.Textbox(
        label="Model answer",
        lines=6
    )

    submit = gr.Button("Ask model")

    submit.click(
        fn=ask_model,
        inputs=[question, max_new_tokens, temperature],
        outputs=answer
    )

    gr.Examples(
        examples=[
            ["Is chocolate safe for dogs?", 60, 0.0],
            ["Can dogs eat grapes?", 60, 0.0],
            ["Are onions safe for dogs?", 60, 0.0],
            ["What is the most important nutrient in a dog's diet?", 60, 0.0],
            ["How often should an adult dog be fed?", 60, 0.0],
        ],
        inputs=[question, max_new_tokens, temperature],
        outputs=answer,
        fn=ask_model,
    )

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