"""Load the supernet of this repository on the CPU, run one sub-network and extract it as a dense model.

Usage, from the downloaded repository:

    python example.py                      # the Low sub-network, float32
    python example.py --tier max --dtype bfloat16
    python example.py --tier low --save low_model   # also write the extracted Low model to ./low_model

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

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_elastic import ElasticForCausalLM  # noqa: E402

# 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 and extract one sub-network of this supernet on the CPU.")
    parser.add_argument("--tier", default="low", choices=["low", "medium", "high", "extra", "max"],
                        help="sub-network (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")
    parser.add_argument("--save", metavar="DIR", help="write the extracted sub-network (config.json and "
                        "model.safetensors, loadable with modeling_dense.py) to DIR")
    args = parser.parse_args()

    model = ElasticForCausalLM.from_pretrained(HERE, device="cpu", dtype=getattr(torch, args.dtype))
    cfg = model.config
    print(f"supernet: {cfg.variant}, {model.num_parameters():,} stored parameters, {args.dtype}")
    print(f"{args.tier}: width {cfg.width(args.tier)}, {cfg.num_heads(args.tier)} heads, FFN width "
          f"{cfg.intermediate_size(args.tier)}, {model.num_parameters(args.tier):,} parameters as a dense model")

    ids = torch.tensor([PROMPT_IDS])
    with torch.no_grad():
        logits, loss = model(ids[:, :-1], tier=args.tier, labels=ids[:, 1:])
        out = model.generate(ids, max_new_tokens=args.new_tokens, tier=args.tier)
    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()}")

    # The sub-network as an ordinary dense model: the tier's blocks of every matrix with its multipliers applied,
    # the FFN divisor e folded into down_proj and the readout scalar alpha folded into lm_head.
    weights = model.extract(args.tier)
    n_extracted = sum(t.numel() for t in weights.values())
    assert n_extracted == model.num_parameters(args.tier)
    print(f"extracted {args.tier}: {len(weights)} tensors, {n_extracted:,} parameters")
    if args.save:
        model.save_subnetwork(args.tier, args.save)
        print(f"wrote {os.path.join(args.save, 'config.json')} and {os.path.join(args.save, 'model.safetensors')}")


if __name__ == "__main__":
    main()