# app.py — Gradio Space для OpenRussianAI/andrey (ZeroGPU)
import torch
import gradio as gr
import spaces
from transformers.dynamic_module_utils import get_class_from_dynamic_module
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer

REPO = "OpenRussianAI/andrey"

# ================= ЗАГРУЗКА =================
print("[init] загружаю конфиг и модель...")

_cfg_cls = get_class_from_dynamic_module(
    "configuration_easyformer.EasyFormerConfig", REPO,
)
_model_cls = get_class_from_dynamic_module(
    "modeling_easyformer.EasyFormerLMHeadModel", REPO,
)

_cfg = _cfg_cls.from_pretrained(REPO)
_model = _model_cls.from_pretrained(REPO, config=_cfg).to("cuda").eval()

MAX_CTX = int(getattr(_cfg, "ctx", 512))
print(f"[init] модель загружена, ctx={MAX_CTX}")

# ================= ТОКЕНИЗАТОР =================
_tok = Tokenizer.from_file(hf_hub_download(REPO, "tokenizer.json"))
_vocab_size = _tok.get_vocab_size()
print(f"[init] vocab={_vocab_size}")

USER, BOT, EOS = "[USER]", "[BOT]", "[EOS]"
_eos_ids = set()
for t in (EOS, "<eos>", "</s>", "<|endoftext|>"):
    tid = _tok.token_to_id(t)
    if tid is not None:
        _eos_ids.add(tid)

STOP_MARKERS = (USER, BOT)

def _enc(s):
    return _tok.encode(s).ids

def _dec(ids):
    """Восстанавливает пробелы между токенами (BPE обучен на Whitespace())."""
    out = ""
    for tid in ids:
        piece = _tok.decode([tid])
        if not piece:
            continue
        no_space = (
            not out
            or out.endswith(" ")
            or piece.startswith(" ")
            or piece in ".,!?;:)]}"
            or piece.startswith(("'", "’", "-"))
        )
        if not no_space:
            out += " "
        out += piece
    return out

def _clean(text):
    cut = len(text)
    for m in STOP_MARKERS:
        i = text.find(m)
        if i != -1 and i < cut:
            cut = i
    for t in (USER, BOT, EOS, "[IDK]", "[PAD]", "[UNK]", "[BOS]"):
        text = text.replace(t, "")
    return text[:cut].strip()

def _build_prompt(message, history, system_message, max_turns=2):
    parts = []
    if system_message and system_message.strip():
        parts.append(f"{USER} {system_message.strip()}{EOS}{BOT} ok{EOS}")
    recent = history[-max_turns * 2:] if len(history) > max_turns * 2 else history
    for turn in recent:
        role = turn.get("role")
        content = turn.get("content", "")
        if role == "user":
            parts.append(f"{USER} {content}{EOS}")
        elif role == "assistant":
            parts.append(f"{BOT} {content}{EOS}")
    parts.append(f"{USER} {message}{EOS}{BOT}")
    return "".join(parts)


# ================= ГЕНЕРАЦИЯ =================
@spaces.GPU(duration=30)
def generate_response(
    message, history, system_message,
    max_tokens, temperature, top_p, top_k,
    repetition_penalty, presence_penalty, frequency_penalty,
):
    prompt = _build_prompt(message, history, system_message)
    ids = _enc(prompt)
    if len(ids) >= MAX_CTX - 4:
        ids = ids[-(MAX_CTX - 8):]

    input_ids = torch.tensor([ids], dtype=torch.long, device="cuda")
    prompt_len = input_ids.size(1)
    max_new = min(int(max_tokens), MAX_CTX - prompt_len - 1)
    if max_new <= 0:
        return "(контекст переполнен)"

    generated = input_ids
    response_ids = []

    for _ in range(max_new):
        if generated.size(1) >= MAX_CTX:
            break

        with torch.no_grad():
            out = _model(input_ids=generated)
        logits = (out.logits if hasattr(out, "logits") else out["logits"])[:, -1, :]
        logits = logits / max(temperature, 1e-5)

        if repetition_penalty != 1.0 and generated.size(1) > prompt_len:
            gen = generated[0, prompt_len:]
            counts = torch.bincount(gen, minlength=_vocab_size).to(logits.dtype)
            mask = counts > 0
            penalized = torch.where(
                logits > 0, logits / repetition_penalty, logits * repetition_penalty,
            )
            logits = torch.where(mask.unsqueeze(0), penalized, logits)

        if top_k and top_k > 0:
            k = min(int(top_k), logits.size(-1))
            v, _ = torch.topk(logits, k)
            logits[logits < v[:, [-1]]] = float("-inf")

        if top_p and 0 < top_p < 1.0:
            sorted_logits, sorted_idx = torch.sort(logits, descending=True)
            probs = torch.softmax(sorted_logits, dim=-1)
            cumprobs = torch.cumsum(probs, dim=-1)
            mask = cumprobs > top_p
            mask[..., 1:] = mask[..., :-1].clone()
            mask[..., 0] = False
            sorted_logits = sorted_logits.masked_fill(mask, float("-inf"))
            logits = torch.empty_like(logits).scatter_(-1, sorted_idx, sorted_logits)

        probs = torch.softmax(logits, dim=-1)
        if not torch.isfinite(probs).all() or probs.sum() <= 0:
            break
        next_id = torch.multinomial(probs, num_samples=1)
        token_id = int(next_id.item())

        if token_id in _eos_ids and response_ids:
            break

        generated = torch.cat([generated, next_id], dim=-1)
        response_ids.append(token_id)

        # ранний стоп по маркерам в сыром тексте
        piece = _tok.decode([token_id])
        if any(m in piece for m in STOP_MARKERS):
            break

    return _clean(_dec(response_ids))


def respond(
    message, history, system_message,
    max_tokens, temperature, top_p, top_k,
    repetition_penalty, presence_penalty, frequency_penalty,
):
    return generate_response(
        message, history, system_message,
        max_tokens, temperature, top_p, top_k,
        repetition_penalty, presence_penalty, frequency_penalty,
    )


chatbot = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="", label="System message"),
        gr.Slider(1, 100, value=40, step=1, label="Max new tokens"),
        gr.Slider(0.1, 2.0, value=0.9, step=0.05, label="Temperature"),
        gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p"),
        gr.Slider(0, 100, value=30, step=1, label="Top-k (0 = выкл)"),
        gr.Slider(1.0, 2.0, value=1.2, step=0.05, label="Repetition penalty"),
        gr.Slider(0.0, 2.0, value=0.1, step=0.05, label="Presence penalty"),
        gr.Slider(0.0, 2.0, value=0.1, step=0.05, label="Frequency penalty"),
    ],
    title="Andrey — EasyFormer 3.65M",
    description="Лёгкий русскоязычный чат-бот на EasyFormer. ZeroGPU.",
    examples=[["привет"], ["кто ты"], ["как дела"], ["пока"]],
)

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