import os

os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"

# Import before torch so ZeroGPU initializes its integration.
import spaces

import copy
from pathlib import Path

import gradio as gr
import torch
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
REVISION = "7ae557604adf67be50417f59c2c2f167def9a775"
ADAPTER_DIR = Path(__file__).resolve().parent / "adapter"

SYSTEM_PROMPT = (
    "Reply like a friend in a casual chat. Keep replies short"
)

torch.set_num_threads(2)

for filename in ("adapter_config.json", "adapter_model.safetensors"):
    if not (ADAPTER_DIR / filename).is_file():
        raise FileNotFoundError(
            f"Missing adapter/{filename}. Add it to the Space repository."
        )

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_ID,
    revision=REVISION,
    trust_remote_code=False,
)

base = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    revision=REVISION,
    dtype=torch.float32,
    attn_implementation="eager",
    use_safetensors=True,
    trust_remote_code=False,
).to("cpu")

model = PeftModel.from_pretrained(
    base,
    str(ADAPTER_DIR),
    torch_device="cpu",
    is_trainable=False,
).eval()

# ZeroGPU supports CUDA placement at startup through its emulation.
# Real GPU execution happens inside the decorated function below.
model.to("cuda")


@spaces.GPU(duration=30)
def reply(text):
    if not isinstance(text, str) or not text.strip():
        raise gr.Error("Please enter a message.")

    text = text.strip()

    if any(token in text for token in tokenizer.all_special_tokens):
        raise gr.Error("Please remove tokenizer control tokens.")

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": text},
    ]

    inputs = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    )

    if inputs["input_ids"].shape[1] > 256:
        raise gr.Error("Please use a shorter message.")

    inputs = inputs.to("cuda")

    generation_config = copy.deepcopy(model.generation_config)
    generation_config.do_sample = False
    generation_config.temperature = 1.0
    generation_config.top_p = 1.0
    generation_config.top_k = 50

    with torch.inference_mode():
        output = model.generate(
            **inputs,
            generation_config=generation_config,
            max_new_tokens=128,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
            use_cache=True,
        )

    generated = output[0, inputs["input_ids"].shape[1]:].cpu()
    return tokenizer.decode(generated, skip_special_tokens=True)


demo = gr.Interface(
    fn=reply,
    inputs=gr.Textbox(
        label="Incoming message",
        placeholder="Kal chai pe milte?",
        lines=3,
    ),
    outputs=gr.Textbox(label="Model reply"),
    title="Hinglish Reply Tester",
    description=(
        "Test individual replies from the fine-tuned model. "
        "Responses may be inaccurate or contain memorized information."
    ),
    flagging_mode="never",
    analytics_enabled=False,
)

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()