diff --git a/tests/test_bugfix_encode_trim_optional.py b/tests/test_bugfix_encode_trim_optional.py new file mode 100644 index 0000000..db04e82 --- /dev/null +++ b/tests/test_bugfix_encode_trim_optional.py @@ -0,0 +1,53 @@ +import importlib.util +import os +import unittest + + +class EncodeTrimBugfixOptionalTests(unittest.TestCase): + @classmethod + def setUpClass(cls): + enabled = os.getenv("VOXMLX_ENABLE_MLX_RUNTIME_TESTS", "").strip().lower() in { + "1", + "true", + "yes", + "on", + } + if not enabled: + raise unittest.SkipTest( + "Set VOXMLX_ENABLE_MLX_RUNTIME_TESTS=1 to run MLX runtime optional tests" + ) + if importlib.util.find_spec("mlx") is None: + raise unittest.SkipTest("mlx is required for runtime optional tests") + + def test_encode_trims_trailing_frames_not_leading_frames(self): + import mlx.core as mx + import numpy as np + + from voxmlx.model import VoxtralRealtime + + class _FakeEncoder: + def __call__(self, mel): + length = int(mel.shape[1]) + return mx.arange(length, dtype=mx.float32).reshape(1, length, 1) + + class _FakeAdapter: + def __call__(self, x): + return x + + class _FakeModel: + def __init__(self): + self.encoder = _FakeEncoder() + self.adapter = _FakeAdapter() + self.downsample_factor = 3 + + fake = _FakeModel() + mel = mx.zeros((128, 9), dtype=mx.float32) + out = VoxtralRealtime.encode(fake, mel) + + # Expected: keep earliest 6 encoded frames (0..5), grouped by 3. + expected = np.array([[0.0, 1.0, 2.0], [3.0, 4.0, 5.0]], dtype=np.float32) + np.testing.assert_allclose(np.array(out), expected, atol=0.0, rtol=0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/voxmlx/model.py b/voxmlx/model.py index e2fd37e..34a5db1 100644 --- a/voxmlx/model.py +++ b/voxmlx/model.py @@ -77,7 +77,7 @@ def encode(self, mel: mx.array) -> mx.array: # Truncate T to be even (for conv stride 2) T = mel.shape[1] if T % 2 != 0: - mel = mel[:, 1:] + mel = mel[:, :-1] x = self.encoder(mel) # [1, T/2, encoder_dim] x = x[0] # [T/2, encoder_dim] @@ -86,7 +86,7 @@ def encode(self, mel: mx.array) -> mx.array: L = x.shape[0] remainder = L % self.downsample_factor if remainder != 0: - x = x[remainder:] + x = x[:-remainder] L = x.shape[0] # Reshape: [T/2, 1280] -> [T/8, 5120]