"""Generate text from a trained Eather-500M checkpoint."""
import argparse
import json
from pathlib import Path
import torch
import torch.nn.functional as F
from tokenizers import Tokenizer
from model import ModelConfig, Eather500M


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--checkpoint", required=True)
    p.add_argument("--tokenizer", default="tokenizer.json")
    p.add_argument("--prompt", default="Hello!")
    p.add_argument("--chat", action="store_true")
    p.add_argument("--system", default="You are a helpful assistant.")
    p.add_argument("--max-new-tokens", type=int, default=128)
    p.add_argument("--temperature", type=float, default=0.8)
    p.add_argument("--top-k", type=int, default=50)
    args = p.parse_args()
    tok = Tokenizer.from_file(args.tokenizer)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    ck = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    cfg = ModelConfig(**ck["config"])
    model = Eather500M(cfg)
    model.load_state_dict(ck["model"])
    del ck
    model.to(device).eval()
    prompt = args.prompt
    if args.chat:
        prompt = f"<|system|>{args.system}<|end|><|user|>{args.prompt}<|end|><|assistant|>"
    ids = tok.encode(prompt, add_special_tokens=False).ids
    if not ids:
        ids = [tok.token_to_id("<|bos|>")]
    stop = {tok.token_to_id("<|end|>"), tok.token_to_id("<|eos|>")}
    result = []
    print("\nOutput:\n", end="", flush=True)
    with torch.inference_mode():
        for _ in range(args.max_new_tokens):
            inp = torch.tensor([ids[-cfg.max_seq_len:]], device=device)
            with torch.autocast(device_type="cuda", dtype=torch.bfloat16,
                                enabled=(device == "cuda" and torch.cuda.is_bf16_supported())):
                logits, _ = model(inp)
            next_logits = logits[0, -1].float()
            if args.temperature <= 0:
                next_id = next_logits.argmax().item()
            else:
                next_logits /= args.temperature
                if args.top_k > 0:
                    v, _ = torch.topk(next_logits, min(args.top_k, next_logits.numel()))
                    next_logits[next_logits < v[-1]] = -float("inf")
                next_id = torch.multinomial(F.softmax(next_logits, dim=-1), 1).item()
            if next_id in stop:
                break
            ids.append(next_id)
            result.append(next_id)
    print(tok.decode(result, skip_special_tokens=True))


if __name__ == "__main__":
    main()