"""Load one model of this repository on the CPU and run it.

Usage, from the downloaded repository:

    python example.py                      # the Low model, float32
    python example.py --tier max --dtype bfloat16

Needs only torch and safetensors (see requirements.txt). Nothing is written to disk.
"""

import argparse
import os
import sys

sys.dont_write_bytecode = True  # leave the downloaded folder unchanged
HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

import torch  # noqa: E402

from modeling_dense import DenseDecoder  # noqa: E402

TIERS = ["low", "medium", "high", "extra", "max"]
# GPT-NeoX-20B token ids of "The capital of France is". Any copy of that tokenizer gives them, for example
# tokenizers.Tokenizer.from_pretrained("EleutherAI/gpt-neox-20b").encode("The capital of France is").ids
PROMPT_IDS = [510, 5347, 273, 6181, 310]


def main() -> None:
    parser = argparse.ArgumentParser(description="Run one model of this repository on the CPU.")
    parser.add_argument("--tier", default="low", choices=TIERS, help="model folder to load (default: low)")
    parser.add_argument("--dtype", default="float32", choices=["float32", "bfloat16"],
                        help="compute precision (the released weights are float32)")
    parser.add_argument("--new-tokens", type=int, default=8, help="greedy tokens to generate")
    args = parser.parse_args()

    model = DenseDecoder.from_pretrained(os.path.join(HERE, args.tier), device="cpu",
                                         dtype=getattr(torch, args.dtype))
    cfg = model.config
    n_params = sum(p.numel() for p in model.parameters())
    print(f"{args.tier}: width {cfg.hidden_size}, {cfg.num_attention_heads} heads, FFN width "
          f"{cfg.intermediate_size}, {n_params:,} parameters, {args.dtype}")

    ids = torch.tensor([PROMPT_IDS])
    with torch.no_grad():
        logits, loss = model(ids[:, :-1], labels=ids[:, 1:])
        out = model.generate(ids, max_new_tokens=args.new_tokens)
    assert logits.shape == (1, len(PROMPT_IDS) - 1, cfg.vocab_size)
    assert torch.isfinite(logits).all() and torch.isfinite(loss)
    print(f"prompt ids: {PROMPT_IDS}")
    print(f"mean next-token loss on the prompt: {loss.item():.4f}")
    print(f"greedy continuation ids: {out[0, len(PROMPT_IDS):].tolist()}")


if __name__ == "__main__":
    main()