# app.py
import gradio as gr
import spaces
import torch
import tiktoken
from huggingface_hub import hf_hub_download
from collections import OrderedDict

from model import GPT, ModelConfig
from inference import generate_stream

# -------------------------
# CPU 上でモデルロード（ZeroGPU重要）
# -------------------------
# Hugging Face からダウンロード
model_path = hf_hub_download(
    repo_id="HayatoHongo/everyoneschat-checkpoints",
    filename="model.pt"
)

# state_dict をロード
state_dict = torch.load(model_path, map_location="cpu")

cfg = checkpoint["config"]
config = ModelConfig(
    embedding_dim=cfg["embedding_dim"],
    hidden_dim=cfg["hidden_dim"],
    num_attention_heads=cfg["num_attention_heads"],
    layer_count=cfg["layer_count"],
    max_sequence_length=cfg["max_sequence_length"],
    rope_theta=cfg["rope_theta"],
    vocab_size=cfg["vocab_size"],
)

# モデル生成 & load
model = GPT(config)
model.load_state_dict(state_dict)
model.eval()

tokenizer = tiktoken.get_encoding("gpt2")
EOS_ID = 50256  # GPT-2 EOS

# -------------------------
# GPU を使う関数だけ ZeroGPU で囲む
# -------------------------
@spaces.GPU
def chat_fn(
    message,
    history,
    temperature,
    top_p,
    top_k,
):
    device = "cuda"
    model_gpu = model.to(device)

    # シングルターンなので毎回 cache を完全リセット
    for block in model_gpu.blocks:
        block.multihead_attention.reset_cache()

    # ---- ここが超シンプルな prompt 整形 ----
    prompt = (
        "<user>\n"
        f"{message}\n"
        "<assistant>\n"
    )

    input_ids = torch.tensor(
        [tokenizer.encode(prompt, allowed_special="all")],
        device=device
    )

    output = ""

    with torch.no_grad(), torch.autocast(
        device_type="cuda",
        dtype=torch.bfloat16,
    ):
        for tid in generate_stream(
            model_gpu,
            input_ids,
            max_new_tokens=256,
            temperature=temperature,
            top_p=top_p if top_p > 0 else None,
            top_k=top_k if top_k > 0 else None,
        ):
            if tid == EOS_ID:
                break
            output += tokenizer.decode([tid])

    model_gpu.to("cpu")
    torch.cuda.empty_cache()

    return output


# -------------------------
# UI 定義
# -------------------------
demo = gr.ChatInterface(
    chat_fn,
    title="EveryonesGPT Pretrained (No Instruction-tuning). Single-turn English-only demo.",
    description=(
        "**Try prompts like:**\n"
        "- What is the capital city of Japan?\n"
        "- What is the element symbol of silver?\n"
        "- Explain AI in simple terms"
    ),
    additional_inputs=[
        gr.Slider(0.1, 2.0, value=0.7, step=0.05, label="Temperature"),
        gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-p"),
        gr.Slider(0, 200, value=0, step=1, label="Top-k"),
    ],
)

demo.launch()