| """tokenizer.pt (+config.yaml) -> w/mlx/<task>/{codec.safetensors,spec.json}.""" |
| import json, os, sys |
| import numpy as np, torch, yaml |
| from safetensors.numpy import save_file |
| sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) |
| from paths import WORK |
| from ref_torch import load_state |
| from codec_mlx import cascade_ops |
|
|
| task = sys.argv[1] |
| src, dst = f"{WORK}/up/{task}", f"{WORK}/mlx/{task}" |
| os.makedirs(dst, exist_ok=True) |
| cfg = yaml.safe_load(open(f"{src}/config.yaml")) |
| v = cfg["tokenizer"]["vqvae"]; vq = v["vector_quantizer_config"]; pt = cfg["tokenizer"]["pre_tokenizer"] |
| sd = {k.removeprefix("module.").removeprefix("_orig_mod."): t.float().numpy() for k, t in load_state(f"{src}/tokenizer.pt").items()} |
| levels = vq["levels"]["manual_levels"] |
| assert vq["levels"]["use_manual_levels"] and vq["quant_resi"]["share_mode"] == 0 and vq["quant_resi"]["enabled"] |
| assert not vq["quant_resi"]["learnable_ratio"] and v["compression_factor"] == 1 and vq["learned_l1_cascade"] and not vq["learned_all_sampling"] |
| assert v["pre_quant_groupnorm"] and v.get("pre_quant_norm_type", "group") == "group" |
| dec_layers = 1 + len({k.split(".")[2] for k in sd if k.startswith("decoder.extra_blocks.")}) |
| spec = dict(levels=levels, groupnorm=v["pre_quant_groupnorm"], n_head=v["n_layers"], phi_ratio=vq["quant_resi"]["ratio"], |
| cascade_base=vq["learned_l1_cascade_base"], cascade_norm=vq["learned_l1_cascade_norm"], |
| cascade_norm_final=vq["learned_l1_cascade_norm_final"], dec_layers=dec_layers, |
| vocab=v["in_vocab_size"], seq_len=v["context_length"], dim=v["embed_dim"], codebook=v["vqvae_vocab_size"], |
| model_id=pt["model_id"], eos_token=pt["special_tokens"]["eos_token"], pad_token=pt["special_tokens"]["pad_token"], |
| prefix_pad_token_id=cfg["model"]["prefix_pad_token_id"]) |
| W, used = {}, set() |
| def take(name, key, f=lambda a: a): |
| W[name] = np.ascontiguousarray(f(sd[key])); used.add(key) |
| take("wte", "encoder.wte.weight"); take("wpe", "encoder.wpe.weight") |
| for i in range(len(levels)): take(f"cb.{i}", f"vq.embeddings.{i}.weight") |
| for j in range(len(levels)): |
| take(f"phi.{j}.w", f"vq.quant_resi.{j}.weight", lambda a: a.transpose(0, 2, 1)); take(f"phi.{j}.b", f"vq.quant_resi.{j}.bias") |
| down, up = cascade_ops(levels[-1] // levels[0], spec["cascade_base"], spec["cascade_norm"], spec["cascade_norm_final"]) |
| for op in down: |
| if op[0] == "d": |
| s, i = op[1], op[2] |
| take(f"down.{i}.w", f"vq.learned_down_l1.{i}.weight", lambda a: a.transpose(2, 1, 0).reshape(-1, a.shape[0])); take(f"down.{i}.b", f"vq.learned_down_l1.{i}.bias") |
| for op in up: |
| if op[0] == "u": |
| s, i = op[1], op[2] |
| take(f"up.{i}.w", f"vq.learned_up_l1.{i}.weight", lambda a: a.transpose(0, 2, 1).reshape(a.shape[0], -1)); take(f"up.{i}.b", f"vq.learned_up_l1.{i}.bias") |
| for n in range(dec_layers): |
| u = "decoder.middle_block" if n == 0 else f"decoder.extra_blocks.{n-1}" |
| for a, b in [("ln1.w", "ln_1.weight"), ("ln1.b", "ln_1.bias"), ("attn.w", "attn.c_attn.weight"), ("attn.b", "attn.c_attn.bias"), |
| ("qn", "attn.q_norm.weight"), ("kn", "attn.k_norm.weight"), ("proj.w", "attn.c_proj.weight"), ("proj.b", "attn.c_proj.bias"), |
| ("ln2.w", "ln_2.weight"), ("ln2.b", "ln_2.bias"), ("fc.w", "mlp.c_fc.weight"), ("fc.b", "mlp.c_fc.bias"), |
| ("fcp.w", "mlp.c_proj.weight"), ("fcp.b", "mlp.c_proj.bias")]: |
| take(f"blk{n}.{a}", f"{u}.{b}") |
| take("lnf.w", "decoder.ln_f.weight"); W["lnf.b"] = np.zeros_like(sd["decoder.ln_f.weight"]) |
| take("head", "decoder.output_layer.weight") |
| left = [k for k in sd if k not in used and not k.startswith(("vq.ema_", "vq.codebook_hit", "vq._"))] |
| print("unused keys:", left) |
| save_file(W, f"{dst}/codec.safetensors"); json.dump(spec, open(f"{dst}/spec.json", "w"), indent=1) |
| print(task, "ok", len(W), "tensors", sum(a.size for a in W.values()) / 1e6, "M params") |
|
|
