File size: 8,887 Bytes
90884df | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """Validate and exactly compare native MiniMax Music 3 token trajectories."""
from __future__ import annotations
import argparse
import json
import pickle
from pathlib import Path
from typing import Any, Literal
import numpy as np
import torch
TokenLayout = Literal["frames_first", "codebooks_first"]
LAYOUTS: tuple[TokenLayout, ...] = ("frames_first", "codebooks_first")
NUM_CODEBOOKS = 8
CODEBOOK_SIZES = (16_384, 1_024, 1_024, 1_024, 1_024, 1_024, 1_024, 1_024)
FRAME_RATE_HZ = 25
def normalize_tokens(tokens: Any, *, layout: TokenLayout) -> np.ndarray:
"""Validate layout explicitly and return a CPU NumPy array ``[frames, 8]``."""
if layout not in LAYOUTS:
raise ValueError(f"layout must be one of {LAYOUTS}, got {layout!r}")
if isinstance(tokens, torch.Tensor):
if tokens.device.type != "cpu":
raise ValueError("token comparison is CPU-only; pass a CPU tensor")
array = tokens.detach().numpy()
else:
array = np.asarray(tokens)
if array.ndim != 2:
raise ValueError(f"tokens must be rank 2, got shape {array.shape}")
expected_axis = 1 if layout == "frames_first" else 0
if array.shape[expected_axis] != NUM_CODEBOOKS:
expected = "[frames, 8]" if layout == "frames_first" else "[8, frames]"
raise ValueError(f"{layout} tokens must have shape {expected}, got {array.shape}")
if not np.issubdtype(array.dtype, np.integer) or np.issubdtype(array.dtype, np.bool_):
raise TypeError(f"tokens must use an integer dtype, got {array.dtype}")
frames_first = array if layout == "frames_first" else array.T
if frames_first.shape[0] == 0:
raise ValueError("token trajectory must contain at least one frame")
return np.ascontiguousarray(frames_first, dtype=np.int64)
def validate_native_tokens(tokens: Any, *, layout: TokenLayout) -> np.ndarray:
"""Return normalized tokens after enforcing every native codebook range."""
normalized = normalize_tokens(tokens, layout=layout)
for index, vocab_size in enumerate(CODEBOOK_SIZES):
values = normalized[:, index]
minimum = int(values.min())
maximum = int(values.max())
if minimum < 0 or maximum >= vocab_size:
raise ValueError(
f"codebook c{index} values must be in [0, {vocab_size - 1}], "
f"observed min={minimum}, max={maximum}"
)
return normalized
def compare_native_tokens(
reference: Any,
recovered: Any,
*,
reference_layout: TokenLayout,
recovered_layout: TokenLayout,
) -> dict[str, Any]:
"""Compare trajectories without truncation, alignment, or tolerance."""
reference_tokens = validate_native_tokens(reference, layout=reference_layout)
recovered_tokens = validate_native_tokens(recovered, layout=recovered_layout)
shape_match = reference_tokens.shape == recovered_tokens.shape
frames = int(reference_tokens.shape[0]) if shape_match else None
per_codebook: list[dict[str, Any]] = []
for index, vocab_size in enumerate(CODEBOOK_SIZES):
entry: dict[str, Any] = {"codebook": index, "vocab_size": vocab_size}
if shape_match:
matches = int(np.count_nonzero(reference_tokens[:, index] == recovered_tokens[:, index]))
entry.update(
{
"matching_frames": matches,
"total_frames": frames,
"agreement": matches / frames,
"agreement_percent": 100.0 * matches / frames,
"exact": matches == frames,
}
)
else:
entry.update(
{
"matching_frames": None,
"total_frames": None,
"agreement": None,
"agreement_percent": None,
"exact": False,
}
)
per_codebook.append(entry)
exact_match = shape_match and bool(np.array_equal(reference_tokens, recovered_tokens))
return {
"status": "MATCH" if exact_match else "MISMATCH",
"exact_match": exact_match,
"shape_match": shape_match,
"reference_shape_frames_first": list(reference_tokens.shape),
"recovered_shape_frames_first": list(recovered_tokens.shape),
"frames": frames,
"duration_seconds_at_25hz": frames / FRAME_RATE_HZ if frames is not None else None,
"frame_rate_hz": FRAME_RATE_HZ,
"codebook_sizes": list(CODEBOOK_SIZES),
"per_codebook": per_codebook,
}
def _select_loaded(value: Any, *, key: str | None, source: Path) -> Any:
if isinstance(value, np.lib.npyio.NpzFile):
available = list(value.files)
selected = key or (available[0] if len(available) == 1 else None)
if selected is None or selected not in available:
raise ValueError(f"{source} contains arrays {available}; select one with --*-key")
return value[selected]
if isinstance(value, dict):
available = list(value)
selected = key or (available[0] if len(available) == 1 else None)
if selected is None or selected not in value:
raise ValueError(f"{source} contains keys {available}; select one with --*-key")
return value[selected]
if key is not None:
raise ValueError(f"{source} is a bare array/tensor, so a key cannot be selected")
return value
def load_token_file(path: str | Path, *, key: str | None = None) -> Any:
source = Path(path)
suffix = source.suffix.lower()
if suffix in {".pt", ".pth"}:
try:
value = torch.load(str(source), map_location="cpu", weights_only=True)
except (pickle.UnpicklingError, EOFError, OSError, RuntimeError) as error:
raise ValueError(f"unable to safely load token file {source}: {error}") from error
return _select_loaded(value, key=key, source=source)
if suffix == ".json":
with source.open("r", encoding="utf-8") as handle:
value = json.load(handle)
return _select_loaded(value, key=key, source=source)
if suffix in {".npy", ".npz"}:
value = np.load(source, allow_pickle=False)
try:
selected = _select_loaded(value, key=key, source=source)
return np.asarray(selected).copy() if isinstance(value, np.lib.npyio.NpzFile) else selected
finally:
if isinstance(value, np.lib.npyio.NpzFile):
value.close()
raise ValueError(f"unsupported token file extension {suffix!r}; use .npy, .npz, .pt, .pth, or .json")
def format_human(result: dict[str, Any]) -> str:
lines = [
f"Status: {result['status']}",
f"Reference shape [frames, codebooks]: {result['reference_shape_frames_first']}",
f"Recovered shape [frames, codebooks]: {result['recovered_shape_frames_first']}",
f"Frame rate: {result['frame_rate_hz']} Hz",
]
if result["duration_seconds_at_25hz"] is not None:
lines.append(f"Duration represented: {result['duration_seconds_at_25hz']:.3f} s")
for entry in result["per_codebook"]:
agreement = entry["agreement_percent"]
rendered = "not comparable (shape mismatch)" if agreement is None else f"{agreement:.6f}%"
lines.append(
f"CB{entry['codebook']} [0, {entry['vocab_size'] - 1}]: {rendered}"
)
return "\n".join(lines)
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("reference", help="Internal Music 3 token trajectory")
parser.add_argument("recovered", help="Tokens recovered by the candidate encoder")
parser.add_argument("--reference-layout", required=True, choices=LAYOUTS)
parser.add_argument("--recovered-layout", required=True, choices=LAYOUTS)
parser.add_argument("--reference-key")
parser.add_argument("--recovered-key")
parser.add_argument("--json", action="store_true")
return parser
def main(argv: list[str] | None = None) -> int:
args = build_parser().parse_args(argv)
try:
reference = load_token_file(args.reference, key=args.reference_key)
recovered = load_token_file(args.recovered, key=args.recovered_key)
result = compare_native_tokens(
reference,
recovered,
reference_layout=args.reference_layout,
recovered_layout=args.recovered_layout,
)
except (FileNotFoundError, OSError, RuntimeError, TypeError, ValueError) as error:
payload = {"status": "ERROR", "error": str(error)}
print(json.dumps(payload, indent=2) if args.json else f"ERROR: {error}")
return 2
print(json.dumps(result, indent=2) if args.json else format_human(result))
return 0 if result["exact_match"] else 1
if __name__ == "__main__":
raise SystemExit(main())
|