From 23a08974f09b0835d9e1879fca5b0612d153c61c Mon Sep 17 00:00:00 2001 From: BP <11394934+benjipeng@users.noreply.github.com> Date: Fri, 14 Aug 2026 18:10:18 -0400 Subject: [PATCH 1/2] Harden test logistics --- .github/workflows/ci.yml | 14 ++- .gitignore | 1 + README.md | 6 +- pyproject.toml | 12 ++ tests/README.md | 108 ++++++++++++++++++ tests/checkpoint/__init__.py | 1 + .../test_cohere_asr_tokenizer_checkpoint.py | 34 ++++++ .../checkpoint/test_moss_delay_checkpoint.py | 26 +++++ .../test_moss_sound_effect_checkpoint.py} | 17 ++- .../test_vibevoice_checkpoint_local.py | 35 ++++++ .../test_vibevoice_config_checkpoint.py | 40 +++++++ tests/conftest.py | 20 ---- tests/integration/__init__.py | 1 + .../test_moss_local.py} | 81 ++++++++++--- tests/integration/test_vibevoice.py | 73 ++++++++++++ tests/runtime/__init__.py | 1 + .../test_moss_delay_generation_runtime.py} | 17 ++- .../test_moss_local_processor.py} | 57 ++++++--- .../test_vibevoice_acoustic_runtime.py | 44 +++++++ tests/test_cohere_asr_generation.py | 62 ---------- .../test_moss_delay_checkpoint_integration.py | 25 ---- tests/test_vibevoice_config.py | 89 --------------- tests/unit/__init__.py | 1 + tests/{ => unit}/test_alignment_report.py | 0 tests/{ => unit}/test_audio_io.py | 0 ..._batch_generate_longcat_audiodit_script.py | 8 +- .../test_benchmark_moss_ttsd_script.py | 7 +- tests/{ => unit}/test_checkpoint_layout.py | 0 .../test_cohere_asr_feature_extraction.py | 2 +- tests/unit/test_cohere_asr_generation.py | 48 ++++++++ .../test_convert_longcat_audiodit_script.py | 2 +- .../unit/test_dots_tts_checkpoint_contract.py | 2 +- tests/unit/test_dots_tts_convert.py | 2 +- tests/unit/test_dots_tts_quantization.py | 2 +- tests/unit/test_dots_tts_vocoder.py | 2 +- tests/unit/test_dots_tts_vocoder_streaming.py | 2 +- .../test_generate_longcat_audiodit_script.py | 2 +- .../test_generate_moss_sound_effect_script.py | 7 +- .../test_generate_moss_ttsd_script.py | 7 +- .../test_generate_step_audio_editx_script.py | 2 +- .../test_generate_vibevoice_script.py | 36 +++--- tests/{ => unit}/test_moss_audio_tokenizer.py | 0 tests/{ => unit}/test_moss_delay_config.py | 20 +++- tests/{ => unit}/test_moss_delay_dialogue.py | 0 .../{ => unit}/test_moss_delay_generation.py | 0 tests/{ => unit}/test_moss_delay_model.py | 0 tests/{ => unit}/test_moss_delay_processor.py | 0 .../test_moss_delay_quantized_checkpoint.py | 0 tests/{ => unit}/test_moss_local_cache.py | 0 tests/{ => unit}/test_moss_local_config.py | 0 .../{ => unit}/test_moss_local_generation.py | 7 +- tests/{ => unit}/test_moss_local_model.py | 0 .../test_moss_local_quantized_checkpoint.py | 0 tests/{ => unit}/test_moss_sound_effect.py | 0 tests/{ => unit}/test_runtime_purity.py | 4 +- tests/{ => unit}/test_sharded_checkpoint.py | 0 tests/unit/test_tts_streaming_protocol.py | 2 +- tests/{ => unit}/test_vibevoice_acoustic.py | 48 ++------ tests/{ => unit}/test_vibevoice_checkpoint.py | 37 ++---- tests/unit/test_vibevoice_config.py | 74 ++++++++++++ tests/{ => unit}/test_vibevoice_diffusion.py | 0 tests/{ => unit}/test_vibevoice_generation.py | 59 ---------- uv.lock | 100 ++++++++++++++++ 63 files changed, 849 insertions(+), 398 deletions(-) create mode 100644 tests/README.md create mode 100644 tests/checkpoint/__init__.py create mode 100644 tests/checkpoint/test_cohere_asr_tokenizer_checkpoint.py create mode 100644 tests/checkpoint/test_moss_delay_checkpoint.py rename tests/{test_moss_sound_effect_integration.py => checkpoint/test_moss_sound_effect_checkpoint.py} (51%) create mode 100644 tests/checkpoint/test_vibevoice_checkpoint_local.py create mode 100644 tests/checkpoint/test_vibevoice_config_checkpoint.py create mode 100644 tests/integration/__init__.py rename tests/{test_moss_local_integration.py => integration/test_moss_local.py} (72%) create mode 100644 tests/integration/test_vibevoice.py create mode 100644 tests/runtime/__init__.py rename tests/{test_moss_delay_generation_integration.py => runtime/test_moss_delay_generation_runtime.py} (85%) rename tests/{test_moss_local_processor_integration.py => runtime/test_moss_local_processor.py} (86%) create mode 100644 tests/runtime/test_vibevoice_acoustic_runtime.py delete mode 100644 tests/test_cohere_asr_generation.py delete mode 100644 tests/test_moss_delay_checkpoint_integration.py delete mode 100644 tests/test_vibevoice_config.py create mode 100644 tests/unit/__init__.py rename tests/{ => unit}/test_alignment_report.py (100%) rename tests/{ => unit}/test_audio_io.py (100%) rename tests/{ => unit}/test_batch_generate_longcat_audiodit_script.py (90%) rename tests/{ => unit}/test_benchmark_moss_ttsd_script.py (95%) rename tests/{ => unit}/test_checkpoint_layout.py (100%) rename tests/{ => unit}/test_cohere_asr_feature_extraction.py (98%) create mode 100644 tests/unit/test_cohere_asr_generation.py rename tests/{ => unit}/test_convert_longcat_audiodit_script.py (85%) rename tests/{ => unit}/test_generate_longcat_audiodit_script.py (83%) rename tests/{ => unit}/test_generate_moss_sound_effect_script.py (95%) rename tests/{ => unit}/test_generate_moss_ttsd_script.py (97%) rename tests/{ => unit}/test_generate_step_audio_editx_script.py (98%) rename tests/{ => unit}/test_generate_vibevoice_script.py (74%) rename tests/{ => unit}/test_moss_audio_tokenizer.py (100%) rename tests/{ => unit}/test_moss_delay_config.py (72%) rename tests/{ => unit}/test_moss_delay_dialogue.py (100%) rename tests/{ => unit}/test_moss_delay_generation.py (100%) rename tests/{ => unit}/test_moss_delay_model.py (100%) rename tests/{ => unit}/test_moss_delay_processor.py (100%) rename tests/{ => unit}/test_moss_delay_quantized_checkpoint.py (100%) rename tests/{ => unit}/test_moss_local_cache.py (100%) rename tests/{ => unit}/test_moss_local_config.py (100%) rename tests/{ => unit}/test_moss_local_generation.py (98%) rename tests/{ => unit}/test_moss_local_model.py (100%) rename tests/{ => unit}/test_moss_local_quantized_checkpoint.py (100%) rename tests/{ => unit}/test_moss_sound_effect.py (100%) rename tests/{ => unit}/test_runtime_purity.py (91%) rename tests/{ => unit}/test_sharded_checkpoint.py (100%) rename tests/{ => unit}/test_vibevoice_acoustic.py (73%) rename tests/{ => unit}/test_vibevoice_checkpoint.py (61%) create mode 100644 tests/unit/test_vibevoice_config.py rename tests/{ => unit}/test_vibevoice_diffusion.py (100%) rename tests/{ => unit}/test_vibevoice_generation.py (57%) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index f7cd309..c545ecd 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -23,12 +23,16 @@ jobs: - name: Install dependencies run: uv sync --all-extras - - name: Run unit tests - run: uv run pytest tests/unit/ -q + - name: Run fast tests + run: >- + uv run pytest -q -ra + --cov=mlx_speech + --cov-report=term-missing:skip-covered env: # GitHub's virtualized macOS runner reports Metal but aborts custom kernels. MLX_SPEECH_DISABLE_CUSTOM_METAL: "1" - - name: Run checkpoint tests - # Skips cleanly when local checkpoints are absent, which is always on CI. - run: uv run pytest tests/checkpoint/ -q + - name: Validate opt-in test collection + run: >- + uv run pytest --collect-only -q + tests/checkpoint/ tests/runtime/ tests/integration/ diff --git a/.gitignore b/.gitignore index 754f800..8491aeb 100644 --- a/.gitignore +++ b/.gitignore @@ -14,6 +14,7 @@ .hypothesis/ .coverage .coverage.* +coverage.xml htmlcov/ __pycache__/ *.py[cod] diff --git a/README.md b/README.md index 1222ec7..6dacbe4 100644 --- a/README.md +++ b/README.md @@ -195,10 +195,14 @@ artifact without PyTorch or `mlx-audio`. git clone https://github.com/appautomaton/mlx-speech.git cd mlx-speech uv sync -uv run pytest tests/unit/ +uv run pytest uv run ruff check . ``` +The default command runs the fast, artifact-free tier. See +[`tests/README.md`](tests/README.md) for checkpoint, runtime, integration, +fixture, and coverage gates. + ```text mlx-speech/ src/mlx_speech/ library code diff --git a/pyproject.toml b/pyproject.toml index 7dc59a2..496380d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,6 +22,7 @@ dependencies = [ [dependency-groups] dev = [ "pytest>=8.3,<9", + "pytest-cov>=7.1,<8", "ruff>=0.11,<0.12", ] @@ -30,6 +31,7 @@ requires = ["uv_build>=0.11.2,<0.12"] build-backend = "uv_build" [tool.pytest.ini_options] +addopts = ["--strict-config", "--strict-markers"] testpaths = ["tests"] markers = [ "checkpoint: needs local model checkpoints (skip if absent)", @@ -38,6 +40,16 @@ markers = [ "local_integration: test requires local model artifacts or repo-specific runtime assets", ] +[tool.coverage.run] +branch = true +source = ["mlx_speech"] + +[tool.coverage.report] +fail_under = 72.8 +precision = 1 +show_missing = true +skip_empty = true + [tool.ruff] target-version = "py313" exclude = [".references", ".venv"] diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..56d7cf6 --- /dev/null +++ b/tests/README.md @@ -0,0 +1,108 @@ +# Testing + +Tests protect existing runtime behavior. Reorganizing tests, fixtures, CI, or +coverage does not authorize changes under `src/`. If test cleanup exposes a +production defect, report it and fix it in a separate behavior change. + +## Tiers + +| Tier | Purpose | Local artifacts | Default CI | +| --- | --- | --- | --- | +| `unit/` | Pure logic, tiny MLX models, synthetic checkpoints, and bounded oracle fixtures | No | Yes | +| `checkpoint/` | Loading and alignment against real local checkpoint assets | Yes | Collection only | +| `runtime/` | Real-weight forward, inference, streaming, and model-level parity | Yes | Collection only | +| `integration/` | Public API through waveform or transcript output | Yes, plus evaluation inputs | Collection only | + +The default command runs only the fast tier: + +```bash +pytest +``` + +Run an opt-in tier by naming its directory: + +```bash +pytest tests/checkpoint/ +pytest tests/runtime/ +RUN_LOCAL_INTEGRATION=1 pytest tests/integration/ +``` + +When an opt-in run is acting as a required gate, make skips fail the session: + +```bash +MLX_SPEECH_REQUIRE_CHECKPOINTS=1 pytest tests/runtime/ +``` + +Use this inexpensive command to validate imports and collection across every +tier without running model inference: + +```bash +pytest --collect-only -q tests/unit/ tests/checkpoint/ tests/runtime/ tests/integration/ +``` + +Do not put `test_*.py` files directly under `tests/`. Tier directories are +Python packages so identically named tests in different tiers cannot collide +during combined collection. + +## Placement rules + +A test belongs in `unit/` only when it is deterministic and needs no network, +local model directory, upstream checkout, or optional Torch environment. Prefer +real tiny MLX modules and tiny safetensors over mocks. + +Checkpoint tests validate real artifact layout, keys, shapes, storage/runtime +dtypes, quantization metadata, and strict alignment. Pure remapping and loader +branches still belong in `unit/` and should use synthetic checkpoint files. + +Runtime tests load real weights and exercise component or inference behavior. +Integration tests cross the public API boundary and must reach waveform or +transcript output. A finite, non-empty waveform assertion is a smoke test, not +a quality gate; WER, CER, speaker similarity, numeric parity, memory, and timing +regressions need their own explicit metrics. + +Every bug fix adds the smallest test that fails before the fix. Put that test in +the lowest tier capable of reproducing the bug, then add a higher-tier contract +test only when the failure can cross a component boundary. + +## Test doubles + +Use stubs or fakes at expensive and nondeterministic boundaries: Hub access, +network calls, subprocesses, clocks, file writers, tokenizers, and heavyweight +model adapters. A test double should reject unexpected calls and preserve the +boundary's input/output contract. + +Do not replace the behavior under test. Loader tests use real tiny files; +generation-state tests use a tiny model or a narrow adapter fake; numeric model +tests run the real layer implementation. + +## Golden fixtures + +Golden fixtures isolate the MLX test from the upstream reference environment. +They avoid a full checkpoint only when the fixture contains the required small +weights or the tested component is weight-free. + +Committed oracle fixtures must record: + +- deterministic input construction and seed; +- reference repository revision and dependency versions; +- array shapes, dtypes, hashes, and bounded file sizes; +- a tolerance chosen for the numeric quantity being compared; +- capture and regeneration commands. + +Use exact equality for discrete schedules and tokens, `allclose` for stable +tensor math, and scale-aware metrics such as correlation or relative RMSE for +waveforms and complex spectra. Do not regenerate a fixture merely because a +test failed. Review the behavioral difference and the pinned reference first. + +## Coverage + +CI measures line and branch coverage for `mlx_speech`. The floor in +`pyproject.toml` is the measured fast-suite baseline, not a claim that every +model family has sufficient behavioral coverage. It may move upward with new +tests and must not be lowered to make CI pass. + +New or changed production behavior needs direct coverage for success, boundary, +and failure paths. Model-family reviews should trace config parsing, checkpoint +mapping, component parity, generation state, real runtime inference, public API +output, and quality/performance gates instead of relying on a single aggregate +percentage. diff --git a/tests/checkpoint/__init__.py b/tests/checkpoint/__init__.py new file mode 100644 index 0000000..3f18d29 --- /dev/null +++ b/tests/checkpoint/__init__.py @@ -0,0 +1 @@ +"""Tests that require local checkpoint artifacts.""" diff --git a/tests/checkpoint/test_cohere_asr_tokenizer_checkpoint.py b/tests/checkpoint/test_cohere_asr_tokenizer_checkpoint.py new file mode 100644 index 0000000..bf406b3 --- /dev/null +++ b/tests/checkpoint/test_cohere_asr_tokenizer_checkpoint.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from mlx_speech.models.cohere_asr.tokenizer import CohereAsrTokenizer + + +MODEL_DIR = Path("models/cohere/cohere_transcribe/original") + +pytestmark = [ + pytest.mark.checkpoint, + pytest.mark.skipif( + not MODEL_DIR.is_dir(), + reason="Cohere ASR tokenizer assets are not present", + ), +] + + +def test_tokenizer_prompt_ids_support_punctuation_and_itn() -> None: + tokenizer = CohereAsrTokenizer.from_dir(MODEL_DIR) + + default_prompt = tokenizer.get_decoder_prompt_ids("en") + no_punctuation_prompt = tokenizer.get_decoder_prompt_ids("en", punctuation=False) + itn_prompt = tokenizer.get_decoder_prompt_ids("en", itn=True) + + assert len(default_prompt) == 10 + assert default_prompt[6] == 5 # <|pnc|> + assert no_punctuation_prompt[6] == 6 # <|nopnc|> + assert default_prompt[7] == 9 # <|noitn|> + assert itn_prompt[7] == 8 # <|itn|> + assert default_prompt[8:] == [11, 13] # <|notimestamp|>, <|nodiarize|> + assert itn_prompt[:7] == default_prompt[:7] diff --git a/tests/checkpoint/test_moss_delay_checkpoint.py b/tests/checkpoint/test_moss_delay_checkpoint.py new file mode 100644 index 0000000..08ce6d2 --- /dev/null +++ b/tests/checkpoint/test_moss_delay_checkpoint.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from mlx_speech.models.moss_delay import load_moss_tts_delay_model + +MODEL_DIR = Path("models/openmoss/moss_ttsd/mlx-int8") + +pytestmark = [ + pytest.mark.checkpoint, + pytest.mark.skipif( + not MODEL_DIR.is_dir(), + reason="MOSS-TTSD checkpoint is not present", + ), +] + + +def test_default_ttsd_runtime_loads_quantized_mlx_model() -> None: + loaded = load_moss_tts_delay_model() + + assert loaded.alignment_report.is_exact_match + assert loaded.model.config.n_vq == 16 + assert loaded.model.language_model.config.num_hidden_layers == 36 + assert loaded.quantization is not None diff --git a/tests/test_moss_sound_effect_integration.py b/tests/checkpoint/test_moss_sound_effect_checkpoint.py similarity index 51% rename from tests/test_moss_sound_effect_integration.py rename to tests/checkpoint/test_moss_sound_effect_checkpoint.py index 0796a88..e11f20b 100644 --- a/tests/test_moss_sound_effect_integration.py +++ b/tests/checkpoint/test_moss_sound_effect_checkpoint.py @@ -1,10 +1,23 @@ from __future__ import annotations +from pathlib import Path + import pytest -from mlx_speech.models.moss_delay import load_moss_sound_effect_model, resolve_moss_sound_effect_model_dir +from mlx_speech.models.moss_delay import ( + load_moss_sound_effect_model, + resolve_moss_sound_effect_model_dir, +) + +MODEL_DIR = Path("models/openmoss/moss_sound_effect/mlx-4bit") -pytestmark = pytest.mark.local_integration +pytestmark = [ + pytest.mark.checkpoint, + pytest.mark.skipif( + not MODEL_DIR.is_dir(), + reason="MOSS sound-effect checkpoint is not present", + ), +] def test_default_moss_sound_effect_runtime_loads_quantized_mlx_model() -> None: diff --git a/tests/checkpoint/test_vibevoice_checkpoint_local.py b/tests/checkpoint/test_vibevoice_checkpoint_local.py new file mode 100644 index 0000000..a4e3389 --- /dev/null +++ b/tests/checkpoint/test_vibevoice_checkpoint_local.py @@ -0,0 +1,35 @@ +"""Local VibeVoice checkpoint loading and alignment coverage.""" + +from pathlib import Path + +import pytest + +from mlx_speech.models.vibevoice.checkpoint import ( + load_vibevoice_checkpoint, + load_vibevoice_model, +) + + +MODEL_DIR = Path("models/vibevoice/mlx-int8") +HAS_CHECKPOINT = MODEL_DIR.is_dir() and any(MODEL_DIR.glob("*.safetensors")) + +pytestmark = [ + pytest.mark.checkpoint, + pytest.mark.skipif( + not HAS_CHECKPOINT, + reason="VibeVoice checkpoint is not present", + ), +] + + +def test_load_checkpoint() -> None: + checkpoint = load_vibevoice_checkpoint(MODEL_DIR) + + assert checkpoint.key_count > 0 + assert checkpoint.config.model_type == "vibevoice" + + +def test_model_alignment() -> None: + loaded = load_vibevoice_model(MODEL_DIR, strict=False) + + assert loaded.alignment_report.is_exact_match diff --git a/tests/checkpoint/test_vibevoice_config_checkpoint.py b/tests/checkpoint/test_vibevoice_config_checkpoint.py new file mode 100644 index 0000000..f5c997d --- /dev/null +++ b/tests/checkpoint/test_vibevoice_config_checkpoint.py @@ -0,0 +1,40 @@ +"""VibeVoice configuration coverage against local upstream assets.""" + +from pathlib import Path + +import pytest + +from mlx_speech.models.vibevoice.config import VibeVoiceConfig + + +MODEL_DIR = Path("models/vibevoice/original") + +pytestmark = [ + pytest.mark.checkpoint, + pytest.mark.skipif( + not (MODEL_DIR / "config.json").is_file(), + reason="VibeVoice original config is not present", + ), +] + + +def test_from_path() -> None: + config = VibeVoiceConfig.from_path(MODEL_DIR) + + assert config.model_type == "vibevoice" + assert config.hidden_size == 3584 + assert config.language_config.num_hidden_layers == 28 + assert config.acoustic_tokenizer_config.vae_dim == 64 + assert config.semantic_tokenizer_config.vae_dim == 128 + assert config.diffusion_config.head_layers == 4 + + +def test_round_trip() -> None: + config = VibeVoiceConfig.from_path(MODEL_DIR) + restored = VibeVoiceConfig.from_dict(config.to_dict()) + + assert restored.hidden_size == config.hidden_size + assert ( + restored.acoustic_tokenizer_config.vae_dim + == config.acoustic_tokenizer_config.vae_dim + ) diff --git a/tests/conftest.py b/tests/conftest.py index da91139..2797a10 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -36,18 +36,6 @@ def _requested_tiers(root: Path, args: tuple[str, ...]) -> set[str]: return enabled -def _is_explicit_target(path: Path, root: Path, args: tuple[str, ...]) -> bool: - for arg in args: - if not arg or arg.startswith("-"): - continue - candidate = Path(arg) - if not candidate.is_absolute(): - candidate = (root / candidate).resolve() - if candidate == path: - return True - return False - - def pytest_ignore_collect(collection_path: Path, config) -> bool: # type: ignore[no-untyped-def] root = Path(str(config.rootpath)).resolve() args = tuple(config.invocation_params.args) @@ -57,14 +45,6 @@ def pytest_ignore_collect(collection_path: Path, config) -> bool: # type: ignor tier_dir = (root / "tests" / tier).resolve() if path == tier_dir or tier_dir in path.parents: return tier not in enabled - if path.parent == (root / "tests").resolve() and path.suffix == ".py": - text = path.read_text(encoding="utf-8") - if ( - 'Path("models/' in text - or 'MODEL_DIR = "models/' in text - or "pytest.mark.local_integration" in text - ): - return not _is_explicit_target(path, root, args) return False diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000..8e76d4a --- /dev/null +++ b/tests/integration/__init__.py @@ -0,0 +1 @@ +"""End-to-end tests that exercise public speech APIs.""" diff --git a/tests/test_moss_local_integration.py b/tests/integration/test_moss_local.py similarity index 72% rename from tests/test_moss_local_integration.py rename to tests/integration/test_moss_local.py index 20f493a..bd1e413 100644 --- a/tests/test_moss_local_integration.py +++ b/tests/integration/test_moss_local.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from pathlib import Path import mlx.core as mx @@ -11,9 +12,25 @@ synthesize_moss_tts_local_conversations, ) from mlx_speech.models.moss_audio_tokenizer import load_moss_audio_tokenizer_model -from mlx_speech.models.moss_local import MossTTSLocalProcessor, load_moss_tts_local_model +from mlx_speech.models.moss_local import ( + MossTTSLocalProcessor, + load_moss_tts_local_model, +) + +MODEL_DIR = Path("models/openmoss/moss_tts_local/mlx-int8") +CODEC_DIR = Path("models/openmoss/moss_audio_tokenizer/mlx-int8") -pytestmark = pytest.mark.local_integration +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + os.environ.get("RUN_LOCAL_INTEGRATION") != "1", + reason="set RUN_LOCAL_INTEGRATION=1 for local waveform generation", + ), + pytest.mark.skipif( + not MODEL_DIR.is_dir() or not CODEC_DIR.is_dir(), + reason="MOSS-TTS-Local checkpoints are not present", + ), +] def _runtime(): @@ -65,14 +82,18 @@ def test_default_quantized_runtime_supports_main_inference_modes( conversations = [ [ processor.build_user_message(**user_kwargs), - processor.build_assistant_message(audio_codes_list=[str(reference_audio_path)]), + processor.build_assistant_message( + audio_codes_list=[str(reference_audio_path)] + ), ] ] else: conversations = [ [ processor.build_user_message(**user_kwargs), - processor.build_assistant_message(audio_codes_list=[str(reference_audio_path)]), + processor.build_assistant_message( + audio_codes_list=[str(reference_audio_path)] + ), ] ] @@ -96,7 +117,11 @@ def test_batch_inference_preserves_output_order(reference_audio_path: Path) -> N config = MossTTSLocalGenerationConfig(max_new_tokens=2, do_sample=False) conversations = [ [processor.build_user_message(text="First sample.", tokens=6)], - [processor.build_user_message(text="Second sample.", reference=[str(reference_audio_path)], tokens=6)], + [ + processor.build_user_message( + text="Second sample.", reference=[str(reference_audio_path)], tokens=6 + ) + ], ] result = synthesize_moss_tts_local_conversations( @@ -141,14 +166,18 @@ def test_cached_and_uncached_single_sample_paths_match( conversations = [ [ processor.build_user_message(**user_kwargs), - processor.build_assistant_message(audio_codes_list=[str(reference_audio_path)]), + processor.build_assistant_message( + audio_codes_list=[str(reference_audio_path)] + ), ] ] else: conversations = [ [ processor.build_user_message(**user_kwargs), - processor.build_assistant_message(audio_codes_list=[str(reference_audio_path)]), + processor.build_assistant_message( + audio_codes_list=[str(reference_audio_path)] + ), ] ] @@ -158,7 +187,9 @@ def test_cached_and_uncached_single_sample_paths_match( codec, conversations=conversations, mode=processor_mode, - config=MossTTSLocalGenerationConfig(max_new_tokens=2, do_sample=False, use_kv_cache=True), + config=MossTTSLocalGenerationConfig( + max_new_tokens=2, do_sample=False, use_kv_cache=True + ), ) uncached = synthesize_moss_tts_local_conversations( model, @@ -173,9 +204,17 @@ def test_cached_and_uncached_single_sample_paths_match( ), ) - assert cached.generation.sequences.tolist() == uncached.generation.sequences.tolist() - assert cached.generation.generated_rows.tolist() == uncached.generation.generated_rows.tolist() - assert cached.generation.audio_codes_list[0].tolist() == uncached.generation.audio_codes_list[0].tolist() + assert ( + cached.generation.sequences.tolist() == uncached.generation.sequences.tolist() + ) + assert ( + cached.generation.generated_rows.tolist() + == uncached.generation.generated_rows.tolist() + ) + assert ( + cached.generation.audio_codes_list[0].tolist() + == uncached.generation.audio_codes_list[0].tolist() + ) def test_batch_kv_cache_flag_falls_back_to_uncached_path( @@ -184,7 +223,13 @@ def test_batch_kv_cache_flag_falls_back_to_uncached_path( model, processor, codec = _runtime() conversations = [ [processor.build_user_message(text="First cached batch item.", tokens=6)], - [processor.build_user_message(text="Second cached batch item.", reference=[str(reference_audio_path)], tokens=6)], + [ + processor.build_user_message( + text="Second cached batch item.", + reference=[str(reference_audio_path)], + tokens=6, + ) + ], ] cached = synthesize_moss_tts_local_conversations( @@ -193,7 +238,9 @@ def test_batch_kv_cache_flag_falls_back_to_uncached_path( codec, conversations=conversations, mode="generation", - config=MossTTSLocalGenerationConfig(max_new_tokens=2, do_sample=False, use_kv_cache=True), + config=MossTTSLocalGenerationConfig( + max_new_tokens=2, do_sample=False, use_kv_cache=True + ), ) uncached = synthesize_moss_tts_local_conversations( model, @@ -208,7 +255,9 @@ def test_batch_kv_cache_flag_falls_back_to_uncached_path( ), ) - assert cached.generation.sequences.tolist() == uncached.generation.sequences.tolist() + assert ( + cached.generation.sequences.tolist() == uncached.generation.sequences.tolist() + ) assert len(cached.outputs) == 2 assert cached.outputs[0].waveform.size > 0 assert cached.outputs[1].waveform.size > 0 @@ -216,7 +265,9 @@ def test_batch_kv_cache_flag_falls_back_to_uncached_path( def test_default_cached_path_runs() -> None: model, processor, codec = _runtime() - conversations = [[processor.build_user_message(text="Local cache sample.", tokens=6)]] + conversations = [ + [processor.build_user_message(text="Local cache sample.", tokens=6)] + ] mx.random.seed(0) result = synthesize_moss_tts_local_conversations( diff --git a/tests/integration/test_vibevoice.py b/tests/integration/test_vibevoice.py new file mode 100644 index 0000000..afb4c5b --- /dev/null +++ b/tests/integration/test_vibevoice.py @@ -0,0 +1,73 @@ +"""Manual end-to-end VibeVoice waveform coverage.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import mlx.core as mx +import pytest + +from mlx_speech.audio.io import load_audio +from mlx_speech.generation.vibevoice import ( + VibeVoiceGenerationConfig, + VibeVoiceSynthesisOutput, + synthesize_vibevoice, +) +from mlx_speech.models.vibevoice.checkpoint import load_vibevoice_model +from mlx_speech.models.vibevoice.tokenizer import VibeVoiceTokenizer + + +INT8_DIR = Path("models/vibevoice/mlx-int8") +ORIGINAL_DIR = Path("models/vibevoice/original") +HAS_INT8 = INT8_DIR.is_dir() and any(INT8_DIR.glob("*.safetensors")) +HAS_ORIGINAL = ORIGINAL_DIR.is_dir() and any(ORIGINAL_DIR.glob("*.safetensors")) +MODEL_DIR = INT8_DIR if HAS_INT8 else ORIGINAL_DIR + +pytestmark = [ + pytest.mark.integration, + pytest.mark.skipif( + os.environ.get("RUN_LOCAL_INTEGRATION") != "1", + reason="set RUN_LOCAL_INTEGRATION=1 for local waveform generation", + ), + pytest.mark.skipif( + not (HAS_INT8 or HAS_ORIGINAL), + reason="VibeVoice checkpoint is not present", + ), +] + + +def test_short_generation() -> None: + loaded = load_vibevoice_model(MODEL_DIR, strict=False) + tokenizer = VibeVoiceTokenizer.from_path(MODEL_DIR) + config = VibeVoiceGenerationConfig(max_new_tokens=20, do_sample=False) + + result = synthesize_vibevoice(loaded.model, tokenizer, "Hello.", config=config) + mx.eval(result.waveform) + + assert isinstance(result, VibeVoiceSynthesisOutput) + assert result.sample_rate == 24_000 + assert result.generated_tokens > 0 + assert result.waveform.shape[0] > 0 + + +def test_voice_cloning() -> None: + reference_path = Path("outputs/source/hank_hill_ref.wav") + if not reference_path.is_file(): + pytest.skip("reference audio is not present") + + loaded = load_vibevoice_model(MODEL_DIR, strict=False) + tokenizer = VibeVoiceTokenizer.from_path(MODEL_DIR) + config = VibeVoiceGenerationConfig(max_new_tokens=20, do_sample=False) + reference, _ = load_audio(str(reference_path), sample_rate=24_000) + + result = synthesize_vibevoice( + loaded.model, + tokenizer, + "Hello.", + reference_audio=reference.reshape(1, 1, -1), + config=config, + ) + mx.eval(result.waveform) + + assert result.waveform.shape[0] > 0 diff --git a/tests/runtime/__init__.py b/tests/runtime/__init__.py new file mode 100644 index 0000000..2bebe21 --- /dev/null +++ b/tests/runtime/__init__.py @@ -0,0 +1 @@ +"""Tests that run inference with local checkpoint artifacts.""" diff --git a/tests/test_moss_delay_generation_integration.py b/tests/runtime/test_moss_delay_generation_runtime.py similarity index 85% rename from tests/test_moss_delay_generation_integration.py rename to tests/runtime/test_moss_delay_generation_runtime.py index 0c98803..bfffd33 100644 --- a/tests/test_moss_delay_generation_integration.py +++ b/tests/runtime/test_moss_delay_generation_runtime.py @@ -1,11 +1,24 @@ from __future__ import annotations +from pathlib import Path + import pytest from mlx_speech.generation import MossTTSDelayGenerationConfig, generate_moss_tts_delay -from mlx_speech.models.moss_delay import MossTTSDelayProcessor, load_moss_tts_delay_model +from mlx_speech.models.moss_delay import ( + MossTTSDelayProcessor, + load_moss_tts_delay_model, +) + +MODEL_DIR = Path("models/openmoss/moss_ttsd/mlx-int8") -pytestmark = pytest.mark.local_integration +pytestmark = [ + pytest.mark.runtime, + pytest.mark.skipif( + not MODEL_DIR.is_dir(), + reason="MOSS-TTSD checkpoint is not present", + ), +] @pytest.fixture(scope="module") diff --git a/tests/test_moss_local_processor_integration.py b/tests/runtime/test_moss_local_processor.py similarity index 86% rename from tests/test_moss_local_processor_integration.py rename to tests/runtime/test_moss_local_processor.py index 872da46..c87d2a6 100644 --- a/tests/test_moss_local_processor_integration.py +++ b/tests/runtime/test_moss_local_processor.py @@ -11,9 +11,15 @@ from mlx_speech.models.moss_local import MossTTSLocalProcessor from mlx_speech.models.moss_local.tokenizer import DEFAULT_MOSS_CHAT_TEMPLATE -pytestmark = pytest.mark.local_integration +MODEL_DIR = Path("models/openmoss/moss_tts_local/mlx-int8") -MODEL_DIR = "models/openmoss/moss_tts_local/mlx-int8" +pytestmark = [ + pytest.mark.runtime, + pytest.mark.skipif( + not MODEL_DIR.is_dir(), + reason="MOSS-TTS-Local tokenizer assets are not present", + ), +] def _tiny_codec_config() -> MossAudioTokenizerConfig: @@ -84,7 +90,7 @@ def _load_upstream_processor(): transformers = pytest.importorskip("transformers") torch = pytest.importorskip("torch") - repo_root = Path(__file__).resolve().parents[1] + repo_root = Path(__file__).resolve().parents[2] upstream_root = repo_root / ".references" / "MOSS-TTS" if str(upstream_root) not in sys.path: sys.path.insert(0, str(upstream_root)) @@ -93,7 +99,9 @@ def _load_upstream_processor(): from moss_tts_local.processing_moss_tts import MossTTSDelayProcessor model_dir = repo_root / "models" / "openmoss" / "moss_tts_local" / "mlx-int8" - tokenizer = transformers.AutoTokenizer.from_pretrained(str(model_dir), trust_remote_code=True) + tokenizer = transformers.AutoTokenizer.from_pretrained( + str(model_dir), trust_remote_code=True + ) tokenizer.chat_template = DEFAULT_MOSS_CHAT_TEMPLATE config = MossTTSDelayConfig.from_pretrained(str(model_dir), trust_remote_code=True) processor = MossTTSDelayProcessor( @@ -117,7 +125,9 @@ def _assert_processor_parity( upstream_batch = upstream_processor([upstream_conversation], mode=mode) assert our_batch.input_ids.tolist() == upstream_batch["input_ids"].tolist() - assert our_batch.attention_mask.tolist() == upstream_batch["attention_mask"].tolist() + assert ( + our_batch.attention_mask.tolist() == upstream_batch["attention_mask"].tolist() + ) def test_processor_loads_local_tokenizer_assets() -> None: @@ -172,7 +182,9 @@ def test_processor_generation_parity_with_upstream_direct() -> None: def test_processor_generation_parity_with_upstream_clone() -> None: processor = MossTTSLocalProcessor.from_path(MODEL_DIR) reference_codes = _make_reference_codes() - our_message = processor.build_user_message(text="Clone this.", reference=[reference_codes]) + our_message = processor.build_user_message( + text="Clone this.", reference=[reference_codes] + ) upstream_processor, torch = _load_upstream_processor() upstream_message = upstream_processor.build_user_message( text="Clone this.", @@ -197,7 +209,9 @@ def test_processor_continuation_parity_with_upstream() -> None: ), ] - _assert_processor_parity(our_conversation, upstream_conversation, mode="continuation") + _assert_processor_parity( + our_conversation, upstream_conversation, mode="continuation" + ) def test_processor_continue_clone_parity_with_upstream() -> None: @@ -210,11 +224,15 @@ def test_processor_continue_clone_parity_with_upstream() -> None: upstream_processor, torch = _load_upstream_processor() upstream_codes = torch.tensor(prefix_codes.tolist(), dtype=torch.long) upstream_conversation = [ - upstream_processor.build_user_message(text="Continue clone.", reference=[upstream_codes]), + upstream_processor.build_user_message( + text="Continue clone.", reference=[upstream_codes] + ), upstream_processor.build_assistant_message(audio_codes_list=[upstream_codes]), ] - _assert_processor_parity(our_conversation, upstream_conversation, mode="continuation") + _assert_processor_parity( + our_conversation, upstream_conversation, mode="continuation" + ) @pytest.mark.parametrize( @@ -249,14 +267,25 @@ def test_processor_continue_clone_parity_with_upstream() -> None: ), ], ) -def test_processor_continuation_modes_ignore_tokens_conditioning(with_tokens, without_tokens) -> None: +def test_processor_continuation_modes_ignore_tokens_conditioning( + with_tokens, without_tokens +) -> None: processor = MossTTSLocalProcessor.from_path(MODEL_DIR) prefix_codes = _make_reference_codes() - with_tokens_batch = processor([with_tokens(processor, prefix_codes)], mode="continuation") - without_tokens_batch = processor([without_tokens(processor, prefix_codes)], mode="continuation") + with_tokens_batch = processor( + [with_tokens(processor, prefix_codes)], mode="continuation" + ) + without_tokens_batch = processor( + [without_tokens(processor, prefix_codes)], mode="continuation" + ) - assert with_tokens_batch.input_ids.tolist() == without_tokens_batch.input_ids.tolist() - assert with_tokens_batch.attention_mask.tolist() == without_tokens_batch.attention_mask.tolist() + assert ( + with_tokens_batch.input_ids.tolist() == without_tokens_batch.input_ids.tolist() + ) + assert ( + with_tokens_batch.attention_mask.tolist() + == without_tokens_batch.attention_mask.tolist() + ) def test_processor_audio_helpers_encode_and_decode_with_bound_codec() -> None: diff --git a/tests/runtime/test_vibevoice_acoustic_runtime.py b/tests/runtime/test_vibevoice_acoustic_runtime.py new file mode 100644 index 0000000..6c418e5 --- /dev/null +++ b/tests/runtime/test_vibevoice_acoustic_runtime.py @@ -0,0 +1,44 @@ +"""Real-checkpoint runtime coverage for the VibeVoice acoustic tokenizer.""" + +from pathlib import Path + +import mlx.core as mx +import pytest + +from mlx_speech.models.vibevoice.checkpoint import load_vibevoice_model + + +MODEL_DIR = Path("models/vibevoice/mlx-int8") + +pytestmark = [ + pytest.mark.runtime, + pytest.mark.skipif( + not (MODEL_DIR / "config.json").is_file(), + reason="VibeVoice checkpoint is not present", + ), +] + + +def test_encoder_output_shape() -> None: + loaded = load_vibevoice_model(MODEL_DIR, strict=False) + encoder = loaded.model.model.acoustic_tokenizer.encoder + waveform = mx.random.normal((1, 1, 24_000)) + output = encoder(waveform) + mx.eval(output) + + assert output.shape[0] == 1 + assert output.shape[1] == 64 + assert output.shape[2] >= 7 + + +def test_encode_decode_roundtrip_shape() -> None: + loaded = load_vibevoice_model(MODEL_DIR, strict=False) + tokenizer = loaded.model.model.acoustic_tokenizer + waveform = mx.random.normal((1, 1, 24_000)) + latent = tokenizer.encode(waveform) + reconstructed = tokenizer.decode(latent) + mx.eval(latent, reconstructed) + + assert reconstructed.shape[0] == 1 + assert reconstructed.shape[1] == 1 + assert abs(reconstructed.shape[2] - 24_000) < 3_200 diff --git a/tests/test_cohere_asr_generation.py b/tests/test_cohere_asr_generation.py deleted file mode 100644 index 90998b5..0000000 --- a/tests/test_cohere_asr_generation.py +++ /dev/null @@ -1,62 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import numpy as np - -from mlx_speech.generation.cohere_asr import CohereAsrModel, CohereAsrResult -from mlx_speech.models.cohere_asr.tokenizer import CohereAsrTokenizer - - -def test_tokenizer_prompt_ids_support_punctuation_and_itn() -> None: - tokenizer = CohereAsrTokenizer.from_dir(Path("models/cohere/cohere_transcribe/original")) - - default_prompt = tokenizer.get_decoder_prompt_ids("en") - no_punctuation_prompt = tokenizer.get_decoder_prompt_ids("en", punctuation=False) - itn_prompt = tokenizer.get_decoder_prompt_ids("en", itn=True) - - assert len(default_prompt) == 10 - assert default_prompt[6] == 5 # <|pnc|> - assert no_punctuation_prompt[6] == 6 # <|nopnc|> - assert default_prompt[7] == 9 # <|noitn|> - assert itn_prompt[7] == 8 # <|itn|> - assert default_prompt[8:] == [11, 13] # <|notimestamp|>, <|nodiarize|> - assert itn_prompt[:7] == default_prompt[:7] - - -def test_transcribe_batch_preserves_order_and_options(monkeypatch) -> None: - model = CohereAsrModel(model=None, feature_extractor=None, tokenizer=None, config=None) - calls: list[tuple[int, int, str, bool, bool, int]] = [] - - def fake_transcribe( - self, - audio: np.ndarray, - *, - sample_rate: int, - language: str, - punctuation: bool, - itn: bool, - max_new_tokens: int, - ) -> CohereAsrResult: - calls.append((len(audio), sample_rate, language, punctuation, itn, max_new_tokens)) - return CohereAsrResult(text=f"len={len(audio)}", tokens=[len(audio)], language=language) - - monkeypatch.setattr(CohereAsrModel, "transcribe", fake_transcribe) - - audios = [np.zeros(8, dtype=np.float32), np.zeros(3, dtype=np.float32)] - results = model.transcribe_batch( - audios, - sample_rate=22050, - language="fr", - punctuation=False, - itn=True, - max_new_tokens=32, - ) - - assert [result.text for result in results] == ["len=8", "len=3"] - assert [result.tokens for result in results] == [[8], [3]] - assert calls == [ - (8, 22050, "fr", False, True, 32), - (3, 22050, "fr", False, True, 32), - ] - diff --git a/tests/test_moss_delay_checkpoint_integration.py b/tests/test_moss_delay_checkpoint_integration.py deleted file mode 100644 index 7d428ca..0000000 --- a/tests/test_moss_delay_checkpoint_integration.py +++ /dev/null @@ -1,25 +0,0 @@ -from __future__ import annotations - -import pytest - -from mlx_speech.models.moss_delay import ( - load_moss_tts_delay_model, - resolve_moss_tts_delay_model_dir, -) - -pytestmark = pytest.mark.local_integration - - -def test_resolve_moss_tts_delay_model_dir_defaults_to_local_quantized_runtime() -> None: - resolved = resolve_moss_tts_delay_model_dir() - - assert resolved.as_posix().endswith("models/openmoss/moss_ttsd/mlx-int8") - - -def test_default_ttsd_runtime_loads_quantized_mlx_model() -> None: - loaded = load_moss_tts_delay_model() - - assert loaded.alignment_report.is_exact_match - assert loaded.model.config.n_vq == 16 - assert loaded.model.language_model.config.num_hidden_layers == 36 - assert loaded.quantization is not None diff --git a/tests/test_vibevoice_config.py b/tests/test_vibevoice_config.py deleted file mode 100644 index fac4011..0000000 --- a/tests/test_vibevoice_config.py +++ /dev/null @@ -1,89 +0,0 @@ -"""Tests for VibeVoice config parsing.""" - -from pathlib import Path - -import pytest - -from mlx_speech.models.vibevoice.config import ( - Qwen2LanguageConfig, - VibeVoiceConfig, - VibeVoiceConvTokenizerConfig, -) - -MODEL_DIR = Path("models/vibevoice/original") -HAS_CHECKPOINT = (MODEL_DIR / "config.json").exists() - - -class TestQwen2LanguageConfig: - def test_from_dict_round_trip(self): - cfg = Qwen2LanguageConfig( - hidden_size=3584, - intermediate_size=18944, - num_hidden_layers=28, - num_attention_heads=28, - num_key_value_heads=4, - vocab_size=152064, - ) - rt = Qwen2LanguageConfig.from_dict(cfg.to_dict()) - assert rt.hidden_size == 3584 - assert rt.head_dim == 128 - - def test_extra_preserved(self): - cfg = Qwen2LanguageConfig.from_dict({ - "hidden_size": 256, - "intermediate_size": 1024, - "num_hidden_layers": 2, - "num_attention_heads": 4, - "num_key_value_heads": 2, - "vocab_size": 100, - "torch_dtype": "bfloat16", - }) - assert cfg.extra["torch_dtype"] == "bfloat16" - - -class TestConvTokenizerConfig: - def test_parsed_depths(self): - cfg = VibeVoiceConvTokenizerConfig(vae_dim=64) - assert cfg.parsed_encoder_depths == [3, 3, 3, 3, 3, 3, 8] - assert cfg.parsed_decoder_depths == [8, 3, 3, 3, 3, 3, 3] - - def test_ratios_as_tuple(self): - cfg = VibeVoiceConvTokenizerConfig.from_dict({ - "vae_dim": 64, - "encoder_ratios": [8, 5, 5, 4, 2, 2], - }) - assert isinstance(cfg.encoder_ratios, tuple) - - -class TestVibeVoiceConfig: - @pytest.mark.skipif(not HAS_CHECKPOINT, reason="checkpoint not available") - def test_from_path(self): - cfg = VibeVoiceConfig.from_path(MODEL_DIR) - assert cfg.model_type == "vibevoice" - assert cfg.hidden_size == 3584 - assert cfg.language_config.num_hidden_layers == 28 - assert cfg.acoustic_tokenizer_config.vae_dim == 64 - assert cfg.semantic_tokenizer_config.vae_dim == 128 - assert cfg.diffusion_config.head_layers == 4 - - @pytest.mark.skipif(not HAS_CHECKPOINT, reason="checkpoint not available") - def test_round_trip(self): - cfg = VibeVoiceConfig.from_path(MODEL_DIR) - rt = VibeVoiceConfig.from_dict(cfg.to_dict()) - assert rt.hidden_size == cfg.hidden_size - assert rt.acoustic_tokenizer_config.vae_dim == cfg.acoustic_tokenizer_config.vae_dim - - def test_acostic_typo_handled(self): - """The upstream config.json has a typo: acostic_vae_dim.""" - cfg = VibeVoiceConfig.from_dict({ - "decoder_config": { - "hidden_size": 256, "intermediate_size": 1024, - "num_hidden_layers": 2, "num_attention_heads": 4, - "num_key_value_heads": 2, "vocab_size": 100, - }, - "acoustic_tokenizer_config": {"vae_dim": 64}, - "semantic_tokenizer_config": {"vae_dim": 128}, - "diffusion_head_config": {}, - "acostic_vae_dim": 64, - }) - assert cfg.acoustic_vae_dim == 64 diff --git a/tests/unit/__init__.py b/tests/unit/__init__.py new file mode 100644 index 0000000..b2e7b06 --- /dev/null +++ b/tests/unit/__init__.py @@ -0,0 +1 @@ +"""Fast tests that do not require local model artifacts.""" diff --git a/tests/test_alignment_report.py b/tests/unit/test_alignment_report.py similarity index 100% rename from tests/test_alignment_report.py rename to tests/unit/test_alignment_report.py diff --git a/tests/test_audio_io.py b/tests/unit/test_audio_io.py similarity index 100% rename from tests/test_audio_io.py rename to tests/unit/test_audio_io.py diff --git a/tests/test_batch_generate_longcat_audiodit_script.py b/tests/unit/test_batch_generate_longcat_audiodit_script.py similarity index 90% rename from tests/test_batch_generate_longcat_audiodit_script.py rename to tests/unit/test_batch_generate_longcat_audiodit_script.py index a481cbe..1bf5b82 100644 --- a/tests/test_batch_generate_longcat_audiodit_script.py +++ b/tests/unit/test_batch_generate_longcat_audiodit_script.py @@ -3,7 +3,7 @@ import mlx.core as mx -from scripts.batch_generate_longcat_audiodit import _build_parser, _run_batch_items +from scripts.generate.batch_longcat_audiodit import _build_parser, _run_batch_items def test_batch_script_builds_parser() -> None: @@ -41,14 +41,14 @@ def _fake_write_wav(path: str, samples, sample_rate: int): written.append(Path(path)) monkeypatch.setattr( - "scripts.batch_generate_longcat_audiodit.load_audio", _fake_load_audio + "scripts.generate.batch_longcat_audiodit.load_audio", _fake_load_audio ) monkeypatch.setattr( - "scripts.batch_generate_longcat_audiodit.synthesize_longcat_audiodit", + "scripts.generate.batch_longcat_audiodit.synthesize_longcat_audiodit", _fake_synthesize_longcat_audiodit, ) monkeypatch.setattr( - "scripts.batch_generate_longcat_audiodit.write_wav", _fake_write_wav + "scripts.generate.batch_longcat_audiodit.write_wav", _fake_write_wav ) model = SimpleNamespace(config=SimpleNamespace(sampling_rate=24000)) diff --git a/tests/test_benchmark_moss_ttsd_script.py b/tests/unit/test_benchmark_moss_ttsd_script.py similarity index 95% rename from tests/test_benchmark_moss_ttsd_script.py rename to tests/unit/test_benchmark_moss_ttsd_script.py index 437ff27..f3689e8 100644 --- a/tests/test_benchmark_moss_ttsd_script.py +++ b/tests/unit/test_benchmark_moss_ttsd_script.py @@ -6,7 +6,12 @@ import sys -SCRIPT_PATH = Path(__file__).resolve().parents[1] / "scripts" / "benchmark_moss_ttsd.py" +SCRIPT_PATH = ( + Path(__file__).resolve().parents[2] + / "scripts" + / "eval" + / "benchmark_moss_ttsd.py" +) def _load_script_module(): diff --git a/tests/test_checkpoint_layout.py b/tests/unit/test_checkpoint_layout.py similarity index 100% rename from tests/test_checkpoint_layout.py rename to tests/unit/test_checkpoint_layout.py diff --git a/tests/test_cohere_asr_feature_extraction.py b/tests/unit/test_cohere_asr_feature_extraction.py similarity index 98% rename from tests/test_cohere_asr_feature_extraction.py rename to tests/unit/test_cohere_asr_feature_extraction.py index 6bf61c5..cb01a64 100644 --- a/tests/test_cohere_asr_feature_extraction.py +++ b/tests/unit/test_cohere_asr_feature_extraction.py @@ -9,7 +9,7 @@ def _load_feature_module(): module_path = ( - Path(__file__).resolve().parents[1] + Path(__file__).resolve().parents[2] / "src" / "mlx_speech" / "models" diff --git a/tests/unit/test_cohere_asr_generation.py b/tests/unit/test_cohere_asr_generation.py new file mode 100644 index 0000000..76ab8b1 --- /dev/null +++ b/tests/unit/test_cohere_asr_generation.py @@ -0,0 +1,48 @@ +from __future__ import annotations + +import numpy as np + +from mlx_speech.generation.cohere_asr import CohereAsrModel, CohereAsrResult + + +def test_transcribe_batch_preserves_order_and_options(monkeypatch) -> None: + model = CohereAsrModel( + model=None, feature_extractor=None, tokenizer=None, config=None + ) + calls: list[tuple[int, int, str, bool, bool, int]] = [] + + def fake_transcribe( + self, + audio: np.ndarray, + *, + sample_rate: int, + language: str, + punctuation: bool, + itn: bool, + max_new_tokens: int, + ) -> CohereAsrResult: + calls.append( + (len(audio), sample_rate, language, punctuation, itn, max_new_tokens) + ) + return CohereAsrResult( + text=f"len={len(audio)}", tokens=[len(audio)], language=language + ) + + monkeypatch.setattr(CohereAsrModel, "transcribe", fake_transcribe) + + audios = [np.zeros(8, dtype=np.float32), np.zeros(3, dtype=np.float32)] + results = model.transcribe_batch( + audios, + sample_rate=22050, + language="fr", + punctuation=False, + itn=True, + max_new_tokens=32, + ) + + assert [result.text for result in results] == ["len=8", "len=3"] + assert [result.tokens for result in results] == [[8], [3]] + assert calls == [ + (8, 22050, "fr", False, True, 32), + (3, 22050, "fr", False, True, 32), + ] diff --git a/tests/test_convert_longcat_audiodit_script.py b/tests/unit/test_convert_longcat_audiodit_script.py similarity index 85% rename from tests/test_convert_longcat_audiodit_script.py rename to tests/unit/test_convert_longcat_audiodit_script.py index dfdfeec..ee9647b 100644 --- a/tests/test_convert_longcat_audiodit_script.py +++ b/tests/unit/test_convert_longcat_audiodit_script.py @@ -1,4 +1,4 @@ -from scripts.convert_longcat_audiodit import _build_parser +from scripts.convert.longcat_audiodit import _build_parser def test_convert_script_defaults_to_longcat_layout() -> None: diff --git a/tests/unit/test_dots_tts_checkpoint_contract.py b/tests/unit/test_dots_tts_checkpoint_contract.py index b05826e..898c4ba 100644 --- a/tests/unit/test_dots_tts_checkpoint_contract.py +++ b/tests/unit/test_dots_tts_checkpoint_contract.py @@ -16,7 +16,7 @@ storage_dtype, validate_artifact_dir, ) -from test_dots_tts_config import dots_config, qwen_config +from tests.unit.test_dots_tts_config import dots_config, qwen_config def _metadata(*, variant: str = "soar", artifact_class: str = "base") -> dict: diff --git a/tests/unit/test_dots_tts_convert.py b/tests/unit/test_dots_tts_convert.py index 50f0ebb..94e1611 100644 --- a/tests/unit/test_dots_tts_convert.py +++ b/tests/unit/test_dots_tts_convert.py @@ -22,7 +22,7 @@ remap_core_weights, remap_speaker_weights, ) -from test_dots_tts_config import dots_config, qwen_config +from tests.unit.test_dots_tts_config import dots_config, qwen_config def test_core_remap_uses_runtime_names_and_native_conv_layout() -> None: diff --git a/tests/unit/test_dots_tts_quantization.py b/tests/unit/test_dots_tts_quantization.py index ae4652c..62ee90c 100644 --- a/tests/unit/test_dots_tts_quantization.py +++ b/tests/unit/test_dots_tts_quantization.py @@ -16,7 +16,7 @@ quantize_dots_tts_core, ) from mlx_speech.models.dots_tts.config import DotsTTSConfig, DotsTTSQwenConfig -from test_dots_tts_config import dots_config +from tests.unit.test_dots_tts_config import dots_config def _core() -> DotsTTSCoreComponents: diff --git a/tests/unit/test_dots_tts_vocoder.py b/tests/unit/test_dots_tts_vocoder.py index 9d675e5..398236d 100644 --- a/tests/unit/test_dots_tts_vocoder.py +++ b/tests/unit/test_dots_tts_vocoder.py @@ -13,7 +13,7 @@ CausalConvTranspose1d, Conv1d, ) -from test_dots_tts_audio_vae import _config +from tests.unit.test_dots_tts_audio_vae import _config def test_causal_convolution_does_not_see_future_inputs() -> None: diff --git a/tests/unit/test_dots_tts_vocoder_streaming.py b/tests/unit/test_dots_tts_vocoder_streaming.py index 8994ad0..0a7345f 100644 --- a/tests/unit/test_dots_tts_vocoder_streaming.py +++ b/tests/unit/test_dots_tts_vocoder_streaming.py @@ -12,7 +12,7 @@ VocoderDecodeState, ) from mlx_speech.models.dots_tts.vocoder import AliasFreeSnakeBeta -from test_dots_tts_audio_vae import _config +from tests.unit.test_dots_tts_audio_vae import _config def _model(seed: int) -> AudioVAE: diff --git a/tests/test_generate_longcat_audiodit_script.py b/tests/unit/test_generate_longcat_audiodit_script.py similarity index 83% rename from tests/test_generate_longcat_audiodit_script.py rename to tests/unit/test_generate_longcat_audiodit_script.py index 36dbfe0..7d26c8d 100644 --- a/tests/test_generate_longcat_audiodit_script.py +++ b/tests/unit/test_generate_longcat_audiodit_script.py @@ -1,4 +1,4 @@ -from scripts.generate_longcat_audiodit import _build_parser +from scripts.generate.longcat_audiodit import _build_parser def test_generate_script_builds_parser() -> None: diff --git a/tests/test_generate_moss_sound_effect_script.py b/tests/unit/test_generate_moss_sound_effect_script.py similarity index 95% rename from tests/test_generate_moss_sound_effect_script.py rename to tests/unit/test_generate_moss_sound_effect_script.py index f5c8b9e..158f206 100644 --- a/tests/test_generate_moss_sound_effect_script.py +++ b/tests/unit/test_generate_moss_sound_effect_script.py @@ -6,7 +6,12 @@ import sys -SCRIPT_PATH = Path(__file__).resolve().parents[1] / "scripts" / "generate_moss_sound_effect.py" +SCRIPT_PATH = ( + Path(__file__).resolve().parents[2] + / "scripts" + / "generate" + / "moss_sound_effect.py" +) def _load_script_module(): diff --git a/tests/test_generate_moss_ttsd_script.py b/tests/unit/test_generate_moss_ttsd_script.py similarity index 97% rename from tests/test_generate_moss_ttsd_script.py rename to tests/unit/test_generate_moss_ttsd_script.py index 1b849f9..1abd912 100644 --- a/tests/test_generate_moss_ttsd_script.py +++ b/tests/unit/test_generate_moss_ttsd_script.py @@ -6,7 +6,12 @@ import sys -SCRIPT_PATH = Path(__file__).resolve().parents[1] / "scripts" / "generate_moss_ttsd.py" +SCRIPT_PATH = ( + Path(__file__).resolve().parents[2] + / "scripts" + / "generate" + / "moss_ttsd.py" +) def _load_script_module(): diff --git a/tests/test_generate_step_audio_editx_script.py b/tests/unit/test_generate_step_audio_editx_script.py similarity index 98% rename from tests/test_generate_step_audio_editx_script.py rename to tests/unit/test_generate_step_audio_editx_script.py index 99f9146..fadf4bf 100644 --- a/tests/test_generate_step_audio_editx_script.py +++ b/tests/unit/test_generate_step_audio_editx_script.py @@ -10,7 +10,7 @@ from mlx_speech.generation.step_audio_editx import StepAudioEditXResult SCRIPT_PATH = ( - Path(__file__).resolve().parents[1] + Path(__file__).resolve().parents[2] / "scripts" / "generate" / "step_audio_editx.py" diff --git a/tests/test_generate_vibevoice_script.py b/tests/unit/test_generate_vibevoice_script.py similarity index 74% rename from tests/test_generate_vibevoice_script.py rename to tests/unit/test_generate_vibevoice_script.py index 46bb376..b765046 100644 --- a/tests/test_generate_vibevoice_script.py +++ b/tests/unit/test_generate_vibevoice_script.py @@ -6,11 +6,15 @@ import sys -SCRIPT_PATH = Path(__file__).resolve().parents[1] / "scripts" / "generate_vibevoice.py" +SCRIPT_PATH = ( + Path(__file__).resolve().parents[2] / "scripts" / "generate" / "vibevoice.py" +) def _load_script_module(): - spec = importlib.util.spec_from_file_location("generate_vibevoice_script", SCRIPT_PATH) + spec = importlib.util.spec_from_file_location( + "generate_vibevoice_script", SCRIPT_PATH + ) if spec is None or spec.loader is None: raise RuntimeError(f"Unable to load script module from {SCRIPT_PATH}.") module = importlib.util.module_from_spec(spec) @@ -31,10 +35,10 @@ def _args(**overrides): "diffusion_steps_fast": None, "diffusion_warmup_frames": 10, "max_new_tokens": 2048, - "temperature": 1.0, - "top_p": 1.0, - "seed": None, - "greedy": False, + "temperature": 0.95, + "top_p": 0.95, + "seed": 42, + "no_greedy": False, } data.update(overrides) return argparse.Namespace(**data) @@ -50,29 +54,29 @@ def test_parser_help_documents_sampling_controls() -> None: assert "--temperature" in help_text assert "--top-p" in help_text assert "--seed" in help_text - assert "--greedy" in help_text + assert "--no-greedy" in help_text -def test_build_generation_config_uses_sampling_defaults() -> None: +def test_build_generation_config_uses_greedy_defaults() -> None: module = _load_script_module() config = module._build_generation_config(_args()) - assert config.do_sample is True - assert config.temperature == 1.0 - assert config.top_p == 1.0 - assert config.seed is None + assert config.do_sample is False + assert config.temperature == 0.0 + assert config.top_p == 0.95 + assert config.seed == 42 assert config.diffusion_steps_fast is None assert config.diffusion_warmup_frames == 10 -def test_build_generation_config_can_force_greedy_with_seed() -> None: +def test_build_generation_config_can_enable_sampling_with_seed() -> None: module = _load_script_module() - config = module._build_generation_config(_args(greedy=True, seed=123, top_p=0.8)) + config = module._build_generation_config(_args(no_greedy=True, seed=123, top_p=0.8)) - assert config.do_sample is False - assert config.temperature == 0.0 + assert config.do_sample is True + assert config.temperature == 0.95 assert config.top_p == 0.8 assert config.seed == 123 diff --git a/tests/test_moss_audio_tokenizer.py b/tests/unit/test_moss_audio_tokenizer.py similarity index 100% rename from tests/test_moss_audio_tokenizer.py rename to tests/unit/test_moss_audio_tokenizer.py diff --git a/tests/test_moss_delay_config.py b/tests/unit/test_moss_delay_config.py similarity index 72% rename from tests/test_moss_delay_config.py rename to tests/unit/test_moss_delay_config.py index 05a488a..e4426d8 100644 --- a/tests/test_moss_delay_config.py +++ b/tests/unit/test_moss_delay_config.py @@ -1,7 +1,25 @@ import json from pathlib import Path -from mlx_speech.models.moss_delay import MossTTSDelayConfig +from mlx_speech.models.moss_delay import ( + MossTTSDelayConfig, + resolve_moss_tts_delay_model_dir, +) +from mlx_speech.models.moss_delay import checkpoint as checkpoint_module + + +def test_resolve_moss_tts_delay_model_dir_defaults_to_local_quantized_runtime( + tmp_path: Path, + monkeypatch, +) -> None: + model_dir = tmp_path / "openmoss" / "moss_ttsd" / "mlx-int8" + model_dir.mkdir(parents=True) + (model_dir / "model.safetensors").write_bytes(b"fixture") + monkeypatch.setattr(checkpoint_module, "MODELS_ROOT", tmp_path) + + resolved = resolve_moss_tts_delay_model_dir() + + assert resolved == model_dir def test_moss_delay_config_from_dict_exposes_derived_properties() -> None: diff --git a/tests/test_moss_delay_dialogue.py b/tests/unit/test_moss_delay_dialogue.py similarity index 100% rename from tests/test_moss_delay_dialogue.py rename to tests/unit/test_moss_delay_dialogue.py diff --git a/tests/test_moss_delay_generation.py b/tests/unit/test_moss_delay_generation.py similarity index 100% rename from tests/test_moss_delay_generation.py rename to tests/unit/test_moss_delay_generation.py diff --git a/tests/test_moss_delay_model.py b/tests/unit/test_moss_delay_model.py similarity index 100% rename from tests/test_moss_delay_model.py rename to tests/unit/test_moss_delay_model.py diff --git a/tests/test_moss_delay_processor.py b/tests/unit/test_moss_delay_processor.py similarity index 100% rename from tests/test_moss_delay_processor.py rename to tests/unit/test_moss_delay_processor.py diff --git a/tests/test_moss_delay_quantized_checkpoint.py b/tests/unit/test_moss_delay_quantized_checkpoint.py similarity index 100% rename from tests/test_moss_delay_quantized_checkpoint.py rename to tests/unit/test_moss_delay_quantized_checkpoint.py diff --git a/tests/test_moss_local_cache.py b/tests/unit/test_moss_local_cache.py similarity index 100% rename from tests/test_moss_local_cache.py rename to tests/unit/test_moss_local_cache.py diff --git a/tests/test_moss_local_config.py b/tests/unit/test_moss_local_config.py similarity index 100% rename from tests/test_moss_local_config.py rename to tests/unit/test_moss_local_config.py diff --git a/tests/test_moss_local_generation.py b/tests/unit/test_moss_local_generation.py similarity index 98% rename from tests/test_moss_local_generation.py rename to tests/unit/test_moss_local_generation.py index dc422d7..68665a6 100644 --- a/tests/test_moss_local_generation.py +++ b/tests/unit/test_moss_local_generation.py @@ -20,7 +20,12 @@ def _load_generate_script_module(): - script_path = Path(__file__).resolve().parents[1] / "scripts" / "generate_moss_local.py" + script_path = ( + Path(__file__).resolve().parents[2] + / "scripts" + / "generate" + / "moss_local.py" + ) spec = importlib.util.spec_from_file_location("generate_moss_local_script", script_path) if spec is None or spec.loader is None: raise RuntimeError(f"Unable to load script module from {script_path}.") diff --git a/tests/test_moss_local_model.py b/tests/unit/test_moss_local_model.py similarity index 100% rename from tests/test_moss_local_model.py rename to tests/unit/test_moss_local_model.py diff --git a/tests/test_moss_local_quantized_checkpoint.py b/tests/unit/test_moss_local_quantized_checkpoint.py similarity index 100% rename from tests/test_moss_local_quantized_checkpoint.py rename to tests/unit/test_moss_local_quantized_checkpoint.py diff --git a/tests/test_moss_sound_effect.py b/tests/unit/test_moss_sound_effect.py similarity index 100% rename from tests/test_moss_sound_effect.py rename to tests/unit/test_moss_sound_effect.py diff --git a/tests/test_runtime_purity.py b/tests/unit/test_runtime_purity.py similarity index 91% rename from tests/test_runtime_purity.py rename to tests/unit/test_runtime_purity.py index 2dbaade..f314688 100644 --- a/tests/test_runtime_purity.py +++ b/tests/unit/test_runtime_purity.py @@ -31,7 +31,7 @@ def _import_roots(path: Path) -> set[str]: def test_runtime_modules_do_not_import_banned_dependency_stacks() -> None: - runtime_root = Path(__file__).resolve().parents[1] / "src" / "mlx_speech" + runtime_root = Path(__file__).resolve().parents[2] / "src" / "mlx_speech" bad_files: list[str] = [] for path in runtime_root.rglob("*.py"): banned = sorted(_import_roots(path) & BANNED_RUNTIME_IMPORTS) @@ -41,7 +41,7 @@ def test_runtime_modules_do_not_import_banned_dependency_stacks() -> None: def test_runtime_modules_do_not_reference_upstream_qwen_asr_distribution() -> None: - runtime_root = Path(__file__).resolve().parents[1] / "src" / "mlx_speech" + runtime_root = Path(__file__).resolve().parents[2] / "src" / "mlx_speech" bad_files: list[str] = [] for path in runtime_root.rglob("*.py"): text = path.read_text(encoding="utf-8") diff --git a/tests/test_sharded_checkpoint.py b/tests/unit/test_sharded_checkpoint.py similarity index 100% rename from tests/test_sharded_checkpoint.py rename to tests/unit/test_sharded_checkpoint.py diff --git a/tests/unit/test_tts_streaming_protocol.py b/tests/unit/test_tts_streaming_protocol.py index 2679ac9..4fa261b 100644 --- a/tests/unit/test_tts_streaming_protocol.py +++ b/tests/unit/test_tts_streaming_protocol.py @@ -3,7 +3,7 @@ from mlx_speech.tts import StreamingTTSModel, TTSOutput from mlx_speech.tts._adapters.dots_tts import DotsTTSAdapter -from test_dots_tts_adapter import _Generator +from tests.unit.test_dots_tts_adapter import _Generator class _NonStreamingModel: diff --git a/tests/test_vibevoice_acoustic.py b/tests/unit/test_vibevoice_acoustic.py similarity index 73% rename from tests/test_vibevoice_acoustic.py rename to tests/unit/test_vibevoice_acoustic.py index e5dadf2..baf8986 100644 --- a/tests/test_vibevoice_acoustic.py +++ b/tests/unit/test_vibevoice_acoustic.py @@ -1,7 +1,5 @@ """Tests for VibeVoice acoustic tokenizer.""" -from pathlib import Path - import mlx.core as mx import pytest @@ -16,9 +14,6 @@ ) from mlx_speech.models.vibevoice.config import VibeVoiceConvTokenizerConfig -MODEL_DIR = Path("models/vibevoice/mlx-int8") -HAS_CHECKPOINT = (MODEL_DIR / "config.json").exists() - class TestCausalConv1d: def test_output_shape_no_stride(self): @@ -121,7 +116,9 @@ def test_streaming(self): class TestEncoder: def test_tiny_config(self): cfg = VibeVoiceConvTokenizerConfig( - vae_dim=8, encoder_ratios=(2, 2), encoder_depths="2-2-2", + vae_dim=8, + encoder_ratios=(2, 2), + encoder_depths="2-2-2", encoder_n_filters=4, ) enc = VibeVoiceConvEncoder(cfg) @@ -134,43 +131,14 @@ def test_tiny_config(self): class TestDecoder: def test_tiny_config(self): cfg = VibeVoiceConvTokenizerConfig( - vae_dim=8, encoder_ratios=(2, 2), encoder_depths="2-2-2", - encoder_n_filters=4, decoder_n_filters=4, + vae_dim=8, + encoder_ratios=(2, 2), + encoder_depths="2-2-2", + encoder_n_filters=4, + decoder_n_filters=4, ) dec = VibeVoiceConvDecoder(cfg) x = mx.random.normal((1, 8, 4)) # (B, vae_dim, T_frames) out = dec(x) assert out.shape[0] == 1 assert out.shape[1] == 1 # channels - - -@pytest.mark.skipif(not HAS_CHECKPOINT, reason="checkpoint not available") -class TestRealCheckpoint: - def test_encoder_output_shape(self): - from mlx_speech.models.vibevoice.checkpoint import load_vibevoice_model - - loaded = load_vibevoice_model(MODEL_DIR, strict=False) - enc = loaded.model.model.acoustic_tokenizer.encoder - x = mx.random.normal((1, 1, 24000)) # 1 second at 24kHz - out = enc(x) - mx.eval(out) - # 24000 / 3200 ≈ 7.5 frames - assert out.shape[0] == 1 - assert out.shape[1] == 64 # vae_dim - assert out.shape[2] >= 7 - - def test_encode_decode_roundtrip_shape(self): - from mlx_speech.models.vibevoice.checkpoint import load_vibevoice_model - - loaded = load_vibevoice_model(MODEL_DIR, strict=False) - at = loaded.model.model.acoustic_tokenizer - x = mx.random.normal((1, 1, 24000)) - latent = at.encode(x) - mx.eval(latent) - # Decode - recon = at.decode(latent) - mx.eval(recon) - assert recon.shape[0] == 1 - assert recon.shape[1] == 1 - # Reconstructed length should be close to original - assert abs(recon.shape[2] - 24000) < 3200 # within one frame diff --git a/tests/test_vibevoice_checkpoint.py b/tests/unit/test_vibevoice_checkpoint.py similarity index 61% rename from tests/test_vibevoice_checkpoint.py rename to tests/unit/test_vibevoice_checkpoint.py index f70c792..5ba375b 100644 --- a/tests/test_vibevoice_checkpoint.py +++ b/tests/unit/test_vibevoice_checkpoint.py @@ -1,25 +1,17 @@ """Tests for VibeVoice checkpoint loading.""" -from pathlib import Path - import mlx.core as mx -import pytest - -from mlx_speech.models.vibevoice.checkpoint import ( - load_vibevoice_checkpoint, - load_vibevoice_model, - sanitize_state_dict, -) -MODEL_DIR = Path("models/vibevoice/mlx-int8") -HAS_CHECKPOINT = any(MODEL_DIR.glob("*.safetensors")) if MODEL_DIR.exists() else False +from mlx_speech.models.vibevoice.checkpoint import sanitize_state_dict class TestSanitize: def test_conv1d_transpose(self): """Conv1d weights should be transposed from (out, in, k) to (out, k, in).""" weights = { - "model.acoustic_tokenizer.decoder.head.conv.conv.weight": mx.zeros((16, 32, 7)), + "model.acoustic_tokenizer.decoder.head.conv.conv.weight": mx.zeros( + (16, 32, 7) + ), "model.acoustic_tokenizer.decoder.head.conv.conv.bias": mx.zeros((16,)), } sanitized, _, _ = sanitize_state_dict(weights) @@ -29,10 +21,14 @@ def test_conv1d_transpose(self): def test_convtr_transpose(self): """ConvTranspose1d weights: (in, out, k) → (out, k, in).""" weights = { - "model.acoustic_tokenizer.decoder.upsample_layers.1.0.convtr.convtr.weight": mx.zeros((64, 32, 16)), + "model.acoustic_tokenizer.decoder.upsample_layers.1.0.convtr.convtr.weight": mx.zeros( + (64, 32, 16) + ), } sanitized, _, _ = sanitize_state_dict(weights) - key = "model.acoustic_tokenizer.decoder.upsample_layers.1.0.convtr.convtr.weight" + key = ( + "model.acoustic_tokenizer.decoder.upsample_layers.1.0.convtr.convtr.weight" + ) w = sanitized[key] assert w.shape == (32, 16, 64) # (out, k, in) @@ -49,16 +45,3 @@ def test_skip_inv_freq(self): sanitized, skipped, _ = sanitize_state_dict(weights) assert "model.language_model.rotary_emb.inv_freq" in skipped assert "model.language_model.rotary_emb.inv_freq" not in sanitized - - -@pytest.mark.skipif(not HAS_CHECKPOINT, reason="checkpoint not available") -@pytest.mark.local_integration -class TestRealCheckpoint: - def test_load_checkpoint(self): - ckpt = load_vibevoice_checkpoint(MODEL_DIR) - assert ckpt.key_count > 0 - assert ckpt.config.model_type == "vibevoice" - - def test_model_alignment(self): - loaded = load_vibevoice_model(MODEL_DIR, strict=False) - assert loaded.alignment_report.is_exact_match diff --git a/tests/unit/test_vibevoice_config.py b/tests/unit/test_vibevoice_config.py new file mode 100644 index 0000000..db95454 --- /dev/null +++ b/tests/unit/test_vibevoice_config.py @@ -0,0 +1,74 @@ +"""Tests for VibeVoice config parsing.""" + +from mlx_speech.models.vibevoice.config import ( + Qwen2LanguageConfig, + VibeVoiceConfig, + VibeVoiceConvTokenizerConfig, +) + + +class TestQwen2LanguageConfig: + def test_from_dict_round_trip(self): + cfg = Qwen2LanguageConfig( + hidden_size=3584, + intermediate_size=18944, + num_hidden_layers=28, + num_attention_heads=28, + num_key_value_heads=4, + vocab_size=152064, + ) + rt = Qwen2LanguageConfig.from_dict(cfg.to_dict()) + assert rt.hidden_size == 3584 + assert rt.head_dim == 128 + + def test_extra_preserved(self): + cfg = Qwen2LanguageConfig.from_dict( + { + "hidden_size": 256, + "intermediate_size": 1024, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 100, + "torch_dtype": "bfloat16", + } + ) + assert cfg.extra["torch_dtype"] == "bfloat16" + + +class TestConvTokenizerConfig: + def test_parsed_depths(self): + cfg = VibeVoiceConvTokenizerConfig(vae_dim=64) + assert cfg.parsed_encoder_depths == [3, 3, 3, 3, 3, 3, 8] + assert cfg.parsed_decoder_depths == [8, 3, 3, 3, 3, 3, 3] + + def test_ratios_as_tuple(self): + cfg = VibeVoiceConvTokenizerConfig.from_dict( + { + "vae_dim": 64, + "encoder_ratios": [8, 5, 5, 4, 2, 2], + } + ) + assert isinstance(cfg.encoder_ratios, tuple) + + +class TestVibeVoiceConfig: + def test_acostic_typo_handled(self): + """The upstream config.json has a typo: acostic_vae_dim.""" + cfg = VibeVoiceConfig.from_dict( + { + "decoder_config": { + "hidden_size": 256, + "intermediate_size": 1024, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 100, + }, + "acoustic_tokenizer_config": {"vae_dim": 64}, + "semantic_tokenizer_config": {"vae_dim": 128}, + "diffusion_head_config": {}, + "acostic_vae_dim": 64, + } + ) + assert cfg.acoustic_vae_dim == 64 diff --git a/tests/test_vibevoice_diffusion.py b/tests/unit/test_vibevoice_diffusion.py similarity index 100% rename from tests/test_vibevoice_diffusion.py rename to tests/unit/test_vibevoice_diffusion.py diff --git a/tests/test_vibevoice_generation.py b/tests/unit/test_vibevoice_generation.py similarity index 57% rename from tests/test_vibevoice_generation.py rename to tests/unit/test_vibevoice_generation.py index 7e62abc..24d3a02 100644 --- a/tests/test_vibevoice_generation.py +++ b/tests/unit/test_vibevoice_generation.py @@ -1,25 +1,15 @@ """Tests for VibeVoice generation loop.""" -from pathlib import Path - import mlx.core as mx -import pytest from mlx_speech.generation.vibevoice import ( VibeVoiceGenerationConfig, - VibeVoiceSynthesisOutput, _apply_top_p, _constrain_logits, _format_text_input, _sample_next_token, ) -MODEL_DIR = Path("models/vibevoice/mlx-int8") -HAS_INT8 = any(MODEL_DIR.glob("*.safetensors")) if MODEL_DIR.exists() else False -ORIGINAL_DIR = Path("models/vibevoice/original") -HAS_ORIGINAL = any(ORIGINAL_DIR.glob("*.safetensors")) if ORIGINAL_DIR.exists() else False -HAS_MODEL = HAS_INT8 or HAS_ORIGINAL - class TestConstrainLogits: def test_masks_invalid_tokens(self): @@ -112,52 +102,3 @@ def test_existing_speaker_labels_are_preserved(self): def test_bracket_speaker_labels_are_treated_as_plain_text(self): text = "[1]: Hello.\n[2]: Hi." assert _format_text_input(text) == "Speaker 0: [1]: Hello. [2]: Hi." - - -@pytest.mark.skipif(not HAS_MODEL, reason="model not available") -@pytest.mark.local_integration -class TestEndToEnd: - def _get_model_dir(self): - return MODEL_DIR if HAS_INT8 else ORIGINAL_DIR - - def test_short_generation(self): - from mlx_speech.models.vibevoice.checkpoint import load_vibevoice_model - from mlx_speech.models.vibevoice.tokenizer import VibeVoiceTokenizer - from mlx_speech.generation.vibevoice import synthesize_vibevoice - - model_dir = self._get_model_dir() - loaded = load_vibevoice_model(model_dir, strict=False) - tok = VibeVoiceTokenizer.from_path(model_dir) - config = VibeVoiceGenerationConfig(max_new_tokens=20, do_sample=False) - - result = synthesize_vibevoice(loaded.model, tok, "Hello.", config=config) - mx.eval(result.waveform) - - assert isinstance(result, VibeVoiceSynthesisOutput) - assert result.sample_rate == 24000 - assert result.generated_tokens > 0 - assert result.waveform.shape[0] > 0 - - def test_voice_cloning(self): - from mlx_speech.models.vibevoice.checkpoint import load_vibevoice_model - from mlx_speech.models.vibevoice.tokenizer import VibeVoiceTokenizer - from mlx_speech.generation.vibevoice import synthesize_vibevoice - from mlx_speech.audio.io import load_audio - - ref_path = Path("outputs/source/hank_hill_ref.wav") - if not ref_path.exists(): - pytest.skip("reference audio not available") - - model_dir = self._get_model_dir() - loaded = load_vibevoice_model(model_dir, strict=False) - tok = VibeVoiceTokenizer.from_path(model_dir) - config = VibeVoiceGenerationConfig(max_new_tokens=20, do_sample=False) - - ref_raw, _ = load_audio(str(ref_path), sample_rate=24000) - result = synthesize_vibevoice( - loaded.model, tok, "Hello.", - reference_audio=ref_raw.reshape(1, 1, -1), - config=config, - ) - mx.eval(result.waveform) - assert result.waveform.shape[0] > 0 diff --git a/uv.lock b/uv.lock index b96b4b8..347a53e 100644 --- a/uv.lock +++ b/uv.lock @@ -98,6 +98,90 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335, upload-time = "2022-10-25T02:36:20.889Z" }, ] +[[package]] +name = "coverage" +version = "7.15.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/be/c3/4f2195f512fb172aa425a8803a874b2baa9ba7f80ff7b6080998761fc701/coverage-7.15.4.tar.gz", hash = "sha256:0548198fff07ccf4faf469520bce1c2eceb1ce3e62891921138dec10907f9d00", size = 936952, upload-time = "2026-08-06T13:50:24.442Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/84/651a9310859673aaa3b3203f1aa1641ca60fcf2494683e1c9474c7172780/coverage-7.15.4-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c705b28feb2775dc82a25f1d473a370bc37ff93f5177f4e29ce2425f560f6921", size = 222565, upload-time = "2026-08-06T13:48:00.796Z" }, + { url = "https://files.pythonhosted.org/packages/82/f9/4dcf700137e8af550670f4d74d1b63828ce93e1e2b05e5f10710eb2ea987/coverage-7.15.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3ff205ab5e3ecc670f6a4dd19d9cbf12ede53dd41cfc1e15716ec961ea6d314e", size = 222936, upload-time = "2026-08-06T13:48:02.391Z" }, + { url = "https://files.pythonhosted.org/packages/07/4a/612ff1e780b3fbfd637486f542f84adc5503873d8b5d279dec1ffeef9414/coverage-7.15.4-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5172326e861a38b48b48befca15e0f477a26b283337a33a739c8fed229934e36", size = 253926, upload-time = "2026-08-06T13:48:04.382Z" }, + { url = "https://files.pythonhosted.org/packages/b0/04/d1cff1c2ead4708a6a79c01d3736b6a25bd38a36678398f72a8dd33dfad9/coverage-7.15.4-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:12b59c90084e3234fb11184886bf4a40f4f16a8c8f867be2e087b81f8e8868d4", size = 256523, upload-time = "2026-08-06T13:48:05.996Z" }, + { url = "https://files.pythonhosted.org/packages/b9/80/d34e13fb4b293cbdb9665838cf5522077b8ad14ef947550631a4bced36a5/coverage-7.15.4-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:349062d66f00b40fa2c1c222438bad25fabf755631b5d82937fe985c8008615c", size = 257759, upload-time = "2026-08-06T13:48:08.036Z" }, + { url = "https://files.pythonhosted.org/packages/0f/e7/2c5fe7636fdb0732fe0f09f308a5b066864078b7fc61f6678e8478554f2e/coverage-7.15.4-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4256ced708e598e05209bc1a8ab4074e04a51dba4c62fb45926a229af675ace7", size = 259890, upload-time = "2026-08-06T13:48:09.834Z" }, + { url = "https://files.pythonhosted.org/packages/92/28/9689f0858dfff59c2ea688938ab9fa2925631235df67126a42b6c5c70ae1/coverage-7.15.4-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d80f974b20782d9612c8b4c9beeca867074c7cf4079d1419843fa25a26428b25", size = 254121, upload-time = "2026-08-06T13:48:11.459Z" }, + { url = "https://files.pythonhosted.org/packages/f9/e2/785077c230c157243eb5aa9a26c3be260ecd02001bead54a3cada3df8e03/coverage-7.15.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2e179f19bfe1d31f8eeeaa12990194d761c4f62f0759661000bca6cd8729f40b", size = 255891, upload-time = "2026-08-06T13:48:13.209Z" }, + { url = "https://files.pythonhosted.org/packages/d4/90/e20371b17b40f912f21305c2db2f30efa3de306f7320fc916804872c85a4/coverage-7.15.4-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:8bc16bb47b7679670eceff71d78bfb7d6e5b143f6c2cd117487ec7c75e0d4b78", size = 253859, upload-time = "2026-08-06T13:48:14.736Z" }, + { url = "https://files.pythonhosted.org/packages/05/49/25371987ee459a5f67c0427fb75c74f9358e65f2c71fe75bf41c1b6c5fcb/coverage-7.15.4-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:1cd685005cd2c4200adfc14cf39a603b9320efab3f18a8f7f156d20c9cc3345f", size = 258011, upload-time = "2026-08-06T13:48:16.464Z" }, + { url = "https://files.pythonhosted.org/packages/30/6e/32e67467f6154bf4f1c4f63b05acc5097cba4237d45bbeeea446b52e8ac1/coverage-7.15.4-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:337399ad2c93b3acd2a937627dae8b3e86b66707cd3d3e856347999aadf1ef8d", size = 253676, upload-time = "2026-08-06T13:48:18.493Z" }, + { url = "https://files.pythonhosted.org/packages/03/c1/8b24192e89286399765155251f99ee9f070a9d637109018ac23d99b99f6f/coverage-7.15.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:96e257121228ec5cd2bb919276e94ac11074471bc37d68dbae0e8308cce15fff", size = 255453, upload-time = "2026-08-06T13:48:20.057Z" }, + { url = "https://files.pythonhosted.org/packages/16/6f/8b41ebdf67c87854e17c035336a90f1cfbad0c14c2a584301be6ff148718/coverage-7.15.4-cp313-cp313-win32.whl", hash = "sha256:c65a9e0dfc6143491879da4e13b5e30f8be192055de508d737fb14601edbd22c", size = 224605, upload-time = "2026-08-06T13:48:21.655Z" }, + { url = "https://files.pythonhosted.org/packages/e0/e2/2946c7f0b42b152ecb21ff1bdad72e3d301e790c0c487e4a86e8c9f69347/coverage-7.15.4-cp313-cp313-win_amd64.whl", hash = "sha256:2ff8f5e9b8f7a94f0c11c45631eee103dbcb7d63274edd12c56efe1be690b3b4", size = 225148, upload-time = "2026-08-06T13:48:23.376Z" }, + { url = "https://files.pythonhosted.org/packages/9e/83/3f4a69957f48ae7a0aba76c34743f88963d607b19e03f3f8e66f91cae0f9/coverage-7.15.4-cp313-cp313-win_arm64.whl", hash = "sha256:6e0a8a5083b096487d6cfced94cdd514d8f5db6f113610fb36c0620edb1028cf", size = 224536, upload-time = "2026-08-06T13:48:25.117Z" }, + { url = "https://files.pythonhosted.org/packages/ea/ac/748cf29eeb2d6be34a3176ce26a4f49e38085ee08e8935f05f6f26ed7e0f/coverage-7.15.4-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:770e9325ab5ea6d56f77e59b29ecfe0ac20b57a82a601876f90494a4dda0386f", size = 222608, upload-time = "2026-08-06T13:48:26.806Z" }, + { url = "https://files.pythonhosted.org/packages/0b/02/1abbf5c984677b0aa439cdacaccbf38d248939d8ef8fe1cc7a50d73edb77/coverage-7.15.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d12b33a3a50a1676b7784dc8d00a0c6d66a9f2add4b85a041c19b6a7e53ef23c", size = 222940, upload-time = "2026-08-06T13:48:28.432Z" }, + { url = "https://files.pythonhosted.org/packages/eb/e1/ff8f9f53d9fcf586125b55d0b1f04ec1c14955fee41e83d5814bee141bb5/coverage-7.15.4-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5669c8378ebde86f5def7a25d29586631b58acc27ffde04399f678f3dfc6e082", size = 253985, upload-time = "2026-08-06T13:48:29.995Z" }, + { url = "https://files.pythonhosted.org/packages/a1/26/595759762e514e81be1d7d01ed03444303bcd152226a6529998d253f9201/coverage-7.15.4-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:ff97a14362eef486483ed44042ca2027ea257df6ff768e62358ee0c9776925ac", size = 256492, upload-time = "2026-08-06T13:48:31.634Z" }, + { url = "https://files.pythonhosted.org/packages/24/68/b79aabac54d482be23b5fcdd4f4662bff24a78edc4ee29201726929936d5/coverage-7.15.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5a325e815318638aed1655d9c06e6d7c2d3d46c09231ce988070428a8762d734", size = 257837, upload-time = "2026-08-06T13:48:33.186Z" }, + { url = "https://files.pythonhosted.org/packages/09/0f/bf7f297885a5bf6fd71e5782404e0ff059ca09e8711ceb3a08544abde45a/coverage-7.15.4-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:474223409d88eb20d2d6a0d37ea60e8647a65a90cc008dc1f0410af5f64f1e0d", size = 260152, upload-time = "2026-08-06T13:48:34.75Z" }, + { url = "https://files.pythonhosted.org/packages/fd/f1/296744e854ff8368542343457414380465e9ceefb9192342feb9d3bc461d/coverage-7.15.4-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7f2f62ae3cd189dd2e13aece758c57b3eecbd27be070dbd4cbd10936049e5dbf", size = 253978, upload-time = "2026-08-06T13:48:36.434Z" }, + { url = "https://files.pythonhosted.org/packages/55/b0/bbdb2e9057493e66220a2e149ca2d301ba0e3a58a83bd6b90de9826d16f3/coverage-7.15.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:39ece820e29e0a2ba34b3ecb3be83c27e997eed8926f2ba6fe7ce7a0bda5843b", size = 255846, upload-time = "2026-08-06T13:48:38.317Z" }, + { url = "https://files.pythonhosted.org/packages/96/e4/38015b2b6d21258713bd17e76b59d033b191efb5703589cffd037dfbca20/coverage-7.15.4-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:f21b56dcace11dfe013014201f577dcd592b2a9b72182d930361b47cf6f73f25", size = 253808, upload-time = "2026-08-06T13:48:39.993Z" }, + { url = "https://files.pythonhosted.org/packages/0b/64/0d515c1e60ee6fbfd1a0e79c07cd87d388a233b7adc37758735677203808/coverage-7.15.4-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:93a3a0b662abcc10c73a47cbc72cd60f63618d6989fb2d1286e50eacd974f303", size = 258081, upload-time = "2026-08-06T13:48:41.971Z" }, + { url = "https://files.pythonhosted.org/packages/91/71/04d9e7a3642146c6351338aef4ef85ab11dbbb54744c13245caba1aad1c0/coverage-7.15.4-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:141fae2cabf5569b782c10afc4c850ce10f618c13f8db54765cba99cc839da1f", size = 253624, upload-time = "2026-08-06T13:48:43.731Z" }, + { url = "https://files.pythonhosted.org/packages/b4/a7/6c28b74c81ebff66987b0e2522ba5cffa3e90b0c33cb6a2eb264d4ee8cf1/coverage-7.15.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:81294c7e6ab30c5f74c0353b11b2fd6320e72d9bee6ac73b357caa8b916323a5", size = 255280, upload-time = "2026-08-06T13:48:45.58Z" }, + { url = "https://files.pythonhosted.org/packages/52/af/bc19996a7014b98d7bbb0f0939453c67074af65784a3aa16a789a07381fa/coverage-7.15.4-cp314-cp314-win32.whl", hash = "sha256:7bbd7d6418e0dab31a206af5203bd43ae36edb8e7fba1940b055d3e9249290d7", size = 224768, upload-time = "2026-08-06T13:48:47.525Z" }, + { url = "https://files.pythonhosted.org/packages/ee/90/219484e476d6e101ba0a444852579e05f5b75c37c611a42ed1190f73ef62/coverage-7.15.4-cp314-cp314-win_amd64.whl", hash = "sha256:f0204ed122758782970526057093f448051a39db9d810d4e344bb87a3546f425", size = 225259, upload-time = "2026-08-06T13:48:49.513Z" }, + { url = "https://files.pythonhosted.org/packages/b7/66/fa77daf4e383e5f776dac62c2409b6af81910ae6fe326bd5170dba74cc63/coverage-7.15.4-cp314-cp314-win_arm64.whl", hash = "sha256:9e71e7bc71c686a123347ae47a0de33a175e797a85bb57b791492adf4eec8ed8", size = 224684, upload-time = "2026-08-06T13:48:51.235Z" }, + { url = "https://files.pythonhosted.org/packages/58/5b/f03bf0ce362bbf3f785fa5219620d00778d4ac6fc9e407734828e9c672f6/coverage-7.15.4-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:7c922735321eef3f87c280a3d39afff6b646723a2880b862cda4ac7a093b8aa8", size = 223338, upload-time = "2026-08-06T13:48:52.896Z" }, + { url = "https://files.pythonhosted.org/packages/0f/76/e77d0ae22501831cc9f92193e8a957a5caa1dd177f90a6d1d9b106242d92/coverage-7.15.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f41c17c4668a655ce96d090d8d5ffdc24ef64b5a02f9753884d08483e8a4a41a", size = 223609, upload-time = "2026-08-06T13:48:54.688Z" }, + { url = "https://files.pythonhosted.org/packages/82/1a/b1f089da8d38ac612fa2dd6dc7f4a1a7657d12f3e261d2996edd3a838d0b/coverage-7.15.4-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:46822e9b6ff1c6a72b518c162c44a8f45a61a1d609c51084bf5b16c023c5037b", size = 264970, upload-time = "2026-08-06T13:48:56.403Z" }, + { url = "https://files.pythonhosted.org/packages/bf/31/e66d98d6e9c7fcc88470f1e234eaf6b1950dc0dfbf797f7282c1c861da24/coverage-7.15.4-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3d6f4955b73b5445271379a59e3792b0d978f42d4a01e0cf7a67d9c33a3bb0a5", size = 267088, upload-time = "2026-08-06T13:48:58.41Z" }, + { url = "https://files.pythonhosted.org/packages/59/a1/ae94eb2c541add426378408379f233591e069040b1e2cdb33df9498a0682/coverage-7.15.4-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3fc9e047706fb4a9abb54f719d3aa643e80e5bb3818182c40aee01ac0f0247ba", size = 269508, upload-time = "2026-08-06T13:49:00.42Z" }, + { url = "https://files.pythonhosted.org/packages/9c/c7/88a10694a1c6a213569766aba9f25847b28155d4ac731b13226db216356d/coverage-7.15.4-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:05e491d4f3165d62d4f5c8fd48dfeabf2ae8f42cbbd484319af33ea851b78982", size = 270629, upload-time = "2026-08-06T13:49:02.234Z" }, + { url = "https://files.pythonhosted.org/packages/b3/34/d8b8232e5e55169933b59aabcef2fedfa4b9d8897361bb80fcbda146505f/coverage-7.15.4-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:226c66e80ec0598d3b9b4874123df167ccca342aca8714f77cac6829688ee09c", size = 264043, upload-time = "2026-08-06T13:49:04.102Z" }, + { url = "https://files.pythonhosted.org/packages/7e/35/58b009dbf8c471c7224716478b9fed4a7e1af15320e1ed41660978504663/coverage-7.15.4-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ac41cc14bebda0dbfb0628036b7f75706935c95bcc07fefe9a0f93614aa60a57", size = 266963, upload-time = "2026-08-06T13:49:05.821Z" }, + { url = "https://files.pythonhosted.org/packages/62/aa/57fbda1b42c892968273c56b6ee9dc0f1310850859230a507bc7873b1f65/coverage-7.15.4-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:8af623e5cd92080acddd02b38f2f406a2c3a0893c38950b211890361448fbf26", size = 264569, upload-time = "2026-08-06T13:49:07.706Z" }, + { url = "https://files.pythonhosted.org/packages/98/8a/360e6e7f24d477b7e889703af0afa878d15b6d4d8d2a822b2835c169a879/coverage-7.15.4-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:07545711d4f0f32852a18f18ad11f76f0109909d09e78b9008b4cfc67e829429", size = 268299, upload-time = "2026-08-06T13:49:09.587Z" }, + { url = "https://files.pythonhosted.org/packages/4e/89/6f701261aee21b6b5fa8f7872229406dc917e125069448292223bf213606/coverage-7.15.4-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:a0865421cfdc53654b342d515e5a233187590882d20b95752150e53f65460017", size = 263413, upload-time = "2026-08-06T13:49:11.604Z" }, + { url = "https://files.pythonhosted.org/packages/3f/0f/6f04036edc260ed425af83e834f627fad48941ce97b50bfe6edd8b6fa623/coverage-7.15.4-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:460115e32ee40566476db5048f9bec1e842c127ad8e6f8be745aad3ac9cbc839", size = 265725, upload-time = "2026-08-06T13:49:13.38Z" }, + { url = "https://files.pythonhosted.org/packages/c4/ce/d19b5d4d5c49a7bfb925fd74310fee7d28bc99520ac3367ccbc54e662518/coverage-7.15.4-cp314-cp314t-win32.whl", hash = "sha256:cbde877ef9dd7baf272b9bfef2b8a25edd45d9170fc326951dd20eb480335e85", size = 225079, upload-time = "2026-08-06T13:49:15.265Z" }, + { url = "https://files.pythonhosted.org/packages/26/bb/7aa1b3b173faee0679037ca950bbbe1247273656697994d8d13f80f8d4b4/coverage-7.15.4-cp314-cp314t-win_amd64.whl", hash = "sha256:3da9e92d1c551fd7563833e9ade686efb0c4b7363ab7681a94283958c950bf5e", size = 225911, upload-time = "2026-08-06T13:49:17.279Z" }, + { url = "https://files.pythonhosted.org/packages/81/1c/4ea9e47426d80038d9222db3c4534cb6021a74b237d3ff97ffd33b6600dd/coverage-7.15.4-cp314-cp314t-win_arm64.whl", hash = "sha256:3a54f5a0d85050c73a38f6793090ee83974531e67fe5e57a1da9bee11398aa5e", size = 225219, upload-time = "2026-08-06T13:49:19.293Z" }, + { url = "https://files.pythonhosted.org/packages/2b/c4/dc5d2ac8f9142e7ec7de66e7bf0591db29d78955a040bd915870d9c0e657/coverage-7.15.4-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:2c9872e4d9dc5d3cf616bf4b382f5a00359305a5be666a3dd0b5cdb4e49597f9", size = 222604, upload-time = "2026-08-06T13:49:21.279Z" }, + { url = "https://files.pythonhosted.org/packages/70/39/33e63df81fe2ee100897451841c821467635923e58e37c6bd4b46dd8106c/coverage-7.15.4-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:e101dbb4b9b72f0cddd8cdc8c9c5b47f456766f5e0ac82dbfb75e5c55409b78a", size = 222944, upload-time = "2026-08-06T13:49:23.187Z" }, + { url = "https://files.pythonhosted.org/packages/99/1f/ef3ffb5557febc75a0d97aa459d0266d7d741110265121cc6d8539343d44/coverage-7.15.4-cp315-cp315-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:7d1abebdb047729e852b9c77a00497dfbeb11eb3a117e037d7dbc3ac8e5f5c54", size = 254050, upload-time = "2026-08-06T13:49:25.008Z" }, + { url = "https://files.pythonhosted.org/packages/6f/f5/1f0f6f77698c3601ca0ae7431e34b24c62ca2f06fecb23b73ed1f651d2be/coverage-7.15.4-cp315-cp315-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:d28a4a899354d0ea6214cc59b4fa19eefbce1b9ff1688ab579acf49e894bd3fb", size = 256967, upload-time = "2026-08-06T13:49:26.896Z" }, + { url = "https://files.pythonhosted.org/packages/03/7a/2ed9bed79925f4367c83c77f66a89e5ca7229c288d2d19ad5f36d1ca0070/coverage-7.15.4-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ffb3c2aacea411cc7e1d27712490c11108e2de1d39019ae32915493a59a8b9ed", size = 258587, upload-time = "2026-08-06T13:49:28.692Z" }, + { url = "https://files.pythonhosted.org/packages/45/8c/fa34044f71b7cc4ecb6da9c2408770959b0591fa9b5fb6fb6bca38f94298/coverage-7.15.4-cp315-cp315-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a9447978a92f405d301123cfd39ff49895490efb769a758fe2734c7f631bf8ce", size = 260785, upload-time = "2026-08-06T13:49:30.472Z" }, + { url = "https://files.pythonhosted.org/packages/4f/54/d5727ce36b4524a7394ab9f5f1df378e1f23affcdab01037dc8655185cc7/coverage-7.15.4-cp315-cp315-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:050467a7983b8e2fe7dd41a78bb30c3e7f8c0b8cafda14b1c46f8b5e3cf2dd3c", size = 254545, upload-time = "2026-08-06T13:49:32.271Z" }, + { url = "https://files.pythonhosted.org/packages/dc/e6/6e3783e576719590194bdffb6dd6d85490801785b7c331e35a245d8cb8b5/coverage-7.15.4-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:d003b7a5708ddad5c206c79607a6b92abb6fc13c57d99d8a4468cc03a2941ced", size = 256682, upload-time = "2026-08-06T13:49:34.089Z" }, + { url = "https://files.pythonhosted.org/packages/dc/f2/bacdbde18b69ed2de424fcf64d9fb0a4913753d4f0eca8bae9daad69f4bd/coverage-7.15.4-cp315-cp315-musllinux_1_2_i686.whl", hash = "sha256:c38efe30fd74e5c19e9433f11fb1f5dc9c6522770971b7c6145bbaa413dc8800", size = 254560, upload-time = "2026-08-06T13:49:36.052Z" }, + { url = "https://files.pythonhosted.org/packages/6c/a3/1fb927196e3477c1b48831169ab58ba08f451ba87ae311ff1de68b26a616/coverage-7.15.4-cp315-cp315-musllinux_1_2_ppc64le.whl", hash = "sha256:1f4f826d70f772ab8b0c052329580d7fe8b8abd191e4ce0c8f81aec6614665d3", size = 258792, upload-time = "2026-08-06T13:49:38.01Z" }, + { url = "https://files.pythonhosted.org/packages/41/58/30d4c149c69053de0edfe325614c1d28d508f62b1783e0e4a234d2e49136/coverage-7.15.4-cp315-cp315-musllinux_1_2_riscv64.whl", hash = "sha256:4a4bf917c9953f57c957be31c1cd504e3bd2f34d4a352b9d391a3025336f6768", size = 253968, upload-time = "2026-08-06T13:49:39.934Z" }, + { url = "https://files.pythonhosted.org/packages/89/e4/77f639371b918aad30dda4051f95404b43578f7f2e2f87ba73e02ed1ff37/coverage-7.15.4-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:1c9bf40ebef178a45192c75c4964760bb261b0e6ad725da5fc4c93f674f19753", size = 255893, upload-time = "2026-08-06T13:49:41.825Z" }, + { url = "https://files.pythonhosted.org/packages/5c/62/13be29b3ddab35f14c87967a4820a05106d2a3eccb4fa4ff550bf30b75e0/coverage-7.15.4-cp315-cp315-win32.whl", hash = "sha256:43619d04c3671792d2c4706ae8bf45e265dc87bbd4078189ef8b847ea1e74be2", size = 224768, upload-time = "2026-08-06T13:49:44.08Z" }, + { url = "https://files.pythonhosted.org/packages/a1/70/af0c6be0f964af6954f6b74bc109b0dbca02824696d2520fb17fe1ab06e3/coverage-7.15.4-cp315-cp315-win_amd64.whl", hash = "sha256:be619439dbcd31a2eab10b32de9fff62c26ed4bab69dc32b8363fdaaa0882809", size = 225242, upload-time = "2026-08-06T13:49:45.899Z" }, + { url = "https://files.pythonhosted.org/packages/4f/2d/f3bd3aab899fc9efc18b53133ee68f5f98574ef480649b23e12962226387/coverage-7.15.4-cp315-cp315-win_arm64.whl", hash = "sha256:def597967dafc2e8d97c9097ea453c464e0bb8ed38f193a43070f10dc623bb6d", size = 224674, upload-time = "2026-08-06T13:49:48.322Z" }, + { url = "https://files.pythonhosted.org/packages/f5/ca/f69251cd63eabc6438321aea22148754cce758a26bde07dd490e3fe7cfc5/coverage-7.15.4-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:c7dbc748ac8a1e3e59a2b28bea47675e6e778081dbbf081bde0d75def2fcbe1d", size = 223333, upload-time = "2026-08-06T13:49:50.293Z" }, + { url = "https://files.pythonhosted.org/packages/a7/a7/037b53b2885b0d8447064432491a4d5a1014cd9f97a594d53acd0c04541a/coverage-7.15.4-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:2413074a5ecbb61a01a7888fc72db0ca324d13588c5b38bc0dd8564cdcdfea26", size = 223630, upload-time = "2026-08-06T13:49:52.637Z" }, + { url = "https://files.pythonhosted.org/packages/80/4f/152b8a4779ae90da11bb24f7467df8a59f0be48a5c52acb856325ca48289/coverage-7.15.4-cp315-cp315t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:4e6f6f632b7b2f714bf7a1346e8f97b650ee71f3c298aaad42a2ab60f0f07645", size = 264489, upload-time = "2026-08-06T13:49:54.52Z" }, + { url = "https://files.pythonhosted.org/packages/10/2d/84b4b9e0e1dd6528a51920ff7031f35b789382e467a28ec6a5a578cb8812/coverage-7.15.4-cp315-cp315t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8df457da2249d3c75ca2e5e835d59c725abfe92d27fdff6cd99eed85b51d5e9a", size = 267567, upload-time = "2026-08-06T13:49:56.721Z" }, + { url = "https://files.pythonhosted.org/packages/53/fc/ba01cc25299f9f8a2c8b02d3b28c53f3543d9fbfbe4e74fa2760b48f163e/coverage-7.15.4-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:050f66a08805acb5b8a23c6d4a517b1ecf82c08e81ed0e4bd727df065e5c6624", size = 270123, upload-time = "2026-08-06T13:49:58.736Z" }, + { url = "https://files.pythonhosted.org/packages/cf/d0/db2647cbf40b14f8c308f94ff7bf89c06d564e59f396906edf50086ec788/coverage-7.15.4-cp315-cp315t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1587fb771d1ccceef708fdde1e5af8c7ed24b486b61d13a321acb7d8145390aa", size = 271107, upload-time = "2026-08-06T13:50:00.811Z" }, + { url = "https://files.pythonhosted.org/packages/70/ff/4d2d17924552c458bb4f77dd631f0e3bc92fbbdf2d2d916cd4b33bbfd5b1/coverage-7.15.4-cp315-cp315t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:8b4f1c3a69ca580f3fbd6b2046915f536d7f586874f25c1bb23add2a3c88d50f", size = 264955, upload-time = "2026-08-06T13:50:03.023Z" }, + { url = "https://files.pythonhosted.org/packages/ee/de/dc010c7a3691f396d93bbc26bfcafa1c2a3a351cd520470f15faf5795bd5/coverage-7.15.4-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:ffb58d7eff5b7f6ecc6fa21d6288ab7f968a212cb67d682c269c09b9eba3b66f", size = 267949, upload-time = "2026-08-06T13:50:05.557Z" }, + { url = "https://files.pythonhosted.org/packages/78/ea/dc96a11375e83c045c2f7c61fb6918277cfe9401db7c0f7b1d111a84b2e5/coverage-7.15.4-cp315-cp315t-musllinux_1_2_i686.whl", hash = "sha256:d9df165544774574ee004b953023d1bebada1894a80b1052a43d798b0f676e67", size = 264421, upload-time = "2026-08-06T13:50:07.612Z" }, + { url = "https://files.pythonhosted.org/packages/c8/86/b77131a0f9503ce461cd577076147d7a9040f0c5dda772686f729e2cc9cb/coverage-7.15.4-cp315-cp315t-musllinux_1_2_ppc64le.whl", hash = "sha256:f9de0a24a4079b53e523b5c5e2c5945ec251ab486652659955187cf255a259bc", size = 269121, upload-time = "2026-08-06T13:50:09.58Z" }, + { url = "https://files.pythonhosted.org/packages/24/24/944bc35007862955e7ebf05754e645419dcf5d7526c52735cfa2715e8ebf/coverage-7.15.4-cp315-cp315t-musllinux_1_2_riscv64.whl", hash = "sha256:150089274bdc9f940628552cb92844e0223c987f1902ab8efe9f45a2ec758d88", size = 264565, upload-time = "2026-08-06T13:50:11.722Z" }, + { url = "https://files.pythonhosted.org/packages/c7/cc/a3bb9f93e7e740659163e2ea584f8196ddcd2c456a5dbe15f6c50105fec1/coverage-7.15.4-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:a58a94fed5da6997d258e8f7668c1e195fbd04a691d781b7558f1e468f9e68bc", size = 266522, upload-time = "2026-08-06T13:50:13.786Z" }, + { url = "https://files.pythonhosted.org/packages/49/dd/e0e40f3560d878d888c580698ff5ad1179f5e1c3ac949684ef66b41a3817/coverage-7.15.4-cp315-cp315t-win32.whl", hash = "sha256:ebd5a6d8466ff30836572f3ba2cae8a5e8f85029b1c6d5e2ed338dc472a5166a", size = 225068, upload-time = "2026-08-06T13:50:15.825Z" }, + { url = "https://files.pythonhosted.org/packages/c6/7e/37732ea80eebc30e976e4cdab15c190bc42d96959a42e38ddf6f8c60468f/coverage-7.15.4-cp315-cp315t-win_amd64.whl", hash = "sha256:288bde2a2d7ab6b6c2d7252fcde8b524387f2d970bdba9658fc6f8bbcaef0f9b", size = 225895, upload-time = "2026-08-06T13:50:17.928Z" }, + { url = "https://files.pythonhosted.org/packages/c6/08/1e00f7923eaaba45fb3d51dd794125fc766304b1df264f3a9c6557bfb30e/coverage-7.15.4-cp315-cp315t-win_arm64.whl", hash = "sha256:68be5e1de60ff13c9095bbec0e5a7fa45b33b101752215b91345ea1f61c4a278", size = 225213, upload-time = "2026-08-06T13:50:19.981Z" }, + { url = "https://files.pythonhosted.org/packages/b4/d9/e70c286c979378f061d8266e279b686ab0b0b688e1fe0af864684f23a77d/coverage-7.15.4-py3-none-any.whl", hash = "sha256:964730a1e9de9c0cf11be6a1a3c79ce419c34882842abd256086ba4698705e84", size = 214332, upload-time = "2026-08-06T13:50:22.192Z" }, +] + [[package]] name = "filelock" version = "3.25.2" @@ -290,6 +374,7 @@ dependencies = [ [package.dev-dependencies] dev = [ { name = "pytest" }, + { name = "pytest-cov" }, { name = "ruff" }, ] @@ -306,6 +391,7 @@ requires-dist = [ [package.metadata.requires-dev] dev = [ { name = "pytest", specifier = ">=8.3,<9" }, + { name = "pytest-cov", specifier = ">=7.1,<8" }, { name = "ruff", specifier = ">=0.11,<0.12" }, ] @@ -411,6 +497,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a8/a4/20da314d277121d6534b3a980b29035dcd51e6744bd79075a6ce8fa4eb8d/pytest-8.4.2-py3-none-any.whl", hash = "sha256:872f880de3fc3a5bdc88a11b39c9710c3497a547cfa9320bc3c5e62fbf272e79", size = 365750, upload-time = "2025-09-04T14:34:20.226Z" }, ] +[[package]] +name = "pytest-cov" +version = "7.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "coverage" }, + { name = "pluggy" }, + { name = "pytest" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/51/a849f96e117386044471c8ec2bd6cfebacda285da9525c9106aeb28da671/pytest_cov-7.1.0.tar.gz", hash = "sha256:30674f2b5f6351aa09702a9c8c364f6a01c27aae0c1366ae8016160d1efc56b2", size = 55592, upload-time = "2026-03-21T20:11:16.284Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/7a/d968e294073affff457b041c2be9868a40c1c71f4a35fcc1e45e5493067b/pytest_cov-7.1.0-py3-none-any.whl", hash = "sha256:a0461110b7865f9a271aa1b51e516c9a95de9d696734a2f71e3e78f46e1d4678", size = 22876, upload-time = "2026-03-21T20:11:14.438Z" }, +] + [[package]] name = "pyyaml" version = "6.0.3" From fe49c7911f726ca24168a954cbe82e4adf6ab9b8 Mon Sep 17 00:00:00 2001 From: BP <11394934+benjipeng@users.noreply.github.com> Date: Fri, 14 Aug 2026 18:17:50 -0400 Subject: [PATCH 2/2] Make Gemma tokenizer tests artifact-free --- tests/unit/test_gemma3_text_tokenizer.py | 71 ++++++++++++++++-------- 1 file changed, 47 insertions(+), 24 deletions(-) diff --git a/tests/unit/test_gemma3_text_tokenizer.py b/tests/unit/test_gemma3_text_tokenizer.py index 671ea9c..a6b91a8 100644 --- a/tests/unit/test_gemma3_text_tokenizer.py +++ b/tests/unit/test_gemma3_text_tokenizer.py @@ -1,33 +1,56 @@ -"""Tokenizer behavior tests for the plain-text Gemma wrapper used by DramaBox. - -Tests run against the real `models/gemma_3_12b_it_backbone/mlx-4bit/tokenizer.json`; they -skip if the model directory is not present locally. -""" +"""Tokenizer behavior tests for the plain-text Gemma wrapper used by DramaBox.""" from __future__ import annotations +import json from pathlib import Path import mlx.core as mx import pytest +from tokenizers import Tokenizer +from tokenizers.models import WordLevel +from tokenizers.pre_tokenizers import Whitespace from mlx_speech.models.gemma3_text import LTXVGemmaTokenizer -GEMMA_DIR = Path("models/gemma_3_12b_it_backbone/mlx-4bit") - -pytestmark = pytest.mark.skipif( - not (GEMMA_DIR / "tokenizer.json").is_file(), - reason="Gemma tokenizer.json not present; skipping (requires local model dir)", -) - -def test_from_dir_loads(): - tok = LTXVGemmaTokenizer.from_dir(GEMMA_DIR) +@pytest.fixture() +def gemma_dir(tmp_path: Path) -> Path: + tokenizer = Tokenizer( + WordLevel( + { + "": 0, + "": 1, + "": 2, + "hello": 3, + "world": 4, + "word": 5, + "short": 6, + "a": 7, + "bit": 8, + "longer": 9, + "sentence": 10, + "here": 11, + }, + unk_token="", + ) + ) + tokenizer.pre_tokenizer = Whitespace() + tokenizer.save(str(tmp_path / "tokenizer.json")) + (tmp_path / "special_tokens_map.json").write_text( + json.dumps({"pad_token": "", "eos_token": ""}), + encoding="utf-8", + ) + return tmp_path + + +def test_from_dir_loads(gemma_dir: Path): + tok = LTXVGemmaTokenizer.from_dir(gemma_dir) assert tok.pad_token_id is not None -def test_encode_returns_left_padded_shape(): - tok = LTXVGemmaTokenizer.from_dir(GEMMA_DIR) +def test_encode_returns_left_padded_shape(gemma_dir: Path): + tok = LTXVGemmaTokenizer.from_dir(gemma_dir) input_ids, attention_mask = tok.encode("hello", max_length=64) assert input_ids.shape == (1, 64) assert attention_mask.shape == (1, 64) @@ -39,8 +62,8 @@ def test_encode_returns_left_padded_shape(): assert mask_list[0] == 0 -def test_encode_long_input_truncates_from_right(): - tok = LTXVGemmaTokenizer.from_dir(GEMMA_DIR) +def test_encode_long_input_truncates_from_right(gemma_dir: Path): + tok = LTXVGemmaTokenizer.from_dir(gemma_dir) long_text = "word " * 5000 # certainly more than max_length=64 tokens input_ids, attention_mask = tok.encode(long_text, max_length=64) assert input_ids.shape == (1, 64) @@ -48,16 +71,16 @@ def test_encode_long_input_truncates_from_right(): assert attention_mask.sum().item() == 64 -def test_encode_strips_whitespace(): +def test_encode_strips_whitespace(gemma_dir: Path): """DramaBox tokenizer strips leading/trailing whitespace before encoding.""" - tok = LTXVGemmaTokenizer.from_dir(GEMMA_DIR) + tok = LTXVGemmaTokenizer.from_dir(gemma_dir) ids_a, _ = tok.encode("hello world", max_length=32) ids_b, _ = tok.encode(" hello world ", max_length=32) assert (ids_a == ids_b).all().item() -def test_encode_batch_uniform_length(): - tok = LTXVGemmaTokenizer.from_dir(GEMMA_DIR) +def test_encode_batch_uniform_length(gemma_dir: Path): + tok = LTXVGemmaTokenizer.from_dir(gemma_dir) ids, mask = tok.encode_batch(["short", "a bit longer sentence here"], max_length=32) assert ids.shape == (2, 32) assert mask.shape == (2, 32) @@ -65,8 +88,8 @@ def test_encode_batch_uniform_length(): assert mask[0].sum().item() < mask[1].sum().item() -def test_encode_dtype_is_int32(): - tok = LTXVGemmaTokenizer.from_dir(GEMMA_DIR) +def test_encode_dtype_is_int32(gemma_dir: Path): + tok = LTXVGemmaTokenizer.from_dir(gemma_dir) input_ids, attention_mask = tok.encode("hello", max_length=16) assert input_ids.dtype == mx.int32 assert attention_mask.dtype == mx.int32