"""Sample a continuation from DWSLM-256M (a base model: it continues text, it doesn't chat).

    python generate.py "The three primary colors are"
    python generate.py "def fibonacci(n):" --temperature 0 --max-new 80

There's no KV cache: each new token re-runs the whole prefix, which is fine at these lengths on a GPU.
"""

import argparse

import torch
from dwslm import EOS_ID, context, load, next_logits

UNFINISHED = "\ufffd"


def main() -> None:
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("prompt")
    p.add_argument("--max-new", type=int, default=64)
    p.add_argument("--temperature", type=float, default=0.7, help="0 = greedy")
    p.add_argument("--top-p", type=float, default=0.9)
    p.add_argument("--repetition-penalty", type=float, default=1.1, help="1 = off")
    p.add_argument("--repo", default=".", help="folder holding config.json and model.safetensors")
    a = p.parse_args()

    model, tok = load(a.repo)
    ids = [EOS_ID, *tok.encode(a.prompt).ids]
    start, shown = len(ids), ""
    print(a.prompt, end="", flush=True)
    for _ in range(min(a.max_new, context(model) - len(ids))):
        z = next_logits(model, [ids])[0]
        nxt = _sample(z, ids, a)
        if nxt == EOS_ID:
            break
        ids.append(nxt)

        # print only what's new; a byte-level token can end mid-character ("\ufffd") until the rest arrives
        text = tok.decode(ids[start:]).rstrip(UNFINISHED)
        print(text[len(shown) :], end="", flush=True)
        shown = text
    print(tok.decode(ids[start:])[len(shown) :])


def _sample(z: torch.Tensor, seen: list[int], a: argparse.Namespace) -> int:
    if a.repetition_penalty != 1.0:
        idx = torch.tensor(sorted(set(seen)), device=z.device)
        prev = z[idx]
        z[idx] = torch.where(prev > 0, prev / a.repetition_penalty, prev * a.repetition_penalty)
    if a.temperature <= 0:
        return int(z.argmax())

    probs = (z / a.temperature).softmax(-1)
    p, order = probs.sort(descending=True)
    p = p * (p.cumsum(-1) - p < a.top_p)
    return int(order[torch.multinomial(p / p.sum(), 1)])


if __name__ == "__main__":
    main()