"""
Hugging Face Gradio Space for your style model.
Upload this as app.py to your HF Space.
"""

import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import torch

# ──────────────────────────────────────────────
MODELS = {
    "3B": "bugaku/igpt",
    "1.5B": "bugaku/igpt-smol",
}
# ──────────────────────────────────────────────

loaded_models = {}

def load_model(model_repo):
    if model_repo in loaded_models:
        return loaded_models[model_repo]

    print(f"Loading model: {model_repo}")
    tokenizer = AutoTokenizer.from_pretrained(model_repo)

    config = AutoConfig.from_pretrained(model_repo)
    if hasattr(config, "quantization_config"):
        delattr(config, "quantization_config")
    if hasattr(config, "_pre_quantization_dtype"):
        delattr(config, "_pre_quantization_dtype")

    model = AutoModelForCausalLM.from_pretrained(
        model_repo,
        config=config,
        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
        device_map="auto",
    )
    model.eval()
    loaded_models[model_repo] = (model, tokenizer)
    print(f"Loaded: {model_repo}")
    return model, tokenizer


def respond(message, history, model_choice, max_new_tokens, temperature):
    model_repo = MODELS[model_choice]
    model, tokenizer = load_model(model_repo)

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

    encoded = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    )
    input_ids = encoded["input_ids"].to(model.device)
    attention_mask = encoded["attention_mask"].to(model.device)

    with torch.no_grad():
        outputs = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=0.9,
            repetition_penalty=1.1,
            do_sample=True,
        )

    response_ids = outputs[0][input_ids.shape[-1]:]
    return tokenizer.decode(response_ids, skip_special_tokens=True)


with gr.Blocks(title="Style Model") as demo:
    gr.Markdown("## 💬 Style Model Chat")
    chatbot = gr.ChatInterface(
        fn=respond,
        additional_inputs=[
            gr.Dropdown(choices=list(MODELS.keys()), value=list(MODELS.keys())[0], label="Model"),
            gr.Slider(32, 512, value=256, step=32, label="Max response length"),
            gr.Slider(0.1, 2.0, value=0.8, step=0.1, label="Temperature"),
        ],
        chatbot=gr.Chatbot(height=500),
    )

demo.launch()