import gradio as gr
import torch

# Some Llama-4 checkpoints store `attn_temperature_tuning` as an int (e.g. 4)
# instead of a bool. Newer huggingface_hub versions strictly validate config
# field types and reject this, even though it loaded fine originally. Relax
# that one check so int values are treated as truthy bools.
import huggingface_hub.dataclasses as _hfd

_original_validate_simple_type = _hfd._validate_simple_type


def _lenient_validate_simple_type(name, value, expected_type):
    if expected_type is bool and isinstance(value, int) and not isinstance(value, bool):
        return
    return _original_validate_simple_type(name, value, expected_type)


_hfd._validate_simple_type = _lenient_validate_simple_type

from transformers import Llama4ForConditionalGeneration, AutoTokenizer
from peft import PeftModel

# The adapter was trained against a text-only Llama4ForCausalLM checkpoint that
# was never published publicly. Llama-4-Scout is normally distributed as the
# full multimodal Llama4ForConditionalGeneration; this public 4-bit re-upload
# has an identical text backbone (same hidden_size/layers/experts/vocab as the
# adapter's config), and PEFT matches target module names by suffix, so the
# LoRA still applies correctly onto its `language_model` submodule.
BASE_MODEL = "unsloth/Llama-4-Scout-17B-16E-Instruct-unsloth-bnb-4bit"
ADAPTER_REPO = "sidddd625/llama4-scout-india-financial-rights-lora"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(ADAPTER_REPO)

print("Loading base model (this can take a few minutes)...")
base_model = Llama4ForConditionalGeneration.from_pretrained(
    BASE_MODEL,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)

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


def respond(message, history):
    messages = []
    for user_msg, assistant_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if assistant_msg:
            messages.append({"role": "assistant", "content": assistant_msg})
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
    ).to(model.device)

    with torch.no_grad():
        output_ids = model.generate(
            inputs,
            max_new_tokens=512,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id or tokenizer.eos_token_id,
        )

    response = tokenizer.decode(
        output_ids[0][inputs.shape[-1]:], skip_special_tokens=True
    )
    return response


demo = gr.ChatInterface(
    fn=respond,
    title="Llama-4-Scout + India Financial Rights LoRA",
    description="Fine-tuned adapter for questions about Indian financial rights.",
)

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