46 lines
1.5 KiB
Python
46 lines
1.5 KiB
Python
import glob
|
|
import os
|
|
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from sglang.srt.rust_extensions._multimodal import inkling as _rs_inkling
|
|
|
|
GOLDEN_DIR = os.environ.get(
|
|
"INKLING_MM_GOLDEN_DIR",
|
|
os.path.join(os.path.dirname(__file__), "..", "tests", "golden"),
|
|
)
|
|
GOLDENS = sorted(glob.glob(os.path.join(GOLDEN_DIR, "golden_*.npz")))
|
|
|
|
|
|
def bf16_bits_to_f32(bits: np.ndarray) -> np.ndarray:
|
|
return (bits.astype(np.uint32) << 16).view(np.float32)
|
|
|
|
|
|
@pytest.mark.parametrize("path", GOLDENS, ids=[os.path.basename(p) for p in GOLDENS])
|
|
def test_patchify_rgb_bit_exact(path):
|
|
g = np.load(path)
|
|
got = _rs_inkling.patchify_rgb(g["arr"], int(g["patch_size"]))
|
|
np.testing.assert_array_equal(got, g["bits"].reshape(-1))
|
|
|
|
|
|
@pytest.mark.parametrize("path", GOLDENS, ids=[os.path.basename(p) for p in GOLDENS])
|
|
def test_decode_patchify_png_bit_exact(path):
|
|
g = np.load(path)
|
|
h_ref, w_ref = g["arr"].shape[:2]
|
|
h, w, got = _rs_inkling.decode_patchify(g["png"].tobytes(), int(g["patch_size"]))
|
|
assert (h, w) == (h_ref, w_ref)
|
|
np.testing.assert_array_equal(got, g["bits"].reshape(-1))
|
|
|
|
|
|
def test_batch_matches_single():
|
|
gs = [np.load(p) for p in GOLDENS]
|
|
data = [g["png"].tobytes() for g in gs]
|
|
ps = int(gs[0]["patch_size"])
|
|
for (h, w, bits), g in zip(_rs_inkling.decode_patchify_batch(data, ps), gs):
|
|
np.testing.assert_array_equal(bits, g["bits"].reshape(-1))
|
|
|
|
|
|
def test_golden_fixtures_exist():
|
|
assert len(GOLDENS) >= 4, f"expected golden fixtures in {GOLDEN_DIR}"
|