"""GoLLeM v6 Instruct: chat helper (ChatML), matching the format and sampling used in training and evaluation.

    import sys; sys.path.insert(0, path)
    from modeling_gollem_v6 import load_gollem_v6
    from chat_gollem_v6 import chat
    model, tok = load_gollem_v6(path)
    print(chat(model, tok, [("user", "Kim jesteś?")], seed=1))

Format: for every message `<|im_start|>{role}\n{content}<|im_end|>\n`, then `<|im_start|>assistant\n`; no system prompt;
generation stops at `<|im_end|>`. Message text is encoded as plain text (the strings "<|im_start|>"/"<|im_end|>" typed by
a user are NOT turned into control tokens). Default sampling: temperature 0.7, nucleus top-p 0.9 (as in our evaluation).
"""
from __future__ import annotations

import torch

IM_START, IM_END = 1, 2


def encode_chat(tok, messages):
    """Token ids of a ChatML conversation ending with an open assistant turn. messages: [(role, text), ...]."""
    prev = tok.encode_special_tokens
    tok.encode_special_tokens = True
    try:
        text = lambda s: tok.encode(s, add_special_tokens=False).ids  # noqa: E731
        ids = []
        for role, content in messages:
            ids += [IM_START] + text(role + "\n") + text(content) + [IM_END] + text("\n")
        return ids + [IM_START] + text("assistant\n")
    finally:
        tok.encode_special_tokens = prev


@torch.no_grad()
def chat(model, tok, messages, max_new_tokens=256, temperature=0.7, top_p=0.9, seed=None):
    """Assistant reply to `messages` ([(role, text), ...], roles "user"/"assistant"). Returns the reply text."""
    g = torch.Generator().manual_seed(seed) if seed is not None else None
    dev = next(model.parameters()).device
    x = torch.tensor([encode_chat(tok, messages)], device=dev)
    out = []
    for _ in range(max_new_tokens):
        logits = model(x[:, -model.block:])[0, -1].float()
        p = torch.softmax(logits / temperature, -1)
        sp, si = torch.sort(p, descending=True)
        sp = sp * (torch.cumsum(sp, 0) - sp < top_p)
        nxt = int(si[torch.multinomial((sp / sp.sum()).cpu(), 1, generator=g)])
        if nxt == IM_END:
            break
        out.append(nxt)
        x = torch.cat([x, torch.tensor([[nxt]], device=dev)], 1)
    return tok.decode(out)