File size: 3,271 Bytes
b9dc61d | 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 | from __future__ import annotations
import json
import numpy as np
import pytest
from native_token_compatibility import (
compare_native_tokens,
main,
validate_native_tokens,
)
def _valid_tokens(frames: int = 5) -> np.ndarray:
tokens = np.zeros((frames, 8), dtype=np.int64)
tokens[:, 0] = np.arange(frames) % 16_384
for index in range(1, 8):
tokens[:, index] = (np.arange(frames) + index) % 1_024
return tokens
def test_exact_cross_layout_comparison():
tokens = _valid_tokens()
result = compare_native_tokens(
tokens,
tokens.T,
reference_layout="frames_first",
recovered_layout="codebooks_first",
)
assert result["exact_match"] is True
assert result["reference_shape_frames_first"] == [5, 8]
assert result["duration_seconds_at_25hz"] == pytest.approx(0.2)
assert all(item["agreement"] == 1.0 for item in result["per_codebook"])
def test_mismatch_reports_exact_per_codebook_agreement():
reference = _valid_tokens(4)
recovered = reference.copy()
recovered[0, 3] += 1
result = compare_native_tokens(
reference,
recovered,
reference_layout="frames_first",
recovered_layout="frames_first",
)
assert result["exact_match"] is False
assert result["per_codebook"][2]["agreement"] == 1.0
assert result["per_codebook"][3]["matching_frames"] == 3
assert result["per_codebook"][3]["agreement"] == 0.75
def test_shape_mismatch_is_not_truncated_or_aligned():
result = compare_native_tokens(
_valid_tokens(4),
_valid_tokens(3),
reference_layout="frames_first",
recovered_layout="frames_first",
)
assert result["shape_match"] is False
assert result["frames"] is None
assert all(item["agreement"] is None for item in result["per_codebook"])
@pytest.mark.parametrize(
("codebook", "value", "message"),
[(0, 16_384, "c0"), (1, 1_024, "c1"), (7, -1, "c7")],
)
def test_codebook_ranges_are_enforced(codebook, value, message):
tokens = _valid_tokens()
tokens[0, codebook] = value
with pytest.raises(ValueError, match=message):
validate_native_tokens(tokens, layout="frames_first")
def test_layout_must_match_declared_axis():
with pytest.raises(ValueError, match="frames_first"):
validate_native_tokens(_valid_tokens().T, layout="frames_first")
def test_floating_point_tokens_are_rejected():
with pytest.raises(TypeError, match="integer dtype"):
validate_native_tokens(_valid_tokens().astype(np.float32), layout="frames_first")
def test_malformed_torch_token_file_json_cli_returns_structured_error(tmp_path, capsys):
malformed = tmp_path / "malformed.pt"
malformed.write_bytes(b"not a torch token file")
valid = tmp_path / "valid.npy"
np.save(valid, _valid_tokens())
result = main(
[
str(malformed),
str(valid),
"--reference-layout",
"frames_first",
"--recovered-layout",
"frames_first",
"--json",
]
)
assert result == 2
payload = json.loads(capsys.readouterr().out)
assert payload["status"] == "ERROR"
assert "unable to safely load token file" in payload["error"]
|