#!/usr/bin/env python3
"""Export every parameter name, shape, dtype and gradient flag from a meta model."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

from configuration_mini_k3 import MiniK3Config
from modeling_mini_k3 import MiniK3ForCausalLM


def main() -> None:
    here = Path(__file__).resolve().parent
    parser = argparse.ArgumentParser()
    parser.add_argument("--config", type=Path, default=here / "config.json")
    parser.add_argument("--output", type=Path, default=here / "parameter_schema.jsonl")
    args = parser.parse_args()

    config = MiniK3Config.from_json(args.config)
    model = MiniK3ForCausalLM(config, device="meta")
    logical = 0
    tensors = 0
    with args.output.open("w") as stream:
        for name, parameter in model.named_parameters():
            record = {
                "name": name,
                "shape": list(parameter.shape),
                "dtype": str(parameter.dtype).removeprefix("torch."),
                "requires_grad": parameter.requires_grad,
                "numel": parameter.numel(),
            }
            stream.write(json.dumps(record, separators=(",", ":")) + "\n")
            logical += parameter.numel()
            tensors += 1
    assert logical == config.expected_logical_parameters
    print(f"wrote {tensors:,} tensors / {logical:,} elements to {args.output}")


if __name__ == "__main__":
    main()