diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index b02ba77..5a581bd 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -7,6 +7,34 @@ on: branches: ["**"] jobs: + dependency-resolution: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: "3.12" + + - name: Resolve dependencies without GitHub credentials + env: + GIT_TERMINAL_PROMPT: "0" + run: | + python -m venv /tmp/pathbench-clean + /tmp/pathbench-clean/bin/python -m pip install --upgrade pip + /tmp/pathbench-clean/bin/python -m pip install \ + --dry-run --ignore-installed --report /tmp/install-report.json . + /tmp/pathbench-clean/bin/python - <<'PY' + import json + from pathlib import Path + + report = json.loads(Path("/tmp/install-report.json").read_text()) + urls = [item["download_info"]["url"] for item in report["install"]] + vcs_urls = [url for url in urls if url.startswith("git+") or "github.com" in url] + assert not vcs_urls, f"VCS/GitHub dependencies remain: {vcs_urls}" + PY + test: runs-on: ubuntu-latest strategy: diff --git a/README.md b/README.md index 2382dc2..842686f 100644 --- a/README.md +++ b/README.md @@ -84,6 +84,25 @@ score = evaluator.score("utt1", "/path/to/audio.wav", transcription="the cat sat print(f"ArtP score: {score}") ``` +ArtP is reference-based: it force-aligns the phonemes in the supplied transcription +with the audio, so the transcription and its language are required. DArtP is +reference-free: a language-specific ASR model and n-gram language model first +produce a transcription, which is then scored by the same phonetic model: + +```python +from pathbench import ArtPDoubleASREvaluator + +evaluator = ArtPDoubleASREvaluator(language="en-us") +score = evaluator.score("utt1", "/path/to/audio.wav") +print(f"DArtP score: {score}") +``` + +DArtP currently supports `en`/`en-us`, `es`, `nl`, `it`, and `cmn`. It also +requires the corresponding file from the [n-gram model download](#n-gram-models) +to be placed in `lms/`; without it, scoring returns `None`. Run commands from +the repository root because DArtP resolves `lms/` relative to the working +directory. + ### I want to contribute a new predictor to this repository, how do I do that? See [CONTRIBUTING.md](CONTRIBUTING.md) for a step-by-step guide. @@ -120,18 +139,85 @@ Results are written to the `results_11/` directory as timestamped text files con ## Installation -We are continously trying to make the installation easier for your use case. +We are continuously trying to make the installation easier for your use case. + +### Complete GPU installation (install missing components only) + +The following Ubuntu procedure is safe to re-run: it installs only absent apt +packages, builds the pinned `espeak-ng` unless its commit marker matches, +clones PathBench only when the checkout is absent, and lets the GPU helper reuse +an existing virtual environment and matching Python packages. + +```bash +# 1. Install missing build prerequisites. +packages=(git python3 python3-venv build-essential cmake libfftw3-dev liblapack-dev) +missing=() +for package in "${packages[@]}"; do + dpkg-query -W -f='${Status}' "$package" 2>/dev/null | grep -q "ok installed" \ + || missing+=("$package") +done +if ((${#missing[@]})); then + sudo apt-get update -qq + sudo apt-get install -y "${missing[@]}" +fi + +# 2. Build the reproducible phonemizer backend unless the pinned commit is installed. +espeak_ng_commit=2ea41210 +espeak_ng_marker=/usr/local/share/pathbench/espeak-ng-commit +if ! command -v espeak-ng >/dev/null \ + || [[ ! -r "$espeak_ng_marker" ]] \ + || [[ "$(cat "$espeak_ng_marker")" != "$espeak_ng_commit" ]]; then + if { test -d /tmp/espeak-ng/.git \ + || git clone https://github.com/espeak-ng/espeak-ng.git /tmp/espeak-ng; } \ + && git -C /tmp/espeak-ng fetch origin "$espeak_ng_commit" \ + && git -C /tmp/espeak-ng checkout --detach "$espeak_ng_commit" \ + && cmake -S /tmp/espeak-ng -B /tmp/espeak-ng/build \ + -DUSE_ASYNC=OFF -DBUILD_SHARED_LIBS=ON \ + && cmake --build /tmp/espeak-ng/build -j"$(nproc)" \ + && sudo cmake --install /tmp/espeak-ng/build \ + && sudo ldconfig \ + && sudo install -d "$(dirname "$espeak_ng_marker")"; then + printf '%s\n' "$espeak_ng_commit" \ + | sudo tee "$espeak_ng_marker" >/dev/null + else + echo "Failed to install pinned espeak-ng; commit marker was not written." >&2 + exit 1 + fi +fi + +# 3. Reuse the current checkout, or clone to a stable absolute destination. +if pathbench_root=$(git rev-parse --show-toplevel 2>/dev/null) \ + && test -f "$pathbench_root/tools/test_gpu_predictors.py"; then + : # Already anywhere inside a PathBench checkout. +else + pathbench_root=${PATHBENCH_ROOT:-"$PWD/pathbench"} + test -d "$pathbench_root/.git" \ + || git clone https://github.com/karkirowle/pathbench.git "$pathbench_root" +fi +cd "$pathbench_root" +python3 tools/test_gpu_predictors.py --download-language-model --cuda-version 12.4 +``` + +This procedure assumes that a working NVIDIA driver is already installed; +`nvidia-smi` must list the assigned GPU. Driver installation is host- and +cloud-specific and is deliberately not attempted by the script. If you have the opportunity to start from a clean AWS/GCE instance, please do so and follow the make installation. If you are working on a highly restricted HPC cluster, I would recommend starting from the singularity container [provided](https://github.com/karkirowle/pathbench/releases/download/v0.1.0/pathbench.sif). -Package installation is the recommended pathway when you are trying to incorporate into your existing stuff. In this case, you are kind of your own figuring out +Package installation is the recommended pathway when incorporating PathBench +into an existing environment. In that case, you are responsible for resolving dependency conflicts. -### Package installation +All Python runtime dependencies are available from package indexes rather than +VCS URLs. In particular, `phonemizer-fork==3.3.2` (which installs the +`phonemizer` import package) and `pyctcdecode==0.5.0` use versioned PyPI +releases, so installing PathBench does not require GitHub credentials. The +system-level espeak-ng revision below remains separately pinned because its +language-specific IPA output is part of the metric definition. -PathBench cannot be published to PyPI because it depends on Git-hosted forks of `phonemizer` and `pyctcdecode`. +### Package installation **System dependencies** (not installable via pip — must be installed separately): - `espeak-ng` at commit [`2ea41210`](https://github.com/espeak-ng/espeak-ng/commit/2ea41210) (post-1.52.0) — required by the phonemizer for grapheme-to-phoneme conversion. The exact commit matters: different espeak-ng versions produce different IPA symbols for some languages (e.g. Italian `ɾ` vs `r`), which affects phoneme-based metrics (PER, dPER, ArtP). Build from source: @@ -141,7 +227,8 @@ PathBench cannot be published to PyPI because it depends on Git-hosted forks of cmake -B build -DUSE_ASYNC=OFF -DBUILD_SHARED_LIBS=ON cmake --build build -j$(nproc) && sudo cmake --install build ``` -- PyTorch with CUDA support — install following [pytorch.org](https://pytorch.org/get-started/locally/) *before* installing pathbench +- PyTorch — install the CPU or CUDA build appropriate for your system by following + [pytorch.org](https://pytorch.org/get-started/locally/) *before* installing PathBench. **Option A — Install from a GitHub Release:** ```bash @@ -160,7 +247,10 @@ pip install "pathbench[scripts] @ git+https://github.com/karkirowle/pathbench.gi ### Make installation -The `make` installation route assumes the default setup of a standard Ubuntu 22.04 image (`ubuntu-2204-jammy`). +The `make` installation route assumes the default setup of a standard Ubuntu +22.04 image (`ubuntu-2204-jammy`) and Python 3.10–3.12. It creates +`tools/venv`. The default is a CPU-only PyTorch installation, which works for +inference and tests but is slower than a supported GPU. ```bash sudo apt-get update -qq @@ -171,12 +261,120 @@ cd /tmp/espeak-ng && git checkout 2ea41210 cmake -B build -DUSE_ASYNC=OFF -DBUILD_SHARED_LIBS=ON cmake --build build -j$(nproc) && sudo cmake --install build && sudo ldconfig cd - -git clone git@github.com:karkirowle/pathbench.git +git clone https://github.com/karkirowle/pathbench.git cd pathbench/tools && make cd .. source tools/venv/bin/activate ``` +For a CUDA build, select a wheel index supported by the pinned PyTorch version. +For example, PyTorch 2.6.0 provides CUDA 12.4 wheels: + +```bash +cd pathbench/tools +make CUDA_VERSION=12.4 +``` + +You can select a particular interpreter with, for example, +`make PYTHON=python3.12`. Re-running `make` resumes after completed stages; +run `make clean` first to rebuild the environment with a different Python, +PyTorch, or CUDA selection. + +### GPU installation and predictor smoke test + +After installing the pinned `espeak-ng` build above, systems with an NVIDIA GPU +and driver can use the helper script to create a separate CUDA environment and +run the focused ArtP and DArtP tests: + +```bash +python tools/test_gpu_predictors.py --download-language-model --cuda-version 12.4 +``` + +The script checks for `nvidia-smi` and `espeak-ng`, creates +`tools/gpu_venv`, installs the CUDA 12.4 builds of PyTorch and torchaudio 2.6.0, +installs PathBench and its test dependencies, verifies that PyTorch can access +the GPU, and runs both predictor tests. It can be invoked from any directory. +The first run downloads the Python packages and model checkpoints and therefore +requires network access and several gigabytes of free disk space. + +Override its defaults with command-line options (or the corresponding +`PYTHON`, `PATHBENCH_CUDA_VERSION`, `PYTORCH_VERSION`, and `VENV` environment variables) +when needed. The selected CUDA wheel must exist for the selected PyTorch release: + +```bash +python tools/test_gpu_predictors.py --python python3.11 --cuda-version 12.6 \ + --pytorch-version 2.6.0 --venv /path/to/pathbench-gpu-venv +``` + +The script requires an NVIDIA driver compatible with the chosen CUDA wheel; +installing the wheel does not install a host GPU driver or the CUDA toolkit. +Language-model download is deliberately opt-in. With +`--download-language-model`, the helper uses HTTP range requests against the +immutable [35 GB Zenodo archive](https://zenodo.org/api/records/18738598/files/lms.zip/content) +to retrieve only the compressed English member: **8,582,666,912 bytes** +(approximately 8.0 GiB). It streams the raw DEFLATE data into a temporary file, +producing a **14,600,342,241-byte** model (approximately 13.6 GiB), and checks +the member metadata, CRC-32, expanded-model SHA-256 +(`d786eec55174c696c0bf3327928ff496684f482194ba3c6ebdf4311acb823d00`), +and KenLM readability before atomically installing it. The server or any proxy +must support standards-compliant byte ranges (HTTP 206 and `Content-Range`); +the helper refuses an HTTP 200 response rather than accidentally downloading +the complete archive. Allow space for the installed model, its temporary +expanded copy, and at least a 1 GiB safety margin. The record is openly +accessible and licensed **CC BY 4.0**, which permits automatic download and +redistribution with attribution. Without the option, a missing model still +causes DArtP to be reported as skipped. ArtP does not need the model. Custom +standalone file/archive mirrors remain supported but must be supplied together with +`--language-model-sha256`; the project-specific environment equivalents are +`PATHBENCH_LANGUAGE_MODEL_URL`, `PATHBENCH_LANGUAGE_MODEL_SHA256`, and +`PATHBENCH_LANGUAGE_MODEL_CACHE`. The built-in range mode retains no compressed +archive and treats the verified model under `lms/` as its cache; the cache +directory applies to custom downloads. + +The built-in Zenodo artifact identity used by the live integration is: + +| Field | Published value | +| --- | --- | +| Archive (`lms.zip`) size | `35017940434` bytes | +| Member | `lms/wiki_en_token.arpa.bin` | +| Decompressed member size | `14600342241` bytes | +| Compressed member size | `8582666912` bytes | +| Member CRC-32 | `5afb90ef` | +| Decompressed member SHA-256 | `d786eec55174c696c0bf3327928ff496684f482194ba3c6ebdf4311acb823d00` | + +The SHA-256 value is specifically the digest of the fully transferred and +decompressed `lms/wiki_en_token.arpa.bin` member. It is **not** a digest of +`lms.zip` or of the member's compressed DEFLATE stream. + +For both tests together, allow **at least 12 GB of system RAM and 8 GB of GPU +VRAM**; **16 GB system RAM and 12–16 GB VRAM are recommended** to leave room +for both wav2vec2 models, the decoder, and transient activations. Any NVIDIA +CUDA GPU supported by the selected PyTorch wheel is acceptable; a T4 (16 GB), +L4 (24 GB), A10/A10G (24 GB), V100 (16/32 GB), or A100 works. Smaller 8 GB +cards may require closing other GPU processes and can run out of memory on +long audio. AMD ROCm GPUs, Apple GPUs, and CPU-only runtimes do not satisfy +this CUDA smoke test. + +Google Colab GPUs can be used. Select a GPU runtime and confirm that +`nvidia-smi` works; the commonly assigned T4 and higher-memory L4/A100 options +meet the recommendation. Colab does not guarantee a particular GPU, RAM +amount, availability, or uninterrupted runtime, and its temporary filesystem +means the environment and downloaded checkpoints may need to be recreated in +a later session. If Colab assigns a smaller GPU or low-RAM runtime, inspect +`nvidia-smi` and available system memory before running the helper. +On a Python 3.12 T4 runtime, use PyTorch 2.6.0's CUDA 12.4 wheels: + +```bash +python tools/test_gpu_predictors.py --download-language-model --cuda-version 12.4 +``` + +A successful run ends with `2 passed`. A second invocation reuses both +`tools/gpu_venv` and the verified model cache rather than downloading them +again. Colab's local disk is ephemeral, however, so the installed 13.6 GiB +model and environment are lost when its runtime is recycled. Persist the +checkout itself if reuse across sessions is important (`PATHBENCH_LANGUAGE_MODEL_CACHE` +only controls custom standalone downloads). + **Without sudo access:** A containerised environment such as Docker is recommended. ## Downloads @@ -213,7 +411,18 @@ find /path/to/your/datasets/easycall/EasyCall/m13 -name "m13 _*" -exec bash -c ' ### N-gram models -The n-gram models required for DArtP and ArtP are included in the [Oral Cancer - YouTube](https://zenodo.org/records/18738598) download. +The n-gram models required by DArtP are included in the +[Oral Cancer - YouTube](https://zenodo.org/records/18738598) download. ArtP does +not require an n-gram model. Create `lms/` at the repository root and copy the +models there with these exact names: + +| Language | Filename | +| --- | --- | +| English | `wiki_en_token.arpa` or `wiki_en_token.arpa.bin` | +| Dutch | `wiki_nl_token.arpa` or `wiki_nl_token.arpa.bin` | +| Spanish | `wiki_es_token.arpa.bin` | +| Italian | `wiki_it_token.arpa.bin` | +| Mandarin Chinese | `wiki_zh_token.arpa` or `wiki_zh_token.arpa.bin` | ## Testing @@ -228,6 +437,19 @@ python -m pytest tests/test_evaluators.py::TestEvaluatorMethods -v All tests should pass. If all evaluator tests fail simultaneously, the reference audio file in `tests/data/test_audio.wav` may be corrupted — the `test_audio_integrity` test will confirm this. +To test only ArtP and DArtP after downloading the English n-gram model, run: + +```bash +python -m pytest \ + tests/test_evaluators.py::TestEvaluatorMethods::test_articulatory_precision \ + tests/test_evaluators.py::TestEvaluatorMethods::test_artp_double_asr -v +``` + +The first run downloads the Hugging Face checkpoints used by the phonetic and +English ASR models and therefore requires network access and several gigabytes +of free disk space. The DArtP test is reported as skipped, rather than failed, +when neither English n-gram filename listed above exists. + > **Note:** During the NAD evaluator tests you will see a `Wav2Vec2Model LOAD REPORT` table listing several keys (e.g. `project_q`, `quantizer`) as **UNEXPECTED**. These warnings are harmless — the keys belong to pre-training heads that are not needed for feature extraction and can be safely ignored. ### Dataset integrity @@ -280,4 +502,3 @@ This work is partly financed by the Dutch Research Council (NWO) under project n ## Author Bence Mark Halpern, Nagoya University - diff --git a/docs/installation.rst b/docs/installation.rst index c855f59..b3bfe30 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -39,5 +39,8 @@ Without sudo access, a containerised environment such as Docker is recommended. .. note:: - PathBench cannot be published to PyPI because it depends on Git-hosted forks - of ``phonemizer`` and ``pyctcdecode``. + PathBench's Python dependencies are available as versioned package-index + releases, including ``phonemizer-fork==3.3.2`` and + ``pyctcdecode==0.5.0``. Installing them does not require GitHub credentials. + The ``espeak-ng`` shared library remains a separate system dependency; use + the revision documented in the project README for reproducible IPA output. diff --git a/pyproject.toml b/pyproject.toml index 641655f..83ec385 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -50,8 +50,10 @@ dependencies = [ "transformers", "dtw-python", "jiwer", - "phonemizer-fork @ git+https://github.com/thewh1teagle/phonemizer-fork.git", - "pyctcdecode @ git+https://github.com/kensho-technologies/pyctcdecode.git", + # PyPI release of the fork; it continues to expose the ``phonemizer`` API. + "phonemizer-fork==3.3.2", + # Keep every runtime dependency installable without VCS/GitHub access. + "pyctcdecode==0.5.0", "praat-parselmouth>=0.4.4", "scikit-learn", ] diff --git a/tests/test_gpu_predictors_tool.py b/tests/test_gpu_predictors_tool.py new file mode 100644 index 0000000..b333c78 --- /dev/null +++ b/tests/test_gpu_predictors_tool.py @@ -0,0 +1,403 @@ +"""Network-free tests for the opt-in GPU smoke-test model downloader.""" + +from __future__ import annotations + +import hashlib +import importlib.util +import io +import os +from pathlib import Path +import sys +import urllib.error +import zipfile +import zlib + +import pytest + + +SPEC = importlib.util.spec_from_file_location( + "test_gpu_predictors", Path(__file__).parents[1] / "tools/test_gpu_predictors.py" +) +assert SPEC and SPEC.loader +tool = importlib.util.module_from_spec(SPEC) +sys.modules[SPEC.name] = tool +SPEC.loader.exec_module(tool) + + +def test_builtin_language_model_checksum_matches_published_value(): + assert tool.LANGUAGE_MODEL_SHA256 == ( + "d786eec55174c696c0bf3327928ff496684f482194ba3c6ebdf4311acb823d00" + ) + + +class Response(io.BytesIO): + status = 206 + + def __init__(self, value=b"", headers=None, status=206): + super().__init__(value) + self.headers = headers or {} + self.status = status + + def getcode(self): + return self.status + + def __enter__(self): + return self + + def __exit__(self, *_args): + self.close() + + +def zip_bytes(files: dict[str, bytes]) -> bytes: + output = io.BytesIO() + with zipfile.ZipFile(output, "w") as archive: + for name, contents in files.items(): + archive.writestr(name, contents) + return output.getvalue() + + +def configure(monkeypatch, tmp_path: Path, payload: bytes): + calls = [] + monkeypatch.setattr(tool, "REPO_ROOT", tmp_path / "repo") + monkeypatch.setattr( + tool.urllib.request, "urlopen", + lambda *_args, **_kwargs: calls.append(True) or Response(payload), + ) + return calls, hashlib.sha256(payload).hexdigest() + + +def test_successful_download_prefers_binary_and_reuses_cache(monkeypatch, tmp_path): + payload = zip_bytes({"models/wiki_en_token.arpa": b"text", "wiki_en_token.arpa.bin": b"binary"}) + calls, digest = configure(monkeypatch, tmp_path, payload) + first = tool.install_language_model("https://example/model.zip", digest, tmp_path / "cache") + second = tool.install_language_model("https://example/model.zip", digest, tmp_path / "cache") + assert first.read_bytes() == second.read_bytes() == b"binary" + assert len(calls) == 1 + + +def test_installing_arpa_removes_stale_preferred_binary(monkeypatch, tmp_path): + payload = zip_bytes({"models/wiki_en_token.arpa": b"replacement arpa"}) + _calls, digest = configure(monkeypatch, tmp_path, payload) + stale_binary = tmp_path / "repo" / "lms" / "wiki_en_token.arpa.bin" + stale_binary.parent.mkdir(parents=True) + stale_binary.write_bytes(b"stale binary") + + installed = tool.install_language_model( + "https://example/model.zip", digest, tmp_path / "cache", + ) + + assert installed.name == "wiki_en_token.arpa" + assert installed.read_bytes() == b"replacement arpa" + assert not stale_binary.exists() + + +def test_checksum_mismatch_removes_partial(monkeypatch, tmp_path): + calls, _ = configure(monkeypatch, tmp_path, b"not expected") + with pytest.raises(RuntimeError, match="SHA-256 mismatch"): + tool.install_language_model("https://example/model", "0" * 64, tmp_path / "cache") + assert calls and not list((tmp_path / "cache").iterdir()) + + +def test_interrupted_download_removes_partial(monkeypatch, tmp_path): + class Interrupted(Response): + def read(self, *_args): + raise OSError("connection reset") + + monkeypatch.setattr(tool.urllib.request, "urlopen", lambda *_a, **_k: Interrupted(b"x")) + with pytest.raises(RuntimeError, match="download failed"): + tool.install_language_model("https://example/model", "0" * 64, tmp_path / "cache") + assert not list((tmp_path / "cache").iterdir()) + + +def test_http_failure(monkeypatch, tmp_path): + def fail(*_args, **_kwargs): + raise urllib.error.HTTPError("url", 503, "unavailable", {}, None) + + monkeypatch.setattr(tool.urllib.request, "urlopen", fail) + with pytest.raises(RuntimeError, match="503"): + tool.install_language_model("https://example/model", "0" * 64, tmp_path / "cache") + + +def test_safe_archive_extracts_only_model(tmp_path): + archive = tmp_path / "models.zip" + archive.write_bytes(zip_bytes({"docs/readme": b"no", "nested/wiki_en_token.arpa": b"yes"})) + output = tmp_path / "out" + output.mkdir() + assert tool._extract_model(archive, output).read_bytes() == b"yes" + assert sorted(path.name for path in output.iterdir()) == ["wiki_en_token.arpa"] + + +@pytest.mark.parametrize("name", ["../wiki_en_token.arpa", "/tmp/wiki_en_token.arpa", "C:\\..\\wiki_en_token.arpa"]) +def test_malicious_archive_path_is_rejected(tmp_path, name): + archive = tmp_path / "bad.zip" + archive.write_bytes(zip_bytes({name: b"bad"})) + output = tmp_path / "out" + output.mkdir() + with pytest.raises(RuntimeError, match="unsafe"): + tool._extract_model(archive, output) + + +def test_corrupt_kenlm_data_is_reported(monkeypatch, tmp_path): + model = tmp_path / "wiki_en_token.arpa.bin" + model.write_bytes(b"corrupt") + monkeypatch.setattr(tool.subprocess, "run", lambda *_a, **_k: (_ for _ in ()).throw( + tool.subprocess.CalledProcessError(1, "python") + )) + with pytest.raises(tool.CommandError, match="Validating the English KenLM model failed"): + tool.validate_language_model(Path("python"), model) + + +def test_already_installed_verified_model(monkeypatch, tmp_path): + monkeypatch.setattr(tool, "REPO_ROOT", tmp_path) + model = tmp_path / "lms" / "wiki_en_token.arpa.bin" + model.parent.mkdir() + model.write_bytes(b"model") + assert tool.installed_language_model(hashlib.sha256(b"model").hexdigest()) == model + assert tool.installed_language_model("0" * 64) is None + + +def test_rejected_preferred_binary_blocks_arpa_fallback(monkeypatch, tmp_path): + monkeypatch.setattr(tool, "REPO_ROOT", tmp_path) + models = tmp_path / "lms" + models.mkdir() + (models / "wiki_en_token.arpa.bin").write_bytes(b"rejected binary") + (models / "wiki_en_token.arpa").write_bytes(b"otherwise valid arpa") + + assert tool.installed_language_model("0" * 64) is None + + +def test_rejected_preferred_binary_is_replaced(monkeypatch, tmp_path): + rejected = tmp_path / "lms" / "wiki_en_token.arpa.bin" + fallback = tmp_path / "lms" / "wiki_en_token.arpa" + replacement = tmp_path / "replacement" / "wiki_en_token.arpa.bin" + rejected.parent.mkdir() + rejected.write_bytes(b"rejected binary") + fallback.write_bytes(b"otherwise valid arpa") + monkeypatch.setattr(tool, "REPO_ROOT", tmp_path) + monkeypatch.setattr( + tool, + "install_zenodo_language_model", + lambda **_kwargs: replacement, + ) + + assert tool.prepare_language_model( + download=True, + url=tool.LANGUAGE_MODEL_URL, + expected_sha256="0" * 64, + cache_dir=tmp_path / "cache", + force=False, + installed_sha256="0" * 64, + ) == replacement + + +def test_forced_download_replaces_an_installed_model(monkeypatch, tmp_path): + installed = tmp_path / "lms" / "wiki_en_token.arpa.bin" + refreshed = tmp_path / "refreshed" / "wiki_en_token.arpa.bin" + monkeypatch.setattr(tool, "installed_language_model", lambda _digest: installed) + calls = [] + + def install(url, digest, cache, *, force): + calls.append((url, digest, cache, force)) + return refreshed + + monkeypatch.setattr(tool, "install_language_model", install) + result = tool.prepare_language_model( + download=True, + url="https://example/custom-model", + expected_sha256="1" * 64, + cache_dir=tmp_path / "cache", + force=True, + installed_sha256=None, + ) + + assert result == refreshed + assert calls == [ + ("https://example/custom-model", "1" * 64, tmp_path / "cache", True) + ] + + +def test_force_without_download_remains_opted_out(monkeypatch, tmp_path): + installed = tmp_path / "lms" / "wiki_en_token.arpa.bin" + monkeypatch.setattr(tool, "installed_language_model", lambda _digest: installed) + monkeypatch.setattr( + tool, + "install_language_model", + lambda *_args, **_kwargs: pytest.fail("download must remain opt-in"), + ) + + assert tool.prepare_language_model( + download=False, + url="https://example/custom-model", + expected_sha256="1" * 64, + cache_dir=tmp_path / "cache", + force=True, + installed_sha256=None, + ) == installed + + +def test_manual_binary_is_not_compared_with_download_digest(): + assert tool.installed_model_digest( + download=False, + url=tool.LANGUAGE_MODEL_URL, + expected_sha256=tool.LANGUAGE_MODEL_SHA256, + ) is None + + +def test_managed_builtin_download_verifies_installed_binary(): + assert tool.installed_model_digest( + download=True, + url=tool.LANGUAGE_MODEL_URL, + expected_sha256=tool.LANGUAGE_MODEL_SHA256, + ) == tool.LANGUAGE_MODEL_SHA256 + + +def test_download_remains_opt_in(monkeypatch): + monkeypatch.setattr(sys, "argv", ["tool"]) + assert tool.parse_args().download_language_model is False + + +def test_range_reader_validates_content_range_and_coalesces(monkeypatch): + payload = bytes(range(256)) * 10000 + calls = [] + + def open_range(request, **_kwargs): + start, end = map(int, request.headers["Range"].removeprefix("bytes=").split("-")) + calls.append((start, end)) + return Response(payload[start:end + 1], { + "Content-Range": f"bytes {start}-{end}/{len(payload)}" + }) + + monkeypatch.setattr(tool.urllib.request, "urlopen", open_range) + reader = tool.BufferedHTTPRangeReader( + tool.HTTPRangeClient("https://example/archive", len(payload)), 1024 * 1024 + ) + reader.seek(17) + assert reader.read(5) == payload[17:22] + reader.seek(800_000) + assert reader.read(5) == payload[800_000:800_005] + assert len(calls) == 1 + + +@pytest.mark.parametrize("header", [None, "bytes 0-8/10", "nonsense"]) +def test_malformed_content_range_is_rejected(monkeypatch, header): + monkeypatch.setattr(tool.urllib.request, "urlopen", lambda *_a, **_k: Response( + b"0123456789", {"Content-Range": header} if header else {} + )) + with pytest.raises(RuntimeError, match="Content-Range"): + tool.HTTPRangeClient("https://example/archive", 10).read(0, 9) + + +def test_http_200_range_response_is_rejected(monkeypatch): + monkeypatch.setattr(tool.urllib.request, "urlopen", lambda *_a, **_k: Response( + b"whole archive", status=200 + )) + with pytest.raises(RuntimeError, match="HTTP 200"): + tool.HTTPRangeClient("https://example/archive", 100).read(0, 9) + + +def _compressed_case(data: bytes): + compressor = zlib.compressobj(wbits=-zlib.MAX_WBITS) + compressed = compressor.compress(data) + compressor.flush() + + class Client: + def chunks(self, _start, _end): + yield compressed + + return Client(), compressed + + +def test_raw_deflate_member_download(tmp_path): + data = b"ordinary ZIP member" * 100 + client, compressed = _compressed_case(data) + output = tmp_path / "model" + tool._expand_deflate_range( + client, 0, len(compressed) - 1, output, + compressed_size=len(compressed), expanded_size=len(data), + expected_crc=zlib.crc32(data), expected_sha256=hashlib.sha256(data).hexdigest(), + ) + assert output.read_bytes() == data + + +@pytest.mark.parametrize("failure", ["truncated", "overlong", "crc", "sha"]) +def test_corrupt_member_downloads_are_rejected(tmp_path, failure): + data = b"model data" * 100 + client, compressed = _compressed_case(data) + kwargs = dict(compressed_size=len(compressed), expanded_size=len(data), + expected_crc=zlib.crc32(data), + expected_sha256=hashlib.sha256(data).hexdigest()) + if failure == "truncated": + kwargs["compressed_size"] += 1 + elif failure == "overlong": + kwargs["expanded_size"] -= 1 + elif failure == "crc": + kwargs["expected_crc"] ^= 1 + else: + kwargs["expected_sha256"] = "0" * 64 + with pytest.raises(RuntimeError): + tool._expand_deflate_range(client, 0, len(compressed) - 1, + tmp_path / failure, **kwargs) + + +def test_zip64_extra_supports_large_sizes(): + large_expanded, large_compressed = 6_000_000_000, 5_000_000_000 + extra = b"\x01\x00\x10\x00" + large_expanded.to_bytes(8, "little") \ + + large_compressed.to_bytes(8, "little") + assert tool._zip64_values(extra) == [large_expanded, large_compressed] + + +def test_zip_metadata_derives_member_bounds_and_checks_local_header(monkeypatch): + data = b"member" * 100 + output = io.BytesIO() + with zipfile.ZipFile(output, "w", compression=zipfile.ZIP_DEFLATED) as archive: + archive.writestr(tool.LANGUAGE_MODEL_MEMBER, data) + payload = bytearray(output.getvalue()) + + class MemoryClient: + size = len(payload) + + def read(self, start, end): + return bytes(payload[start:end + 1]) + + with zipfile.ZipFile(io.BytesIO(payload)) as archive: + expected = archive.getinfo(tool.LANGUAGE_MODEL_MEMBER) + monkeypatch.setattr(tool, "LANGUAGE_MODEL_SIZE", len(data)) + monkeypatch.setattr(tool, "LANGUAGE_MODEL_COMPRESSED_SIZE", expected.compress_size) + monkeypatch.setattr(tool, "LANGUAGE_MODEL_CRC32", expected.CRC) + info, start, end = tool.inspect_zenodo_member(MemoryClient()) + assert end - start + 1 == info.compress_size + + payload[info.header_offset + 30] ^= 1 + with pytest.raises(RuntimeError, match="local and central"): + tool.inspect_zenodo_member(MemoryClient()) + + +def test_insufficient_disk_space(monkeypatch, tmp_path): + usage = type("Usage", (), {"free": 1})() + monkeypatch.setattr(tool.shutil, "disk_usage", lambda _path: usage) + with pytest.raises(RuntimeError, match="Insufficient disk space"): + tool._check_model_space(tmp_path) + + +def test_retry_exhaustion(monkeypatch): + calls = [] + monkeypatch.setattr(tool.time, "sleep", lambda _seconds: None) + + def fail(*_args, **_kwargs): + calls.append(1) + raise urllib.error.URLError("interrupted") + + monkeypatch.setattr(tool.urllib.request, "urlopen", fail) + with pytest.raises(RuntimeError, match="range request failed"): + tool.HTTPRangeClient("https://example/archive", 10, retries=2).read(0, 9) + assert len(calls) == 3 + + +@pytest.mark.skipif(not os.environ.get("PATHBENCH_LIVE_ZENODO_METADATA"), + reason="set PATHBENCH_LIVE_ZENODO_METADATA=1 for live range check") +def test_live_zenodo_metadata(): + client = tool.HTTPRangeClient(tool.LANGUAGE_MODEL_URL, + tool.LANGUAGE_MODEL_ARCHIVE_SIZE) + info, start, end = tool.inspect_zenodo_member(client) + assert info.filename == tool.LANGUAGE_MODEL_MEMBER + assert end - start + 1 == tool.LANGUAGE_MODEL_COMPRESSED_SIZE + assert (start, end) == (18_177_078_384, 26_759_745_295) diff --git a/tests/test_string_clean.py b/tests/test_string_clean.py new file mode 100644 index 0000000..002801e --- /dev/null +++ b/tests/test_string_clean.py @@ -0,0 +1,37 @@ +"""Tests for text cleaning and the pinned phonemizer API.""" + +import importlib.util +import os +from pathlib import Path + +import pytest + + +_SPEC = importlib.util.spec_from_file_location( + "pathbench_string_clean", Path(__file__).parents[1] / "pathbench" / "string_clean.py" +) +assert _SPEC and _SPEC.loader +_STRING_CLEAN = importlib.util.module_from_spec(_SPEC) +_SPEC.loader.exec_module(_STRING_CLEAN) +cached_phonemize = _STRING_CLEAN.cached_phonemize + + +def test_phonemizer_fork_exposes_expected_api(): + """The PyPI fork must retain the imports used by ``string_clean``.""" + from phonemizer.phonemize import phonemize + from phonemizer.separator import Separator + + assert callable(phonemize) + assert Separator(phone=" ", word="|").phone == " " + + +@pytest.mark.skipif( + os.environ.get("PATHBENCH_TEST_PINNED_ESPEAK") != "1", + reason="requires the README-pinned espeak-ng commit", +) +def test_pinned_espeak_preserves_language_specific_ipa(): + """The documented espeak-ng revision uses an Italian tap, unlike English.""" + cached_phonemize.cache_clear() + + assert cached_phonemize("Roma", "it").strip() == "ɾ o m a" + assert cached_phonemize("Roma", "en-us").strip() == "ɹ oʊ m ə" diff --git a/tools/test_gpu_predictors.py b/tools/test_gpu_predictors.py new file mode 100755 index 0000000..b5a13a1 --- /dev/null +++ b/tools/test_gpu_predictors.py @@ -0,0 +1,750 @@ +#!/usr/bin/env python3 +"""Install a CUDA-enabled PathBench environment and smoke-test ArtP and DArtP. + +Stepwise installation (each step skips components that are already suitable): + +1. Install the Ubuntu prerequisites listed in README.md if they are missing. +2. Build the pinned espeak-ng commit unless the installation marker confirms it. +3. Ensure an NVIDIA driver is installed and ``nvidia-smi --list-gpus`` works. + Host driver installation is intentionally left to the machine or cloud provider. +4. From the PathBench checkout, run ``python3 tools/test_gpu_predictors.py``. +5. This script then reuses or creates ``tools/gpu_venv``; installs PyTorch, + torchaudio, PathBench, and test dependencies only when its import/version + checks fail; verifies CUDA access; and runs the focused ArtP and DArtP tests. +6. Download the English n-gram model described in README.md to exercise DArtP; + without it, pytest reports the DArtP test as skipped. + +Run with ``--help`` to select another Python, CUDA/PyTorch version, or venv. +Allow at least 12 GB system RAM and 8 GB GPU VRAM (16 GB RAM and 12--16 GB +VRAM recommended). Google Colab NVIDIA GPU runtimes are supported when +``nvidia-smi`` works; T4, L4, and A100 assignments meet the recommendation. +""" + +from __future__ import annotations + +import argparse +import hashlib +import http.client +import io +import os +from pathlib import Path +import shutil +import subprocess +import sys +import tarfile +import tempfile +import time +import urllib.error +import urllib.parse +import urllib.request +import zipfile +import binascii +import re +import struct +import zlib + +SCRIPT_DIR = Path(__file__).resolve().parent +REPO_ROOT = SCRIPT_DIR.parent +MODEL_NAMES = ("wiki_en_token.arpa.bin", "wiki_en_token.arpa") +# The versioned record URL is deliberately not the mutable ``latest`` link. +LANGUAGE_MODEL_URL = ( + "https://zenodo.org/api/records/18738598/files/lms.zip/content" +) +# The MD5 describes the complete archive. Range extraction cannot verify it. +LANGUAGE_MODEL_ARCHIVE_SIZE = 35_017_940_434 +LANGUAGE_MODEL_ARCHIVE_MD5 = "01d62027902e93270e5f0d00806c473c" +LANGUAGE_MODEL_MEMBER = "lms/wiki_en_token.arpa.bin" +LANGUAGE_MODEL_SIZE = 14_600_342_241 +LANGUAGE_MODEL_COMPRESSED_SIZE = 8_582_666_912 +LANGUAGE_MODEL_CRC32 = 0x5AFB90EF +# Calculated from the decompressed member, not lms.zip or its DEFLATE stream. +LANGUAGE_MODEL_SHA256 = "d786eec55174c696c0bf3327928ff496684f482194ba3c6ebdf4311acb823d00" +USER_AGENT = "PathBench GPU predictor model installer/1.0" +RANGE_BLOCK_SIZE = 1024 * 1024 +RANGE_RETRIES = 4 + + +class CommandError(RuntimeError): + """A child command failed and its exit status should be preserved.""" + + def __init__(self, message: str, returncode: int) -> None: + super().__init__(message) + self.returncode = returncode + + +def run( + command: list[str], description: str, *, cwd: Path | None = None +) -> None: + """Run a visible command and turn a nonzero status into a useful error.""" + print(f"\n==> {description}", flush=True) + print("+ " + " ".join(command), flush=True) + try: + subprocess.run(command, check=True, cwd=cwd) + except subprocess.CalledProcessError as error: + raise CommandError( + f"{description} failed with exit status {error.returncode}. " + f"Review the command output above: {' '.join(command)}", + error.returncode, + ) from error + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Create a CUDA-enabled PathBench virtual environment, verify GPU " + "access, and run the ArtP and DArtP smoke tests." + ) + ) + parser.add_argument( + "--python", default=os.environ.get("PYTHON", "python3"), + help="Python 3.10-3.12 executable (default: PYTHON or python3)", + ) + parser.add_argument( + "--pytorch-version", default=os.environ.get("PYTORCH_VERSION", "2.6.0"), + help="matching PyTorch and torchaudio version (default: 2.6.0)", + ) + parser.add_argument( + "--cuda-version", default=os.environ.get("PATHBENCH_CUDA_VERSION", "12.4"), + help="CUDA wheel version, such as 12.4 (default: 12.4)", + ) + parser.add_argument( + "--venv", type=Path, + default=Path(os.environ.get("VENV", SCRIPT_DIR / "gpu_venv")), + help="virtual-environment destination (default: tools/gpu_venv)", + ) + parser.add_argument( + "--download-language-model", action="store_true", + help="download and install the English DArtP model (large; opt in)", + ) + parser.add_argument( + "--language-model-url", + default=os.environ.get("PATHBENCH_LANGUAGE_MODEL_URL", LANGUAGE_MODEL_URL), + help="model/archive URL (PATHBENCH_LANGUAGE_MODEL_URL)", + ) + parser.add_argument( + "--language-model-sha256", + default=os.environ.get("PATHBENCH_LANGUAGE_MODEL_SHA256"), + help="SHA-256 for a custom URL (PATHBENCH_LANGUAGE_MODEL_SHA256)", + ) + parser.add_argument( + "--language-model-cache", type=Path, + default=Path(os.environ.get("PATHBENCH_LANGUAGE_MODEL_CACHE", "~/.cache/pathbench")), + help="download cache (PATHBENCH_LANGUAGE_MODEL_CACHE; default: ~/.cache/pathbench)", + ) + parser.add_argument( + "--force-language-model-download", action="store_true", + help="discard a cached artifact and download it again", + ) + return parser.parse_args() + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +class HTTPRangeClient: + """Strict, retrying HTTP byte-range client for a pinned immutable object.""" + + def __init__(self, url: str, size: int, *, retries: int = RANGE_RETRIES) -> None: + self.url = url + self.size = size + self.retries = retries + + def _open(self, start: int, end: int): + if start < 0 or end < start or end >= self.size: + raise RuntimeError(f"Invalid archive byte range {start}-{end}") + request = urllib.request.Request( + self.url, + headers={"Range": f"bytes={start}-{end}", "User-Agent": USER_AGENT}, + ) + response = urllib.request.urlopen(request, timeout=60) + status = getattr(response, "status", None) or response.getcode() + if status != 206: + response.close() + raise RuntimeError( + f"Range server returned HTTP {status}, not 206; refusing a possible " + f"full {self.size:,}-byte archive response" + ) + value = response.headers.get("Content-Range") + match = re.fullmatch(r"bytes (\d+)-(\d+)/(\d+)", value or "") + if not match: + response.close() + raise RuntimeError(f"Malformed or missing Content-Range: {value!r}") + reported = tuple(map(int, match.groups())) + if reported != (start, end, self.size): + response.close() + raise RuntimeError( + f"Unexpected Content-Range {value!r}; expected bytes " + f"{start}-{end}/{self.size}" + ) + return response + + @staticmethod + def _transient(error: BaseException) -> bool: + if isinstance(error, RuntimeError): + return False + if isinstance(error, urllib.error.HTTPError): + return error.code in {408, 425, 429, 500, 502, 503, 504} + return isinstance(error, (OSError, http.client.HTTPException, + urllib.error.URLError)) + + def read(self, start: int, end: int) -> bytes: + for attempt in range(self.retries + 1): + try: + with self._open(start, end) as response: + data = response.read(end - start + 1) + if response.read(1) or len(data) != end - start + 1: + raise OSError("truncated or overlong HTTP range response") + return data + except Exception as error: + if not self._transient(error) or attempt == self.retries: + raise RuntimeError(f"HTTP range request failed: {error}") from error + time.sleep(min(2 ** attempt, 8)) + raise AssertionError("unreachable") + + def chunks(self, start: int, end: int, chunk_size: int = 4 * 1024 * 1024): + """Yield an exact interval, resuming interrupted responses with a new range.""" + position = start + failures = 0 + while position <= end: + response = None + try: + response = self._open(position, end) + while position <= end: + chunk = response.read(min(chunk_size, end - position + 1)) + if not chunk: + raise OSError("truncated HTTP range response") + position += len(chunk) + yield chunk + if response.read(1): + raise OSError("overlong HTTP range response") + failures = 0 + except Exception as error: + if not self._transient(error) or failures >= self.retries: + raise RuntimeError(f"HTTP range download failed: {error}") from error + time.sleep(min(2 ** failures, 8)) + failures += 1 + finally: + if response is not None: + response.close() + + +class BufferedHTTPRangeReader(io.RawIOBase): + """Seekable range-backed file with block coalescing for ZIP metadata reads.""" + + def __init__(self, client: HTTPRangeClient, block_size: int = RANGE_BLOCK_SIZE): + self.client = client + self.block_size = block_size + self.position = 0 + self.cache: dict[int, bytes] = {} + + def readable(self): + return True + + def seekable(self): + return True + + def tell(self): + return self.position + + def seek(self, offset, whence=os.SEEK_SET): + positions = {os.SEEK_SET: offset, os.SEEK_CUR: self.position + offset, + os.SEEK_END: self.client.size + offset} + if whence not in positions or positions[whence] < 0: + raise ValueError("invalid seek") + self.position = positions[whence] + return self.position + + def read(self, size=-1): + if size is None or size < 0: + size = self.client.size - self.position + size = min(size, self.client.size - self.position) + output = bytearray() + while size > 0: + block = self.position // self.block_size + if block not in self.cache: + start = block * self.block_size + end = min(start + self.block_size, self.client.size) - 1 + self.cache[block] = self.client.read(start, end) + data = self.cache[block] + within = self.position - block * self.block_size + take = min(size, len(data) - within) + output.extend(data[within:within + take]) + self.position += take + size -= take + return bytes(output) + + +def _zip64_values(extra: bytes) -> list[int]: + position = 0 + while position + 4 <= len(extra): + kind, length = struct.unpack_from(" tuple[zipfile.ZipInfo, int, int]: + """Read ZIP/ZIP64 metadata remotely and return the raw DEFLATE interval.""" + reader = BufferedHTTPRangeReader(client) + try: + with zipfile.ZipFile(reader) as archive: + info = archive.getinfo(LANGUAGE_MODEL_MEMBER) + except (KeyError, zipfile.BadZipFile, OSError) as error: + raise RuntimeError(f"Cannot inspect Zenodo ZIP metadata: {error}") from error + if not _safe_member(info.filename) or info.filename != LANGUAGE_MODEL_MEMBER: + raise RuntimeError(f"Unsafe or unexpected ZIP member name: {info.filename!r}") + expected = (zipfile.ZIP_DEFLATED, LANGUAGE_MODEL_SIZE, + LANGUAGE_MODEL_COMPRESSED_SIZE, LANGUAGE_MODEL_CRC32) + actual = (info.compress_type, info.file_size, info.compress_size, info.CRC) + if actual != expected: + raise RuntimeError(f"Zenodo member metadata changed: {actual!r} != {expected!r}") + if info.flag_bits & 1: + raise RuntimeError("Encrypted ZIP members are unsupported") + + fixed = client.read(info.header_offset, info.header_offset + 29) + signature, _version, flags, method, _time, _date, crc, compressed, expanded, name_len, extra_len = \ + struct.unpack(" bool: + path = Path(name.replace("\\", "/")) + return not path.is_absolute() and ".." not in path.parts + + +def _extract_model(archive: Path, destination: Path) -> Path: + """Extract only the preferred English model, after checking every path.""" + members: dict[str, object] + opener: object + if zipfile.is_zipfile(archive): + opener = zipfile.ZipFile(archive) + members = {item.filename: item for item in opener.infolist()} + elif tarfile.is_tarfile(archive): + opener = tarfile.open(archive) + members = {item.name: item for item in opener.getmembers()} + else: + # A direct ARPA or KenLM binary needs no extraction. + target = destination / next( + (name for name in MODEL_NAMES if name in archive.name), MODEL_NAMES[0] + ) + shutil.copyfile(archive, target) + return target + with opener: + if any(not _safe_member(name) for name in members): + raise RuntimeError("Language-model archive contains an unsafe path") + selected = next( + (name for wanted in MODEL_NAMES for name in members if Path(name).name == wanted), + None, + ) + if selected is None: + raise RuntimeError("Archive contains neither wiki_en_token.arpa.bin nor wiki_en_token.arpa") + target = destination / Path(selected).name + source = opener.open(members[selected]) if isinstance(opener, zipfile.ZipFile) else opener.extractfile(members[selected]) + if source is None: + raise RuntimeError("Selected language-model archive member is not a file") + with source, target.open("wb") as output: + shutil.copyfileobj(source, output) + return target + + +def _check_model_space(directory: Path, required_size: int = LANGUAGE_MODEL_SIZE) -> None: + """Require room for the staged model plus a conservative one-GiB margin.""" + available = shutil.disk_usage(directory).free + required = required_size + 1024 ** 3 + if available < required: + raise RuntimeError( + f"Insufficient disk space: {available:,} bytes free; {required:,} required " + "for the temporary expanded model and safety margin" + ) + + +def _expand_deflate_range( + client: HTTPRangeClient, start: int, end: int, destination: Path, *, + compressed_size: int = LANGUAGE_MODEL_COMPRESSED_SIZE, + expanded_size: int = LANGUAGE_MODEL_SIZE, + expected_crc: int = LANGUAGE_MODEL_CRC32, + expected_sha256: str = LANGUAGE_MODEL_SHA256, +) -> None: + """Download and authenticate one raw-DEFLATE ZIP member.""" + inflater = zlib.decompressobj(-zlib.MAX_WBITS) + transferred = expanded = crc = 0 + digest = hashlib.sha256() + last_report = time.monotonic() + with destination.open("wb") as output: + for chunk in client.chunks(start, end): + transferred += len(chunk) + if transferred > compressed_size: + raise RuntimeError("Compressed member exceeds its advertised size") + data = inflater.decompress(chunk, expanded_size - expanded + 1) + if inflater.unconsumed_tail: + raise RuntimeError("Expanded model exceeds its advertised size") + expanded += len(data) + if expanded > expanded_size: + raise RuntimeError("Expanded model exceeds its advertised size") + output.write(data) + crc = binascii.crc32(data, crc) + digest.update(data) + if time.monotonic() - last_report >= 5: + print(f"Transferred {transferred:,}/{compressed_size:,}; expanded " + f"{expanded:,}/{expanded_size:,} bytes...", flush=True) + last_report = time.monotonic() + tail = inflater.flush() + expanded += len(tail) + if expanded > expanded_size: + raise RuntimeError("Expanded model exceeds its advertised size") + output.write(tail) + crc = binascii.crc32(tail, crc) + digest.update(tail) + print(f"Transferred {transferred:,}/{compressed_size:,}; expanded " + f"{expanded:,}/{expanded_size:,} bytes.", flush=True) + if transferred != compressed_size: + raise RuntimeError(f"Compressed size mismatch: {transferred:,} != {compressed_size:,}") + if not inflater.eof or inflater.unused_data: + raise RuntimeError("Truncated or trailing compressed member data") + if expanded != expanded_size: + raise RuntimeError(f"Expanded size mismatch: {expanded:,} != {expanded_size:,}") + if crc & 0xFFFFFFFF != expected_crc: + raise RuntimeError(f"Model CRC-32 mismatch: {crc & 0xFFFFFFFF:08x}") + actual_sha256 = digest.hexdigest() + if actual_sha256 != expected_sha256.lower(): + raise RuntimeError( + f"Model SHA-256 mismatch (expected {expected_sha256}, got {actual_sha256})" + ) + + +def install_zenodo_language_model( + *, force: bool = False, validator=None, +) -> Path: + """Range-extract, validate, and atomically install the pinned ZIP member.""" + models_dir = REPO_ROOT / "lms" + models_dir.mkdir(parents=True, exist_ok=True) + installed = models_dir / Path(LANGUAGE_MODEL_MEMBER).name + if not force and installed.is_file() and installed.stat().st_size == LANGUAGE_MODEL_SIZE \ + and sha256(installed) == LANGUAGE_MODEL_SHA256: + print(f"Reusing verified installed language model: {installed}") + return installed + print( + "The Zenodo range download will transfer approximately 8.0 GiB and install " + "a 13.6 GiB English language model. Additional temporary space and a 1 GiB " + "safety margin are required.", flush=True, + ) + _check_model_space(models_dir) + client = HTTPRangeClient(LANGUAGE_MODEL_URL, LANGUAGE_MODEL_ARCHIVE_SIZE) + _info, start, end = inspect_zenodo_member(client) + temporary = Path(tempfile.mkstemp(prefix=".wiki-en-", dir=models_dir)[1]) + try: + _expand_deflate_range(client, start, end, temporary) + if validator is not None: + validator(temporary) + os.replace(temporary, installed) + return installed + finally: + temporary.unlink(missing_ok=True) + + +def install_language_model( + url: str, expected_sha256: str, cache_dir: Path, *, force: bool = False, + validator=None, +) -> Path: + """Install a custom standalone file/archive with a user-supplied digest.""" + if not expected_sha256 or len(expected_sha256) != 64: + raise RuntimeError("A pinned 64-character SHA-256 is required for the language model") + expected_sha256 = expected_sha256.lower() + cache_dir = cache_dir.expanduser().resolve() + cache_dir.mkdir(parents=True, exist_ok=True) + source_name = Path(urllib.parse.unquote(urllib.parse.urlparse(url).path)).name + source_name = source_name if source_name else "download" + artifact = cache_dir / f"language-model-{expected_sha256}-{source_name}" + if force: + artifact.unlink(missing_ok=True) + if artifact.exists() and sha256(artifact) != expected_sha256: + artifact.unlink() + if not artifact.exists(): + temporary = Path(tempfile.mkstemp(prefix=".language-model-", dir=cache_dir)[1]) + try: + request = urllib.request.Request(url, headers={"User-Agent": "PathBench/0.1"}) + with urllib.request.urlopen(request, timeout=60) as response, temporary.open("wb") as output: + total = 0 + last_report = time.monotonic() + while chunk := response.read(1024 * 1024): + output.write(chunk) + total += len(chunk) + if time.monotonic() - last_report >= 5: + print(f"Downloaded {total:,} bytes...", flush=True) + last_report = time.monotonic() + print(f"Downloaded {total:,} bytes.", flush=True) + actual = sha256(temporary) + if actual != expected_sha256: + raise RuntimeError( + f"Language-model SHA-256 mismatch (expected {expected_sha256}, got {actual})" + ) + os.replace(temporary, artifact) + except (OSError, urllib.error.URLError, http.client.HTTPException) as error: + raise RuntimeError(f"Language-model download failed: {error}") from error + finally: + temporary.unlink(missing_ok=True) + else: + print(f"Reusing verified language-model download: {artifact}") + + models_dir = REPO_ROOT / "lms" + models_dir.mkdir(parents=True, exist_ok=True) + staging = Path(tempfile.mkdtemp(prefix=".model-", dir=models_dir)) + try: + extracted = _extract_model(artifact, staging) + if extracted.stat().st_size == 0: + raise RuntimeError("Downloaded language model is empty") + if validator is not None: + validator(extracted) + installed = models_dir / extracted.name + if installed.name == "wiki_en_token.arpa": + # The evaluator prefers the binary whenever it exists. Removing a + # stale binary ensures it actually consumes the verified ARPA we + # are about to install and validate. + (models_dir / "wiki_en_token.arpa.bin").unlink(missing_ok=True) + os.replace(extracted, installed) + return installed + finally: + shutil.rmtree(staging, ignore_errors=True) + + +def installed_language_model(expected_sha256: str | None = None) -> Path | None: + """Return the preferred installed model, rejecting empty/known-bad files.""" + for name in MODEL_NAMES: + path = REPO_ROOT / "lms" / name + if not path.is_file() or path.stat().st_size == 0: + continue + # The built-in artifact is the compiled model itself, so its digest also + # authenticates an already-installed copy. Archive digests do not. + if expected_sha256 and name == "wiki_en_token.arpa.bin": + if sha256(path) != expected_sha256.lower(): + # The evaluator always prefers this binary when it exists. Do + # not fall back to an ARPA that the smoke test would not use; + # make the opted-in preparation path replace the bad binary. + return None + return path + return None + + +def validate_language_model(venv_python: Path, model: Path) -> None: + """Have the target environment parse the model before expensive tests start.""" + run([ + str(venv_python), "-c", + "import kenlm, pathlib, sys; p=pathlib.Path(sys.argv[1]); " + "assert p.stat().st_size, 'model is empty'; kenlm.Model(str(p))", + str(model), + ], "Validating the English KenLM model") + + +def prepare_language_model( + *, + download: bool, + url: str, + expected_sha256: str, + cache_dir: Path, + force: bool, + installed_sha256: str | None, + validator=None, +) -> Path | None: + """Reuse an installed model, unless an opted-in forced refresh was requested.""" + model = installed_language_model(installed_sha256) + if download and (model is None or force): + if url == LANGUAGE_MODEL_URL: + return install_zenodo_language_model(force=force, validator=validator) + kwargs = {"force": force} + if validator is not None: + kwargs["validator"] = validator + return install_language_model(url, expected_sha256, cache_dir, **kwargs) + return model + + +def installed_model_digest( + *, download: bool, url: str, expected_sha256: str | None, +) -> str | None: + """Return a digest only when this run is managing the pinned artifact. + + Manually installed models are supported and need not be byte-for-byte + identical to the project's default Zenodo binary. They are checked by + KenLM's parser instead of against the download artifact's digest. + """ + if download and url == LANGUAGE_MODEL_URL: + return expected_sha256 + return None + + +def require_program(program: str, explanation: str) -> str: + path = shutil.which(program) + if path is None: + raise RuntimeError(f"Required program '{program}' was not found. {explanation}") + return path + + +def python_succeeds(python: Path | str, code: str) -> bool: + """Return whether an environment already satisfies a Python-side check.""" + return subprocess.run( + [str(python), "-c", code], stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ).returncode == 0 + + +def main() -> int: + args = parse_args() + venv = args.venv.expanduser().resolve() + try: + python = require_program( + args.python, "Set --python to a Python 3.10-3.12 executable." + ) + nvidia_smi = require_program( + "nvidia-smi", "Install an NVIDIA driver and expose the GPU to this environment." + ) + require_program( + "espeak-ng", "Install the pinned version described in README.md first." + ) + + version_check = subprocess.run( + [ + python, "-c", + "import sys; print(f'{sys.version_info.major}.{sys.version_info.minor}'); " + "raise SystemExit(0 if (3, 10) <= sys.version_info[:2] <= (3, 12) else 1)", + ], + text=True, capture_output=True, + ) + if version_check.returncode != 0: + detected = version_check.stdout.strip() or "unknown" + raise RuntimeError( + f"Python 3.10-3.12 is required, but '{python}' reports {detected}." + ) + + run([nvidia_smi, "--list-gpus"], "Checking the NVIDIA driver and visible GPUs") + if not (venv / "bin" / "python").is_file(): + run([python, "-m", "venv", str(venv)], "Creating the GPU virtual environment") + else: + print(f"\n==> Reusing existing virtual environment: {venv}") + venv_python = venv / "bin" / "python" + cuda_tag = f"cu{args.cuda_version.replace('.', '')}" + + torch_check = ( + "import torch, torchaudio; " + f"assert torch.__version__.split('+')[0] == '{args.pytorch_version}'; " + f"assert torchaudio.__version__.split('+')[0] == '{args.pytorch_version}'; " + f"assert (torch.version.cuda or '').replace('.', '') == '{cuda_tag[2:]}'" + ) + if not python_succeeds(venv_python, torch_check): + run([ + str(venv_python), "-m", "pip", "install", + "--force-reinstall", + f"torch=={args.pytorch_version}", f"torchaudio=={args.pytorch_version}", + "--index-url", f"https://download.pytorch.org/whl/{cuda_tag}", + ], "Installing missing or mismatched CUDA-enabled PyTorch packages") + else: + print("\n==> Matching CUDA-enabled PyTorch packages are already installed") + + project_check = ( + "from pathlib import Path; import pathbench, pytest, kenlm; " + f"assert Path(pathbench.__file__).resolve().is_relative_to(Path({str(REPO_ROOT)!r}))" + ) + if not python_succeeds(venv_python, project_check): + run([ + str(venv_python), "-m", "pip", "install", "-e", f"{REPO_ROOT}[all]", + "pytest", "kenlm", + ], "Installing missing PathBench or test dependencies") + else: + print("\n==> PathBench and test dependencies are already installed") + run([ + str(venv_python), "-c", + "import torch; " + "print(f'PyTorch: {torch.__version__}'); " + "print(f'PyTorch CUDA runtime: {torch.version.cuda}'); " + "assert torch.cuda.is_available(), " + "'CUDA build installed, but no GPU is available; check the driver and container access'; " + "print(f'GPU: {torch.cuda.get_device_name(0)}')", + ], "Verifying that PyTorch can use the GPU") + + expected_digest = args.language_model_sha256 + if args.language_model_url == LANGUAGE_MODEL_URL and not expected_digest: + expected_digest = LANGUAGE_MODEL_SHA256 + known_installed_digest = installed_model_digest( + download=args.download_language_model, + url=args.language_model_url, + expected_sha256=expected_digest, + ) + model = prepare_language_model( + download=args.download_language_model, + url=args.language_model_url, + expected_sha256=expected_digest or "", + cache_dir=args.language_model_cache, + force=args.force_language_model_download, + installed_sha256=known_installed_digest, + validator=lambda path: validate_language_model(venv_python, path), + ) + if model is None: + print( + "\nWarning: the English n-gram model is absent. The DArtP test " + "will be reported as skipped; see README.md's N-gram models section.", + file=sys.stderr, + ) + if model is not None: + validate_language_model(venv_python, model) + + pytest_environment = os.environ.copy() + pytest_environment["MPLBACKEND"] = "Agg" + print("\n==> Running the ArtP and DArtP smoke tests", flush=True) + command = [ + str(venv_python), "-m", "pytest", + "tests/test_evaluators.py::TestEvaluatorMethods::test_articulatory_precision", + "tests/test_evaluators.py::TestEvaluatorMethods::test_artp_double_asr", "-v", + ] + print("+ " + " ".join(command), flush=True) + try: + subprocess.run(command, check=True, cwd=REPO_ROOT, env=pytest_environment) + except subprocess.CalledProcessError as error: + raise CommandError( + f"Running the ArtP and DArtP smoke tests failed with exit status " + f"{error.returncode}. Review the command output above: {' '.join(command)}", + error.returncode, + ) from error + except CommandError as error: + print(f"\nError: {error}", file=sys.stderr) + return error.returncode + except RuntimeError as error: + print(f"\nError: {error}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())