"""Generate a tiny, random, deterministic DeepSeek-V4.1-Flash HF checkpoint.

Same tensor names, dtypes and layout conventions as the released checkpoint
(deepseek-ai/DeepSeek-V4.1-Flash-1017); only the sizes shrink. The result is meant to be loadable by
both DeepInfer's HF loader (deepinfer/hugging_face/*) and the official reference
(inference/convert.py + inference/model.py).

Conventions (taken from the official inference/ code):
  * FP8 weights: F8_E4M3 [out, in] + F8_E8M0 scale [ceil(out/32), ceil(in/32)] (32x32 blocks).
    scale = 2**ceil(log2(amax/448)), q = e4m3(clamp(w/scale, +-448)); dequant = q * scale.
  * Routed experts: FP4 e2m1, two codes per byte stored as I8 [out, in/2]; element 2j is the LOW nibble,
    element 2j+1 the HIGH nibble (convert.py cast_e2m1fn_to_e4m3fn: low = x & 0xF -> first, high -> second);
    code table = FP4_TABLE in convert.py (bit 3 = sign, 0..7 = 0,.5,1,1.5,2,3,4,6).
    F8_E8M0 scale [out, in/32] (1x32 blocks along K), scale = 2**ceil(log2(amax/6)), RNE to the e2m1 grid.
  * Engram table: F8_E4M3 [rows, head_dim] + F8_E8M0 [rows, head_dim/32] (1x32 blocks), rows =
    sum of the per-(ngram, head) prime bucket sizes (EngramLayout in DeepInfer and in inference/engram.py).
  * hc_* and attn_sink and gate biases are F32; norms, compressor, indexer wk/weights_proj, gate weight,
    embed/head, vision, markov/confidence heads are BF16 (as in the real index).

Random weights: W ~ N(0, 1/in_features) in fp32 -> bf16 (-> quantized), norms = 1, biases/sinks small; the
projections that write into the residual stream get an extra gain (attn wo_b: --attn-out-scale 0.5; shared/routed
expert w2 and engram wkv: --ffn-out-scale 1/16) so that the reference is numerically stable enough to be a golden.
Every tensor uses its own generator seeded from (seed, crc32(name)), so the content of one tensor does not
depend on which other tensors exist.

Usage:
    python make_tiny_checkpoint.py --out ~/di/tiny_ckpt --tokenizer-dir <dir with tokenizer.json> \
        --real-config <real HF config.json>
"""

from __future__ import annotations

import argparse
import copy
import json
import math
import os
import shutil
import zlib
from collections import OrderedDict

import torch
from safetensors.torch import save_file

# ----------------------------------------------------------------------------------------------
# Tiny config. Keys not listed keep the value from the real HF config.
# ----------------------------------------------------------------------------------------------
TINY_TEXT_OVERRIDES = dict(
    hidden_size=1024,
    moe_intermediate_size=256,
    num_hidden_layers=10,
    num_attention_heads=8,
    head_dim=512,
    qk_rope_head_dim=64,
    q_lora_rank=256,
    o_lora_rank=256,
    o_groups=2,
    n_routed_experts=16,
    num_experts_per_tok=4,
    # 10 backbone layers + 3 nextn layers. Layout mirrors the real one: 2 SWA-only layers, a ratio-2
    # segment, a ratio-1 segment starting at the candidate source, 3 SWA-only MTP layers.
    compress_ratios=[0, 0, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0],
    kv_source_layer_ids=[2, 4, 6],
    index_source_layer_ids=[2, 4, 6, 8],
    candidate_source_layer_id=6,
    # DeepInfer: candidate_tokens * sizeof(bf16 psum) % 1024 == 0 (index_attn.py:109) -> topk_blocks*8 % 512 == 0,
    # so 64 is the smallest legal value with block size 8.
    candidate_topk_blocks=64,
    index_n_heads=4,
    index_head_dim=128,
    index_topk=64,
    engram_layer_ids=[1, 4],
    engram_vocab_size=4096,
    engram_n_heads=2,
    engram_head_dim=256,
    num_nextn_predict_layers=3,
    dspark_target_layer_ids=[7, 8, 9],
    dspark_n_routed_experts=8,
    dspark_num_experts_per_tok=2,
)
# DeepInfer's VisionTransformer hard-codes DEPTH=32 / dim 1024 / 16 heads / mlp 2816
# (deepinfer/model_executor/models/vision_encoder/vit.py), always builds it, and the loader requires every
# model parameter to be covered (model_loader.load_hf_checkpoint). So the vision tower keeps its real size.
TINY_VISION_OVERRIDES = dict()

FP8_MAX = 448.0
FP4_MAX = 6.0
BLOCK = 32
FP4_TABLE = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0])


# ----------------------------------------------------------------------------------------------
# Quantization helpers (match inference/kernel.py fast_round_scale: 2**ceil(log2(amax * (1/max))))
# ----------------------------------------------------------------------------------------------
def pow2_ceil_scale(amax: torch.Tensor, max_inv: float) -> torch.Tensor:
    """2**ceil(log2(amax * max_inv)) computed from the fp32 bit pattern, like kernel.py fast_log2_ceil."""
    v = (amax.float() * torch.tensor(max_inv, dtype=torch.float32)).contiguous()
    bits = v.view(torch.int32)
    exp = (bits >> 23) & 0xFF
    man = bits & ((1 << 23) - 1)
    e = exp - 127 + (man != 0).to(torch.int32)
    return torch.ldexp(torch.ones_like(v), e)


def to_e8m0(scale: torch.Tensor) -> torch.Tensor:
    """Power-of-two fp32 -> float8_e8m0fnu (exact)."""
    bits = (scale.float().contiguous().view(torch.int32) >> 23) & 0xFF
    return bits.to(torch.uint8).view(torch.float8_e8m0fnu)


def quant_fp8_block(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """[out, in] -> e4m3 weight + e8m0 scale [ceil(out/32), ceil(in/32)]."""
    out_f, in_f = w.shape
    bo, bi = math.ceil(out_f / BLOCK), math.ceil(in_f / BLOCK)
    wp = torch.zeros(bo * BLOCK, bi * BLOCK, dtype=torch.float32)
    wp[:out_f, :in_f] = w.float()
    blocks = wp.view(bo, BLOCK, bi, BLOCK)
    amax = blocks.abs().amax(dim=(1, 3)).clamp_min(1e-4)
    s = pow2_ceil_scale(amax, 1.0 / FP8_MAX)
    q = (blocks / s[:, None, :, None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn)
    q = q.view(bo * BLOCK, bi * BLOCK)[:out_f, :in_f].contiguous()
    return q, to_e8m0(s)


def round_e2m1_codes(v: torch.Tensor) -> torch.Tensor:
    """fp32 (already scaled, |v|<=6) -> 4-bit e2m1 codes, round-to-nearest-even."""
    a = v.abs()
    code = torch.zeros_like(a, dtype=torch.int32)
    # midpoints whose tie goes DOWN to an even code: 0.25(0|1) 1.25(2|3) 2.5(4|5) 5.0(6|7)
    for b in (0.25, 1.25, 2.5, 5.0):
        code += (a > b).to(torch.int32)
    # midpoints whose tie goes UP to an even code: 0.75(1|2) 1.75(3|4) 3.5(5|6)
    for b in (0.75, 1.75, 3.5):
        code += (a >= b).to(torch.int32)
    sign = (v < 0) | ((v == 0) & torch.signbit(v))
    return code | (sign.to(torch.int32) << 3)


def quant_fp4_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """[out, in] -> packed int8 [out, in/2] (low nibble = even k) + e8m0 scale [out, in/32]."""
    out_f, in_f = w.shape
    assert in_f % BLOCK == 0
    blocks = w.float().view(out_f, in_f // BLOCK, BLOCK)
    amax = blocks.abs().amax(-1).clamp_min(FP4_MAX * 2.0**-126)
    s = pow2_ceil_scale(amax, 1.0 / FP4_MAX)
    codes = round_e2m1_codes((blocks / s[..., None]).clamp(-FP4_MAX, FP4_MAX)).view(out_f, in_f)
    packed = (codes[:, 0::2] | (codes[:, 1::2] << 4)).to(torch.uint8).view(torch.int8).contiguous()
    return packed, to_e8m0(s)


def dequant_fp4_row(packed: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    x = packed.view(torch.uint8)
    lo, hi = (x & 0xF).long(), (x >> 4).long()
    vals = torch.stack([FP4_TABLE[lo], FP4_TABLE[hi]], dim=-1).flatten(-2)
    return vals * scale.float().repeat_interleave(BLOCK, dim=-1)


def quant_fp8_row(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Engram table: [rows, d] -> e4m3 + e8m0 [rows, d/32] (1x32 blocks)."""
    rows, d = w.shape
    blocks = w.float().view(rows, d // BLOCK, BLOCK)
    amax = blocks.abs().amax(-1).clamp_min(1e-4)
    s = pow2_ceil_scale(amax, 1.0 / FP8_MAX)
    q = (blocks / s[..., None]).clamp(-FP8_MAX, FP8_MAX).to(torch.float8_e4m3fn).view(rows, d)
    return q.contiguous(), to_e8m0(s)


# ----------------------------------------------------------------------------------------------
# Engram layout (identical in DeepInfer engram_metadata.EngramLayout.from_cfg and inference/engram.py)
# ----------------------------------------------------------------------------------------------
def engram_primes(layer_ids, max_ngram_size, n_heads, vocab_size):
    from sympy import isprime, nextprime

    # reference (inference/engram.py find_next_prime)
    seen, ref = set(), []
    for _ in layer_ids:
        per_ngram = []
        for _ in range(max_ngram_size - 1):
            sizes, current = [], vocab_size - 1
            for _ in range(n_heads):
                c = current + 1
                while not isprime(c) or c in seen:
                    c += 1
                current = c
                seen.add(c)
                sizes.append(c)
            per_ngram.append(sizes)
        ref.append(per_ngram)
    # DeepInfer (engram_metadata.EngramLayout.from_cfg)
    seen, di = set(), []
    for _ in layer_ids:
        per_ngram = []
        for _ in range(max_ngram_size - 1):
            start, sizes = vocab_size - 1, []
            for _ in range(n_heads):
                p = int(nextprime(start))
                while p in seen:
                    p = int(nextprime(p))
                seen.add(p)
                sizes.append(p)
                start = p
            per_ngram.append(sizes)
        di.append(per_ngram)
    assert ref == di, (ref, di)
    return ref


# ----------------------------------------------------------------------------------------------
# Tensor factory
# ----------------------------------------------------------------------------------------------
class Factory:
    def __init__(self, seed: int, norm_noise: float):
        self.seed = seed
        self.norm_noise = norm_noise
        self.tensors: "OrderedDict[str, torch.Tensor]" = OrderedDict()

    def gen(self, name: str) -> torch.Generator:
        g = torch.Generator()
        g.manual_seed((self.seed * 1_000_003 + zlib.crc32(name.encode())) % (2**63))
        return g

    def randn(self, name, shape, std):
        return torch.randn(shape, generator=self.gen(name), dtype=torch.float32) * std

    def add(self, name, t):
        assert name not in self.tensors, name
        self.tensors[name] = t.contiguous()

    def bf16(self, name, out_f, in_f, std=None):
        std = (1.0 / in_f) ** 0.5 if std is None else std
        self.add(name, self.randn(name, (out_f, in_f), std).bfloat16())

    def norm(self, name, dim):
        w = torch.ones(dim)
        if self.norm_noise:
            w = w + self.randn(name, (dim,), self.norm_noise)
        self.add(name, w.bfloat16())

    def bias(self, name, dim, std=0.02, dtype=torch.bfloat16):
        self.add(name, self.randn(name, (dim,), std).to(dtype))

    def fp8(self, prefix, out_f, in_f, gain=1.0):
        w = self.randn(prefix + ".weight", (out_f, in_f), gain * (1.0 / in_f) ** 0.5).bfloat16()
        q, s = quant_fp8_block(w)
        self.add(prefix + ".weight", q)
        self.add(prefix + ".scale", s)

    def fp4(self, prefix, out_f, in_f, gain=1.0):
        w = self.randn(prefix + ".weight", (out_f, in_f), gain * (1.0 / in_f) ** 0.5).bfloat16()
        q, s = quant_fp4_row(w)
        self.add(prefix + ".weight", q)
        self.add(prefix + ".scale", s)

    def f32(self, name, t):
        self.add(name, t.float())


def build_tensors(cfg: dict, seed: int, norm_noise: float, attn_out_scale: float, ffn_out_scale: float) -> "OrderedDict[str, torch.Tensor]":
    t = cfg["text_config"]
    v = cfg["vision_config"]
    f = Factory(seed, norm_noise)
    dim, H, hd = t["hidden_size"], t["num_attention_heads"], t["head_dim"]
    qr, olr, og = t["q_lora_rank"], t["o_lora_rank"], t["o_groups"]
    inter, hc = t["moe_intermediate_size"], t["hc_mult"]
    n_layer, nextn = t["num_hidden_layers"], t["num_nextn_predict_layers"]
    ratios = t["compress_ratios"]
    kv_src, idx_src = set(t["kv_source_layer_ids"]), set(t["index_source_layer_ids"])
    ih, ihd = t["index_n_heads"], t["index_head_dim"]
    vocab = t["vocab_size"]
    mix_hc = (2 + hc) * hc

    # ---- vision tower + aligner (BF16), real shapes ----
    vd, vinter, ps = v["hidden_size"], v["intermediate_size"], v["patch_size"]
    f.bf16("vision.patch_embed.proj.weight", vd, 3 * ps * ps)
    f.bias("vision.patch_embed.proj.bias", vd)
    for i in range(v["num_hidden_layers"]):
        p = f"vision.blocks.{i}."
        f.norm(p + "norm1.weight", vd)
        f.bf16(p + "attn.wqkv.weight", 3 * vd, vd)
        f.bias(p + "attn.wqkv.bias", 3 * vd)
        f.bf16(p + "attn.wo.weight", vd, vd)
        f.bias(p + "attn.wo.bias", vd)
        f.norm(p + "norm2.weight", vd)
        f.bf16(p + "mlp.w1.weight", 2 * vinter, vd)
        f.bf16(p + "mlp.w2.weight", vd, vinter)
    f.norm("vision.norm.weight", vd)
    ds = v["downsample_ratio"]
    f.bf16("aligner.w1.weight", dim, vd * ds * ds)
    f.bias("aligner.w1.bias", dim)
    f.bf16("aligner.w2.weight", dim, dim)
    f.bias("aligner.w2.bias", dim)

    # ---- global ----
    f.bf16("embed.weight", vocab, dim, std=1.0)
    for n in ("image_start", "image_end", "image_newline"):
        f.add(n, f.randn(n, (dim,), 1.0).bfloat16())
    f.norm("norm.weight", dim)
    f.bf16("head.weight", vocab, dim)

    def hc_params(p):
        for which in ("attn", "ffn"):
            f.f32(p + f"hc_{which}_fn", f.randn(p + f"hc_{which}_fn", (mix_hc, hc * dim), (1.0 / (hc * dim)) ** 0.5))
            f.f32(p + f"hc_{which}_scale", torch.full((3,), 0.3) + f.randn(p + f"hc_{which}_scale", (3,), 0.05))
            base = f.randn(p + f"hc_{which}_base", (mix_hc,), 0.5)
            base[2 * hc :] += 2.0 * torch.eye(hc).flatten()  # comb biased towards identity
            f.f32(p + f"hc_{which}_base", base)

    def attn_params(p, layer_id):
        f.f32(p + "attn.attn_sink", f.randn(p + "attn.attn_sink", (H,), 0.1))
        f.fp8(p + "attn.wq_a", qr, dim)
        f.norm(p + "attn.q_norm.weight", qr)
        f.fp8(p + "attn.wq_b", H * hd, qr)
        f.fp8(p + "attn.wkv", hd, dim)
        f.norm(p + "attn.kv_norm.weight", hd)
        f.fp8(p + "attn.wo_a", og * olr, H * hd // og)
        f.fp8(p + "attn.wo_b", dim, og * olr, gain=attn_out_scale)
        backbone = layer_id is not None
        if backbone and layer_id in kv_src:
            f.bf16(p + "attn.compressor.wkv.weight", hd, dim)
            if ratios[layer_id] > 1:
                f.bf16(p + "attn.compressor.wgate.weight", hd, dim)
            f.norm(p + "attn.compressor.norm.weight", hd)
        if backbone and layer_id in idx_src:
            f.fp8(p + "attn.indexer.wq_b", ih * ihd, qr)
            f.bf16(p + "attn.indexer.weights_proj.weight", ih, dim)
            if layer_id in kv_src:  # only index-K owners carry wk / k_norm
                f.bf16(p + "attn.indexer.wk.weight", ihd, hd)
                f.norm(p + "attn.indexer.k_norm.weight", ihd)

    def ffn_params(p, n_exp):
        f.norm(p + "attn_norm.weight", dim)
        f.norm(p + "ffn_norm.weight", dim)
        f.bf16(p + "ffn.gate.weight", n_exp, dim)
        f.bias(p + "ffn.gate.bias", n_exp, std=0.02, dtype=torch.float32)
        f.bias(p + "ffn.gate.bias_vl", n_exp, std=0.02, dtype=torch.float32)
        f.fp8(p + "ffn.shared_experts.w1", inter, dim)
        f.fp8(p + "ffn.shared_experts.w3", inter, dim)
        f.fp8(p + "ffn.shared_experts.w2", dim, inter, gain=ffn_out_scale)
        for e in range(n_exp):
            q = p + f"ffn.experts.{e}."
            f.fp4(q + "w1", inter, dim)
            f.fp4(q + "w3", inter, dim)
            f.fp4(q + "w2", dim, inter, gain=ffn_out_scale)

    primes = engram_primes(t["engram_layer_ids"], t["engram_max_ngram_size"], t["engram_n_heads"], t["engram_vocab_size"])
    n_hash_cols = (t["engram_max_ngram_size"] - 1) * t["engram_n_heads"]
    ehd = t["engram_head_dim"]

    for L in range(n_layer):
        p = f"layers.{L}."
        hc_params(p)
        attn_params(p, L)
        ffn_params(p, t["n_routed_experts"])
        if L in t["engram_layer_ids"]:
            li = t["engram_layer_ids"].index(L)
            rows = sum(sum(x) for x in primes[li])
            assert rows == t["engram_num_embeddings"][li]
            w = f.randn(p + "engram.embed.weight", (rows, ehd), 1.0).bfloat16()
            q, s = quant_fp8_row(w)
            f.add(p + "engram.embed.weight", q)
            f.add(p + "engram.embed.scale", s)
            f.add(p + "engram.q_weight", (1.0 + f.randn(p + "engram.q_weight", (hc, dim), 0.1)).bfloat16())
            f.add(p + "engram.k_weight", (1.0 + f.randn(p + "engram.k_weight", (hc, dim), 0.1)).bfloat16())
            # key half is RMS-normalised inside the gate, so the gain only shrinks the value written to the stream
            f.fp8(p + "engram.wkv", dim * (hc + 1), n_hash_cols * ehd, gain=ffn_out_scale)

    for i in range(nextn):
        p = f"mtp.{i}."
        hc_params(p)
        attn_params(p, None)
        ffn_params(p, t["dspark_n_routed_experts"])
        if i == 0:
            f.fp8(p + "main_proj", dim, dim * len(t["dspark_target_layer_ids"]))
            f.norm(p + "main_norm.weight", dim)
        if i == nextn - 1:
            r = t["dspark_markov_rank"]
            f.norm(p + "norm.weight", dim)
            f.bf16(p + "markov_head.embed.weight", vocab, r, std=1.0)
            f.bf16(p + "markov_head.head.weight", vocab, r)
            f.bf16(p + "confidence_head.proj.weight", 1, dim + r)
    return f.tensors


def make_config(real: dict) -> dict:
    cfg = copy.deepcopy(real)
    t = cfg["text_config"]
    for k, val in TINY_TEXT_OVERRIDES.items():
        assert k in t, k
        t[k] = val
    for k, val in TINY_VISION_OVERRIDES.items():
        cfg["vision_config"][k] = val
    assert len(t["compress_ratios"]) == t["num_hidden_layers"] + t["num_nextn_predict_layers"]
    primes = engram_primes(t["engram_layer_ids"], t["engram_max_ngram_size"], t["engram_n_heads"], t["engram_vocab_size"])
    t["engram_num_embeddings"] = [sum(sum(x) for x in per_layer) for per_layer in primes]
    return cfg


def shard_of(name: str) -> int:
    if name.startswith(("vision.", "aligner.")):
        return 0
    if name.startswith(("embed.", "head.", "norm.", "image_")):
        return 1
    if name.startswith("layers."):
        return 2 if int(name.split(".")[1]) < 5 else 3
    return 4  # mtp.*


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--tokenizer-dir", required=True)
    ap.add_argument("--real-config", required=True, help="the released HF config.json")
    ap.add_argument("--seed", type=int, default=20261020)
    ap.add_argument("--norm-noise", type=float, default=0.0, help="std of N(1, s) norm weights (0 = all ones)")
    # Gains on the projections that write into the residual stream. With plain N(0, 1/in) everywhere (1.0, 1.0)
    # the reference is chaotic: FP8 activation quantization + top-k routing turn sub-ulp differences (e.g. another
    # CPU thread count) into 20-50% logit differences. Small residual updates mimic trained models; the MoE/engram
    # branch (discrete routing) is damped most, attention less so that the logits still depend on the context.
    # Measured (threads 16 vs 7, teacher forced): logits rel err <= 0.04 while short/medium/long contexts ending in
    # the same token differ by 0.24-0.49.
    ap.add_argument("--attn-out-scale", type=float, default=0.5, help="gain of attn.wo_b")
    ap.add_argument("--ffn-out-scale", type=float, default=1.0 / 16, help="gain of shared/routed expert w2 and engram wkv")
    args = ap.parse_args()
    torch.set_num_threads(min(16, os.cpu_count() or 1))

    real = json.load(open(args.real_config))
    cfg = make_config(real)
    os.makedirs(args.out, exist_ok=True)
    tensors = build_tensors(cfg, args.seed, args.norm_noise, args.attn_out_scale, args.ffn_out_scale)

    # self-check of the FP4 packing against convert.py's decoding convention
    for name, ten in tensors.items():
        if name.endswith("experts.0.w1.weight") and name.startswith("layers.0."):
            deq = dequant_fp4_row(ten, tensors[name.replace("weight", "scale")])
            ref = torch.randn((ten.shape[0], ten.shape[1] * 2), generator=Factory(args.seed, 0).gen(name)) * (1.0 / (ten.shape[1] * 2)) ** 0.5
            rel = (deq - ref.bfloat16().float()).norm() / ref.norm()
            print(f"fp4 self-check {name}: rel err {rel:.4f}")
            assert rel < 0.2

    n_shards = 5
    shards = [OrderedDict() for _ in range(n_shards)]
    for name, ten in tensors.items():
        shards[shard_of(name)][name] = ten
    weight_map, total = {}, 0
    for i, sd in enumerate(shards):
        fname = f"model-{i + 1:05d}-of-{n_shards:05d}.safetensors"
        save_file(sd, os.path.join(args.out, fname), metadata={"format": "pt"})
        for name, ten in sd.items():
            weight_map[name] = fname
            total += ten.numel() * ten.element_size()
    with open(os.path.join(args.out, "model.safetensors.index.json"), "w") as fh:
        json.dump({"metadata": {"total_size": total}, "weight_map": dict(sorted(weight_map.items()))}, fh, indent=2)
    with open(os.path.join(args.out, "config.json"), "w") as fh:
        json.dump(cfg, fh, indent=2)
    with open(os.path.join(args.out, "tiny_checkpoint.json"), "w") as fh:
        json.dump({"generator": "examples/neuron/make_tiny_checkpoint.py", "seed": args.seed, "norm_noise": args.norm_noise,
                   "attn_out_scale": args.attn_out_scale, "ffn_out_scale": args.ffn_out_scale, "n_tensors": len(tensors), "total_bytes": total}, fh, indent=2)
    for fn in ("tokenizer.json", "tokenizer_config.json"):
        shutil.copyfile(os.path.join(args.tokenizer_dir, fn), os.path.join(args.out, fn))
    print(f"wrote {len(tensors)} tensors, {total / 2**30:.2f} GiB to {args.out}")
    print("engram_num_embeddings", cfg["text_config"]["engram_num_embeddings"])


if __name__ == "__main__":
    main()