"""Load the ShellElastic supernet of this repository on the CPU, run one sub-network and extract it.

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_shellelastic import ShellElasticForCausalLM  # 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_shellelastic.py) to DIR")
    args = parser.parse_args()

    model = ShellElasticForCausalLM.from_pretrained(HERE, device="cpu", dtype=getattr(torch, args.dtype))
    cfg = model.config
    print(f"supernet: {model.num_parameters():,} stored parameters, {args.dtype}")
    print(f"{args.tier}: width {cfg.width(args.tier)}, {cfg.width(args.tier) // cfg.head_dim} heads, "
          f"{model.num_parameters(args.tier):,} parameters")

    ids = torch.tensor([PROMPT_IDS])
    with torch.no_grad():
        logits, loss = model(ids[:, :-1], tier=args.tier, labels=ids[:, 1:])
        all_tiers = model(ids, tiers=cfg.tiers)          # every sub-network from one full-width trunk pass
        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.float()).all() and torch.isfinite(loss)
    print(f"prompt ids: {PROMPT_IDS}")
    print(f"mean next-token loss on the prompt: {loss.item():.4f}")
    print("next-token id after the prompt, per sub-network: "
          + ", ".join(f"{t} {all_tiers[t][0, -1].argmax().item()}" for t in cfg.tiers))
    print(f"greedy continuation ids: {out[0, len(PROMPT_IDS):].tolist()}")

    # The sub-network as a standalone model: its embedding columns, the core and the shells it contains, and its own
    # terminal norm and readout. By Proposition 1 of the paper its logits equal the supernet's for this tier.
    small = model.extract(args.tier)
    with torch.no_grad():
        diff = (small(ids, tier=args.tier).float() - model(ids, tier=args.tier).float()).abs().max().item()
    print(f"extracted {args.tier}: {small.num_parameters():,} parameters, max |logit difference| to the "
          f"supernet {diff:g}")
    assert diff == 0.0
    if args.save:
        small.save_pretrained(args.save)
        print(f"wrote {os.path.join(args.save, 'config.json')} and {os.path.join(args.save, 'model.safetensors')}")


if __name__ == "__main__":
    main()