"""Gradio demo for the New Prussian translator (Apertus-8B-int8 + LoRA).

Loads the pre-quantized int8 base model `strfry/Apertus-8B-Instruct-2509-int8`
and applies the LoRA adapter `strfry/apertus-8b-prussian-youtube` — all at
module level on CPU. ZeroGPU transfers everything to VRAM automatically when
a request enters @spaces.GPU. No lazy loading, no safetensors CUDA issues.
"""

import json
import torch
import gradio as gr
import safetensors.torch
from huggingface_hub import hf_hub_download
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel, LoraConfig, set_peft_model_state_dict

# ZeroGPU decorator if available; no-op fallback so the app also runs locally
try:
    import spaces

    GPU = spaces.GPU
except ImportError:  # local / non-ZeroGPU

    def GPU(func=None, **_kwargs):
        if func is None:
            return lambda f: f
        return func


BASE_MODEL = "strfry/Apertus-8B-Instruct-2509-int8"
ADAPTER = "strfry/apertus-8b-prussian-youtube"

SYSTEM_PROMPT = "Translate to reconstructed neo-prussian:"
MAX_NEW_TOKENS = 100

# ── All initialisation happens at module level (ZeroGPU best practice) ─────
# Tokenizer from the adapter repo (carries the chatml template).
_tokenizer = AutoTokenizer.from_pretrained(ADAPTER, trust_remote_code=True)

# Pre-quantized int8 base model — quantization_config is baked into config.json
# so no BitsAndBytesConfig is needed. device_map="auto" works with ZeroGPU's
# CUDA emulation at import time, then ZeroGPU swaps in real CUDA at request time.
_base = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL,
    device_map="auto",
    trust_remote_code=True,
)

# Attach the LoRA adapter using EXPLICIT CPU loading to avoid the safetensors
# CUDA path that fails in ZeroGPU's emulation mode (RuntimeError: No CUDA GPUs
# are available). PeftModel.from_pretrained() triggers __torch_dispatch__ hooks;
# loading weights manually on CPU bypasses that.
_config_path = hf_hub_download(ADAPTER, "adapter_config.json")
with open(_config_path) as f:
    _peft_config = LoraConfig(**json.load(f))

_model = PeftModel(_base, _peft_config)

_weights_path = hf_hub_download(ADAPTER, "adapter_model.safetensors")
_adapter_weights = safetensors.torch.load_file(_weights_path, device="cpu")

# Use PEFT's own loader, NOT _model.load_state_dict(). The safetensors keys are
# named `....lora_A.weight`, but the PeftModel expects the adapter name baked in
# (`....lora_A.default.weight`). A plain load_state_dict(strict=False) matches
# none of them and silently leaves lora_B at its zero-init — i.e. the adapter
# has zero effect. set_peft_model_state_dict performs the key remapping.
_load_result = set_peft_model_state_dict(_model, _adapter_weights)
assert not _load_result.unexpected_keys, (
    f"Adapter weights did not load: {_load_result.unexpected_keys[:5]}"
)
_model.eval()


@GPU(duration=20)
def translate(text: str) -> str:
    """Tokenize, generate, decode — the GPU-heavy work lives here."""
    if not text.strip():
        return ""

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": text.strip()},
    ]
    inputs = (
        _tokenizer.apply_chat_template(
            messages,
            add_generation_prompt=True,
            return_tensors="pt",
            return_dict=True,
        )
        .to(_base.device)
    )

    # <|im_end|> is NOT a single special token in the base tokenizer — the
    # model was trained to output it as subword pieces.  eos_token_id cannot
    # help here (it points to <|assistant_end|>, which the adapter never
    # generates).  We clean up after decoding instead.
    with torch.no_grad():
        out = _model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=False,
            repetition_penalty=1.2,
            pad_token_id=_tokenizer.eos_token_id,
        )
    new_tokens = out[0][inputs["input_ids"].shape[1] :]
    result = _tokenizer.decode(new_tokens, skip_special_tokens=False).strip()
    result = result.replace("<|im_end|>", "").strip()
    cutoff = result.find("<|im_start|>")
    if cutoff != -1:
        result = result[:cutoff].strip()
    return result


with gr.Blocks(title="New Prussian Translator") as demo:
    gr.Markdown(
        "# New Prussian Translator\n"
        "Apertus-8B + LoRA. Translates **into** reconstructed neo-Prussian "
        "from German, English, Lithuanian, Latvian, … Model is pre-loaded — "
        "queries complete in a few seconds."
    )
    gr.Markdown(f"**Fixed system prompt:** `{SYSTEM_PROMPT}`")
    with gr.Row():
        with gr.Column():
            text = gr.Textbox(lines=4, label="Source text")
            btn = gr.Button("Translate", variant="primary")
        with gr.Column():
            output = gr.Textbox(lines=4, label="New Prussian")

    gr.Examples(
        examples=[
            ["Ich gehe in den Wald"],
            ["All is very white."],
            ["Wie heißt du?"],
        ],
        inputs=[text],
    )

    btn.click(translate, text, output)
    text.submit(translate, text, output)

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