| """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()) |
|
|
