Add Canopy-M inference script and package

README.md

@@ -1,147 +1,178 @@
----
-language:
-  - ur
-  - sd
-license: apache-2.0
-pretty_name: Canopy-M
-pipeline_tag: automatic-speech-recognition
-library_name: pytorch
-tags:
-  - automatic-speech-recognition
-  - speech
-  - ctc
-  - conformer
-  - urdu
-  - sindhi
----
-
-# Canopy-M
-
-
-<div align="center">
-  <img src="figures/canopy-m-logo.png" alt="Canopy-M logo" width="220"/>
-</div>
-
-
-Canopy-M was trained on Urdu and Sindhi. The numbers below are the evaluation.
-
-
-![Canopy-M versus Whisper on the Urdu and Sindhi development sets](figures/overview_cer_wer.png)
-
-
-## Use
-
-```python
-from canopy.inference.transcriber import Transcriber
-
-asr = Transcriber.from_checkpoint("canopy-m.pt", device="cuda")
-print(asr.transcribe("urdu.wav", language="urd"))
-print(asr.transcribe("sindhi.wav", language="snd"))
-```
-
-Pass `urd` or `snd`. That selects the language embedding and restricts the CTC vocabulary to that language.
-
-## Training data
-
-About **186.7 hours** used in training.
-
-| Language | Source | Train hours | Dev hours | Labels |
-|---|---|---:|---:|---|
-| Urdu | Common Voice + mahwiz + UrduSpeech | ~79.6 | ~17 | human |
-| Sindhi | Common Voice | 71.5 | 0.10 | human; mostly unvalidated `other` |
-| Sindhi | FLEURS `sd_in` | 12.3 | 1.33 | human, CC-BY-4.0 |
-
-## Benchmark
-
-Same development clips and the same `eval_norm` for every system. Whisper is forced to the `urdu` or `sindhi` token, fp16, greedy, `max_new_tokens=225`.
-
-| System | Params | Urdu CER | Urdu WER | Sindhi CER | Sindhi WER | Macro CER |
-|---|---:|---:|---:|---:|---:|---:|
-| **Canopy-M** | 57.9M | **0.132** | **0.267** | **0.107** | **0.271** | **0.119** |
-| Whisper-tiny | 39M | 1.024 | 1.477 | 1.430 | 1.599 | 1.227 |
-| Whisper-base | 74M | 0.855 | 1.247 | 0.624 | 1.083 | 0.739 |
-| whisper-small-urdu | 244M | 0.177 | 0.330 | not run | not run | — |
-
-`whisper-small-urdu` is `khawajaaliarshad/whisper-small-urdu`, an Urdu-only fine-tune, so it was not run on Sindhi. A value above 1.0 means the output is longer than the reference, usually a repetition or the wrong script.
-
-![Model size against median utterance CER](figures/size_vs_cer.png)
-
-### Urdu by source
-
-14,681 utterances, 63 speakers: Common Voice 7,772 / 54 speakers, UrduSpeech 4,296 / 7, mahwiz 2,613 / 2.
-
-| System | CV CER | UrduSpeech CER | mahwiz CER |
-|---|---:|---:|---:|
-| Canopy-M | 0.248 | 0.022 | 0.027 |
-| Whisper-tiny | 1.505 | 0.642 | 0.451 |
-| Whisper-base | 1.520 | 0.209 | 0.283 |
-| whisper-small-urdu | 0.278 | 0.069 | 0.111 |
-
-![Urdu development CER and WER by source](figures/urdu_by_source.png)
-
-Paired speaker-cluster test, Δ = whisper-small-urdu minus Canopy-M. Positive means Canopy-M makes fewer errors. “No clear difference” means the 95% interval includes zero.
-
-| Subset | ΔCER [95% CI] | Verdict |
-|---|---|---|
-| All Urdu | +0.046 [−0.003, +0.118] | no clear difference |
-| Common Voice | +0.030 [−0.042, +0.184] | no clear difference |
-| UrduSpeech | +0.047 [+0.021, +0.090] | Canopy-M lower CER |
-| mahwiz | +0.084 [+0.074, +0.099] | Canopy-M lower CER |
-
-whisper-small-urdu was fine-tuned on Common Voice Urdu 24.0 with its own split, so its Common Voice number may include clips it trained on. Canopy-M trained on the UrduSpeech and mahwiz training splits, so those two rows are in-domain for Canopy-M.
-
-### Sindhi by source
-
-496 utterances, 164 speakers: Common Voice 70 / 14 speakers, FLEURS 426 / 150 speakers. FLEURS `sd_in` is India-recorded Sindhi.
-
-| System | CV CER | FLEURS CER |
-|---|---:|---:|
-| Canopy-M | 0.286 | 0.095 |
-| Whisper-tiny | 1.911 | 1.399 |
-| Whisper-base | 0.690 | 0.619 |
-
-![Sindhi development CER and WER by source](figures/sindhi_by_source.png)
-
-Against zero-shot Whisper-base, Canopy-M is lower on every Sindhi subset. All Sindhi ΔCER is +0.517 [+0.470, +0.568]. Whisper-base did not train on this Sindhi data. Canopy-M did.
-
-### Robustness
-
-Median utterance CER is 0.059 on Urdu and 0.078 on Sindhi. Canopy-M produced no runaway outputs (hypothesis longer than 3× the reference) and no empty outputs. Whisper-tiny runaway rate is 10.3% on Urdu and 9.1% on Sindhi.
-
-![Per-utterance CER](figures/utterance_cer_distribution.png)
-
-![Runaway rate](figures/runaway_rate.png)
-
-![Paired differences](figures/paired_differences.png)
-
-Decode real-time factor on an RTX 3090 is about 0.0006 for Canopy-M, 0.0218 for Whisper-tiny, 0.0100 for Whisper-base, and 0.0069 for whisper-small-urdu. Batch sizes differed, so treat the speed column as indicative.
-
-## Intended use
-
-Research transcription of read Urdu and Sindhi at 16 kHz, one language per utterance, with the language given.
-
-## Out of scope
-
-Surveillance, speaker identification, and any decision about a person made without review. Production deployment while the test split is unscored. Telephone, broadcast, conversational, and code-switched speech.
-
-## Limitations
-
-- These scores are the development split. The test split (Urdu ~37 h, Sindhi FLEURS ~3.3 h plus a small Common Voice test) has not been opened.
-- Sindhi Common Voice development audio is 0.10 hours (70 utterances). Most of the Sindhi score is FLEURS.
-- Common Voice Sindhi training audio is mostly unvalidated `other`, included with `--include-unvalidated`.
-- Urdu Common Voice CER is 0.248. The headline 0.132 is pulled down by UrduSpeech and mahwiz.
-- A wrong language id can cross Urdu and Sindhi, which share a Perso-Arabic script.
-- No word n-gram was used in this table. Fusion can still overwrite a correct acoustic hypothesis.
-
-## Citation
-
-```bibtex
-@software{canopy_m_2026,
-  title = {Canopy-M: Conformer-CTC speech recognition for Urdu and Sindhi},
-  year = {2026},
-  url = {https://github.com/mwzkhalil/canopy},
-  license = {Apache-2.0}
-}
-```
-
-Also cite Common Voice, FLEURS, and the Metro-ASR / Conformer designs the block follows.
+---
+language:
+  - ur
+  - sd
+license: apache-2.0
+pretty_name: Canopy-M
+pipeline_tag: automatic-speech-recognition
+library_name: pytorch
+tags:
+  - automatic-speech-recognition
+  - speech
+  - ctc
+  - conformer
+  - urdu
+  - sindhi
+---
+
+# Canopy-M
+
+Canopy-M was trained on **Urdu** and **Sindhi**. The numbers below are the full **development** evaluation (greedy CTC, language id on, no language model). The held-out test split has not been scored.
+
+| | Urdu `urd` | Sindhi `snd` | Macro |
+|---|---:|---:|---:|
+| CER | 0.132 [0.065, 0.211] | 0.107 [0.091, 0.123] | 0.119 |
+| WER | 0.267 [0.166, 0.383] | 0.271 [0.240, 0.304] | 0.269 |
+| Utterances | 14,681 | 496 | |
+| Speakers | 63 | 164 | |
+
+Intervals are 95% speaker-cluster bootstrap. Lower is better. Run `canopy_exp_002_urd_snd_20261005-120848`, checkpoint selected at step 40,000, evaluated 2026-10-06.
+
+![Canopy-M versus Whisper on the Urdu and Sindhi development sets](figures/overview_cer_wer.png)
+
+## Model
+
+| | |
+|---|---|
+| Architecture | Conformer-CTC, Metro-style Macaron blocks |
+| Size | 57.9M inference parameters, d=384, 12 layers, 6 heads |
+| Languages | Urdu (id 1), Sindhi (id 2). Id 0 is unconditioned |
+| Conditioning | language embedding and a language-ID head |
+| Tokenizer | character, `char_urd_snd` |
+| Frontend | 16 kHz, 80-bin log-mel, 4× conv subsampling, 25 frames/s |
+| Objective | `(1 − 0.3) · CTC + 0.3 · auxiliary CTC` at layer 6 |
+| Decoder in this table | greedy CTC, no n-gram |
+| Training | 40,000 steps, bf16, RTX 3090, about 9,114 hours of audio seen, ~50/50 Urdu/Sindhi |
+| Code | Apache-2.0 |
+
+One block is half SwiGLU, rotary attention, a squeeze-excite depthwise convolution (kernel 31), and a second half SwiGLU, then RMSNorm. The auxiliary head is dropped at export. Sindhi keeps its own letters: Urdu yeh, kaf, and heh mappings are not applied to Sindhi.
+
+## Use
+
+The model repository is private. Log in once with a token that can read `ProximaAI/Canopy-M`.
+
+```bash
+pip install torch torchaudio soundfile numpy pyyaml huggingface_hub
+hf auth login
+hf download ProximaAI/Canopy-M --local-dir Canopy-M
+python Canopy-M/canopy_m.py clip.wav --language urd
+python Canopy-M/canopy_m.py clip.wav --language snd --device cpu
+```
+
+`--language` accepts `urd`, `ur`, `urdu`, `snd`, `sd`, or `sindhi`. That selects the language embedding and keeps the CTC output inside that language’s letters. Decoding is greedy CTC, the same setup as the table above.
+
+```python
+from canopy_m import CanopyM
+
+asr = CanopyM.from_pretrained("ProximaAI/Canopy-M")  # cuda when available, otherwise cpu
+print(asr.transcribe("clip.wav", language="urd"))
+print(asr.transcribe("clip.wav", language="snd"))
+```
+
+Run that from the `Canopy-M` directory, or add it to `PYTHONPATH`. `canopy_m.py` and the `canopy` package are both in this repository, next to `canopy-m.pt`. A 16 kHz mono file is the expected input. Other sample rates are resampled. Files longer than 30 seconds are transcribed in chunks.
+
+## Training data
+
+About **186.7 hours** used in training. YouTube Sindhi is train-only and is excluded from the scores.
+
+| Language | Source | Train hours | Dev hours | Labels |
+|---|---|---:|---:|---|
+| Urdu | Common Voice + mahwiz + UrduSpeech | ~79.6 | ~17 | human |
+| Sindhi | Common Voice | 71.5 | 0.10 | human; mostly unvalidated `other` |
+| Sindhi | FLEURS `sd_in` | 12.3 | 1.33 | human, CC-BY-4.0 |
+| Sindhi | YouTube | 23.7 | — | Whisper pseudo-labels, train only |
+
+## Benchmark
+
+Same development clips and the same `eval_norm` for every system. Whisper is forced to the `urdu` or `sindhi` token, fp16, greedy, `max_new_tokens=225`.
+
+| System | Params | Urdu CER | Urdu WER | Sindhi CER | Sindhi WER | Macro CER |
+|---|---:|---:|---:|---:|---:|---:|
+| **Canopy-M** | 57.9M | **0.132** | **0.267** | **0.107** | **0.271** | **0.119** |
+| Whisper-tiny | 39M | 1.024 | 1.477 | 1.430 | 1.599 | 1.227 |
+| Whisper-base | 74M | 0.855 | 1.247 | 0.624 | 1.083 | 0.739 |
+| whisper-small-urdu | 244M | 0.177 | 0.330 | not run | not run | — |
+
+`whisper-small-urdu` is `khawajaaliarshad/whisper-small-urdu`, an Urdu-only fine-tune, so it was not run on Sindhi. A value above 1.0 means the output is longer than the reference, usually a repetition or the wrong script.
+
+![Model size against median utterance CER](figures/size_vs_cer.png)
+
+### Urdu by source
+
+14,681 utterances, 63 speakers: Common Voice 7,772 / 54 speakers, UrduSpeech 4,296 / 7, mahwiz 2,613 / 2.
+
+| System | CV CER | UrduSpeech CER | mahwiz CER |
+|---|---:|---:|---:|
+| Canopy-M | 0.248 | 0.022 | 0.027 |
+| Whisper-tiny | 1.505 | 0.642 | 0.451 |
+| Whisper-base | 1.520 | 0.209 | 0.283 |
+| whisper-small-urdu | 0.278 | 0.069 | 0.111 |
+
+![Urdu development CER and WER by source](figures/urdu_by_source.png)
+
+Paired speaker-cluster test, Δ = whisper-small-urdu minus Canopy-M. Positive means Canopy-M makes fewer errors. “No clear difference” means the 95% interval includes zero.
+
+| Subset | ΔCER [95% CI] | Verdict |
+|---|---|---|
+| All Urdu | +0.046 [−0.003, +0.118] | no clear difference |
+| Common Voice | +0.030 [−0.042, +0.184] | no clear difference |
+| UrduSpeech | +0.047 [+0.021, +0.090] | Canopy-M lower CER |
+| mahwiz | +0.084 [+0.074, +0.099] | Canopy-M lower CER |
+
+whisper-small-urdu was fine-tuned on Common Voice Urdu 24.0 with its own split, so its Common Voice number may include clips it trained on. Canopy-M trained on the UrduSpeech and mahwiz training splits, so those two rows are in-domain for Canopy-M.
+
+### Sindhi by source
+
+496 utterances, 164 speakers: Common Voice 70 / 14 speakers, FLEURS 426 / 150 speakers. FLEURS `sd_in` is India-recorded Sindhi.
+
+| System | CV CER | FLEURS CER |
+|---|---:|---:|
+| Canopy-M | 0.286 | 0.095 |
+| Whisper-tiny | 1.911 | 1.399 |
+| Whisper-base | 0.690 | 0.619 |
+
+![Sindhi development CER and WER by source](figures/sindhi_by_source.png)
+
+Against zero-shot Whisper-base, Canopy-M is lower on every Sindhi subset. All Sindhi ΔCER is +0.517 [+0.470, +0.568]. Whisper-base did not train on this Sindhi data. Canopy-M did.
+
+### Robustness
+
+Median utterance CER is 0.059 on Urdu and 0.078 on Sindhi. Canopy-M produced no runaway outputs (hypothesis longer than 3× the reference) and no empty outputs. Whisper-tiny runaway rate is 10.3% on Urdu and 9.1% on Sindhi.
+
+![Per-utterance CER](figures/utterance_cer_distribution.png)
+
+![Runaway rate](figures/runaway_rate.png)
+
+![Paired differences](figures/paired_differences.png)
+
+Decode real-time factor on an RTX 3090 is about 0.0006 for Canopy-M, 0.0218 for Whisper-tiny, 0.0100 for Whisper-base, and 0.0069 for whisper-small-urdu. Batch sizes differed, so treat the speed column as indicative.
+
+## Intended use
+
+Research transcription of read Urdu and Sindhi at 16 kHz, one language per utterance, with the language given.
+
+## Out of scope
+
+Surveillance, speaker identification, and any decision about a person made without review. Production deployment while the test split is unscored. Telephone, broadcast, conversational, and code-switched speech.
+
+## Limitations
+
+- These scores are the development split. The test split (Urdu ~37 h, Sindhi FLEURS ~3.3 h plus a small Common Voice test) has not been opened.
+- Sindhi Common Voice development audio is 0.10 hours (70 utterances). Most of the Sindhi score is FLEURS.
+- Common Voice Sindhi training audio is mostly unvalidated `other`, included with `--include-unvalidated`.
+- Urdu Common Voice CER is 0.248. The headline 0.132 is pulled down by UrduSpeech and mahwiz.
+- A wrong language id can cross Urdu and Sindhi, which share a Perso-Arabic script.
+- No word n-gram was used in this table. Fusion can still overwrite a correct acoustic hypothesis.
+
+## Citation
+
+```bibtex
+@software{canopy_m_2026,
+  title = {Canopy-M: Conformer-CTC speech recognition for Urdu and Sindhi},
+  year = {2026},
+  url = {https://github.com/mwzkhalil/canopy},
+  license = {Apache-2.0}
+}
+```
+
+Also cite Common Voice, FLEURS, and the Metro-ASR / Conformer designs the block follows.

canopy/__init__.py

@@ -0,0 +1,6 @@
+"""Canopy: multilingual non-autoregressive CTC ASR for Pakistani languages."""
+
+__version__ = "0.1.0.dev0"
+
+BLANK_ID = 0
+SAMPLE_RATE = 16000

canopy/data/__init__.py

@@ -0,0 +1,3 @@
+from canopy.data.manifest import PROVENANCES, REQUIRED_FIELDS, load_manifest, make_row, validate_row
+
+__all__ = ["PROVENANCES", "REQUIRED_FIELDS", "load_manifest", "make_row", "validate_row"]

canopy/data/commonvoice.py

@@ -0,0 +1,113 @@
+"""Import a locally downloaded Common Voice locale into a Canopy manifest.
+
+Common Voice is obtained by each user from Mozilla Data Collective under the
+platform terms, which forbid re-hosting any part of the dataset and forbid
+attempting to identify speakers. This importer only reads local files. The
+CV ``client_id`` is re-hashed into an opaque ``speaker_id`` used solely for
+split construction.
+"""
+
+from __future__ import annotations
+
+import csv
+import hashlib
+import os
+import sys
+from pathlib import Path
+
+from canopy.data.manifest import make_row
+
+CV_LICENSE = "CC0-1.0 (Mozilla Data Collective terms apply: no redistribution)"
+CV_SOURCE = "cv_scripted_speech"
+
+# Which TSV files map to which validation status. `invalidated.tsv` is never imported.
+CV_TABLES = {"validated": "validated.tsv", "other": "other.tsv"}
+
+csv.field_size_limit(min(sys.maxsize, 2**31 - 1))
+
+
+def _speaker_hash(client_id: str) -> str:
+    return "cv:" + hashlib.sha256(client_id.encode("utf-8")).hexdigest()[:16]
+
+
+def _read_tsv(path: Path) -> list[dict]:
+    with open(path, encoding="utf-8", newline="") as f:
+        return list(csv.DictReader(f, delimiter="\t", quoting=csv.QUOTE_NONE))
+
+
+def read_clip_durations(locale_dir: Path) -> dict[str, float]:
+    p = locale_dir / "clip_durations.tsv"
+    if not p.exists():
+        return {}
+    out = {}
+    for row in _read_tsv(p):
+        clip = row.get("clip")
+        dur = row.get("duration[ms]") or row.get("duration")
+        if clip and dur:
+            out[clip] = float(dur) / 1000.0
+    return out
+
+
+def import_locale(locale_dir: str | os.PathLike, *, language: str, version: str, tables: tuple[str, ...] = ("validated",),
+                  probe_missing_durations: bool = True, audio_root: str | os.PathLike | None = None) -> tuple[list[dict], dict]:
+    """Return (rows, stats) for one CV locale directory (the one containing ``clips/``)."""
+    from canopy.data.generic import check_importable
+
+    check_importable(CV_SOURCE)
+    locale_dir = Path(locale_dir)
+    clips_dir = Path(audio_root) if audio_root else locale_dir / "clips"
+    durations = read_clip_durations(locale_dir)
+    stats = {"locale_dir": str(locale_dir), "language": language, "version": version, "tables": {},
+             "missing_audio": 0, "missing_duration": 0, "empty_text": 0}
+    rows: list[dict] = []
+    seen: set[str] = set()
+    for table in tables:
+        if table not in CV_TABLES:
+            raise ValueError(f"unknown CV table {table!r}; choose from {list(CV_TABLES)} ('invalidated' is never imported)")
+        path = locale_dir / CV_TABLES[table]
+        if not path.exists():
+            raise FileNotFoundError(f"{path} not found; is {locale_dir} an extracted CV locale directory?")
+        n = 0
+        for rec in _read_tsv(path):
+            clip = rec["path"]
+            if clip in seen:
+                continue
+            seen.add(clip)
+            text = (rec.get("sentence") or "").strip()
+            if not text:
+                stats["empty_text"] += 1
+                continue
+            audio = clips_dir / clip
+            if not audio.exists():
+                stats["missing_audio"] += 1
+                continue
+            dur = durations.get(clip)
+            if dur is None and probe_missing_durations:
+                try:
+                    from canopy.frontend.audio import audio_info
+
+                    dur = audio_info(audio).duration
+                except Exception:
+                    dur = None
+            if not dur:
+                stats["missing_duration"] += 1
+                continue
+            extra = {k: rec[k] for k in ("sentence_id", "variant", "accents", "segment") if rec.get(k)}
+            extra["cv_locale"] = rec.get("locale") or locale_dir.name
+            rows.append(make_row(
+                id=f"cv{version}:{language}:{Path(clip).stem}",
+                audio_path=str(audio.resolve()),
+                duration=dur,
+                language=language,
+                speaker_id=_speaker_hash(rec["client_id"]),
+                text_raw=text,
+                source=CV_SOURCE,
+                source_version=version,
+                license=CV_LICENSE,
+                provenance="human",
+                validation="validated" if table == "validated" else "unvalidated",
+                extra=extra,
+            ))
+            n += 1
+        stats["tables"][table] = n
+    return rows, stats

canopy/data/dataset.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+
+import torch
+from torch.utils.data import Dataset
+
+from canopy.frontend.audio import load_audio
+from canopy.frontend.augment import speed_perturb
+
+
+def ctc_required_frames(tokens: list[int]) -> int:
+    """Minimum CTC frames: one per token plus a blank between adjacent repeats."""
+    repeats = sum(1 for a, b in zip(tokens, tokens[1:]) if a == b)
+    return len(tokens) + repeats
+
+
+@dataclass
+class FilterStats:
+    kept: int = 0
+    empty_target: int = 0
+    ctc_infeasible: int = 0
+    unknown_language: int = 0
+
+    def as_dict(self) -> dict:
+        return dict(self.__dict__)
+
+
+class ASRDataset(Dataset):
+    """Utterance dataset. Features are computed on the device in the trainer.
+
+    Rows whose target cannot be emitted by CTC at the encoder frame rate (even
+    at the fastest speed perturbation) are removed up front and counted.
+    """
+
+    def __init__(self, rows: list[dict], tokenizer, lang2id: dict[str, int], *, encoder_fps: float = 25.0,
+                 speed_factors: tuple[float, ...] = (1.0,), train: bool = False, filter_targets: bool = True):
+        self.tokenizer = tokenizer
+        self.lang2id = lang2id
+        self.train = train
+        self.speed_factors = tuple(speed_factors) if train else (1.0,)
+        self.stats = FilterStats()
+        fastest = max(self.speed_factors)
+        self.rows: list[dict] = []
+        self.targets: list[list[int]] = []
+        for r in rows:
+            if r["language"] not in lang2id:
+                self.stats.unknown_language += 1
+                continue
+            ids = tokenizer.encode(r["text_train"], language=r["language"])
+            if filter_targets:
+                if not ids:
+                    self.stats.empty_target += 1
+                    continue
+                frames = int(r["duration"] / fastest * encoder_fps)
+                if ctc_required_frames(ids) > frames:
+                    self.stats.ctc_infeasible += 1
+                    continue
+            self.rows.append(r)
+            self.targets.append(ids)
+        self.stats.kept = len(self.rows)
+
+    def __len__(self) -> int:
+        return len(self.rows)
+
+    def duration(self, i: int) -> float:
+        return self.rows[i]["duration"]
+
+    def __getitem__(self, i: int) -> dict:
+        r = self.rows[i]
+        wav = load_audio(r["audio_path"])
+        if self.train and len(self.speed_factors) > 1:
+            f = random.choice(self.speed_factors)
+            wav = speed_perturb(wav, f)
+        return {
+            "id": r["id"],
+            "wav": wav,
+            "tokens": torch.tensor(self.targets[i], dtype=torch.long),
+            "lang_id": self.lang2id[r["language"]],
+            "language": r["language"],
+            "speaker_id": r["speaker_id"],
+        }
+
+
+def collate(batch: list[dict]) -> dict:
+    B = len(batch)
+    wav_lens = torch.tensor([b["wav"].numel() for b in batch], dtype=torch.long)
+    tgt_lens = torch.tensor([b["tokens"].numel() for b in batch], dtype=torch.long)
+    wavs = torch.zeros(B, int(wav_lens.max()))
+    tgts = torch.zeros(B, int(tgt_lens.max()), dtype=torch.long)
+    for i, b in enumerate(batch):
+        wavs[i, : b["wav"].numel()] = b["wav"]
+        tgts[i, : b["tokens"].numel()] = b["tokens"]
+    return {
+        "ids": [b["id"] for b in batch],
+        "languages": [b["language"] for b in batch],
+        "speaker_ids": [b["speaker_id"] for b in batch],
+        "wav": wavs,
+        "wav_lengths": wav_lens,
+        "targets": tgts,
+        "target_lengths": tgt_lens,
+        "lang_ids": torch.tensor([b["lang_id"] for b in batch], dtype=torch.long),
+    }

canopy/data/generic.py

@@ -0,0 +1,75 @@
+"""Import a corpus described by a simple TSV (path, text, speaker[, duration][, split]).
+
+Used for corpora other than Common Voice (e.g. FLEURS exported locally). The
+caller must supply the dataset id, version and licence exactly as recorded in
+data/LICENSES.json; the importer refuses datasets whose release status there is
+blocked or excluded.
+"""
+
+from __future__ import annotations
+
+import csv
+import json
+import os
+from pathlib import Path
+
+from canopy.data.manifest import make_row
+
+BLOCKED_STATUSES = {"blocked_pending_review", "excluded", "eval_only_pending_review"}
+DEFAULT_LICENSES = Path(__file__).resolve().parents[2] / "data" / "LICENSES.json"
+
+
+def license_record(dataset_id: str, licenses_path: str | os.PathLike = DEFAULT_LICENSES) -> dict:
+    with open(licenses_path, encoding="utf-8") as f:
+        db = json.load(f)
+    datasets = db["datasets"]
+    if isinstance(datasets, dict):
+        rec = datasets.get(dataset_id)
+    else:
+        rec = next((d for d in datasets if d.get("id") == dataset_id), None)
+    if rec is None:
+        raise KeyError(f"dataset {dataset_id!r} is not in {licenses_path}; audit its licence before importing")
+    return rec
+
+
+def check_importable(dataset_id: str, licenses_path: str | os.PathLike = DEFAULT_LICENSES) -> dict:
+    rec = license_record(dataset_id, licenses_path)
+    status = rec.get("release_status")
+    if status in BLOCKED_STATUSES:
+        raise PermissionError(f"dataset {dataset_id!r} has release_status={status!r} in {licenses_path}; "
+                              "it must not be used until the licence review is resolved")
+    return rec
+
+
+def import_tsv(tsv: str | os.PathLike, *, dataset_id: str, language: str, version: str, audio_root: str | os.PathLike | None = None,
+               licenses_path: str | os.PathLike = DEFAULT_LICENSES, provenance: str = "human") -> list[dict]:
+    rec = check_importable(dataset_id, licenses_path)
+    license_str = rec.get("license", "unknown")
+    tsv = Path(tsv)
+    root = Path(audio_root) if audio_root else tsv.parent
+    rows = []
+    with open(tsv, encoding="utf-8", newline="") as f:
+        for i, r in enumerate(csv.DictReader(f, delimiter="\t", quoting=csv.QUOTE_NONE)):
+            audio = Path(r["path"])
+            if not audio.is_absolute():
+                audio = root / audio
+            dur = r.get("duration")
+            if not dur:
+                from canopy.frontend.audio import audio_info
+
+                dur = audio_info(audio).duration
+            extra = {"split": r["split"]} if r.get("split") else None
+            rows.append(make_row(
+                id=f"{dataset_id}:{language}:{r.get('id') or i}",
+                audio_path=str(audio.resolve()),
+                duration=float(dur),
+                language=language,
+                speaker_id=f"{dataset_id}:{r['speaker']}",
+                text_raw=r["text"],
+                source=dataset_id,
+                source_version=version,
+                license=license_str,
+                provenance=provenance,
+                extra=extra,
+            ))
+    return rows

canopy/data/manifest.py

@@ -0,0 +1,124 @@
+"""Manifest schema.
+
+A manifest is a JSONL file with one utterance per line. Manifests reference
+audio on the local disk and are never committed or published with text:
+Common Voice terms forbid re-hosting any part of the corpus.
+"""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Iterable
+
+from canopy.text.normalize import get_normalizer
+from canopy.utils.io import read_jsonl, sha256_text
+
+REQUIRED_FIELDS = (
+    "id", "audio_path", "duration", "language", "speaker_id", "text_raw", "text_train", "text_eval",
+    "sentence_key", "source", "source_version", "license", "provenance",
+)
+# `human`: human transcript; `pseudo`: model-generated label; `synthetic`: synthetic audio (e.g. TTS).
+PROVENANCES = ("human", "pseudo", "synthetic")
+SPLITS = ("train", "dev", "test")
+
+
+def sentence_key(language: str, text_eval: str) -> str:
+    """Identity of a sentence for leakage control: same language + same scoring text."""
+    return sha256_text(f"{language}\t{text_eval}")[:20]
+
+
+def make_row(*, id: str, audio_path: str, duration: float, language: str, speaker_id: str, text_raw: str,
+             source: str, source_version: str, license: str, provenance: str = "human",
+             validation: str | None = None, extra: dict | None = None) -> dict:
+    if provenance not in PROVENANCES:
+        raise ValueError(f"provenance must be one of {PROVENANCES}, got {provenance!r}")
+    norm = get_normalizer(language)
+    text_train = norm.train_norm(text_raw)
+    text_eval = norm.eval_norm(text_raw)
+    row = {
+        "id": id,
+        "audio_path": audio_path,
+        "duration": round(float(duration), 4),
+        "language": language,
+        "speaker_id": speaker_id,
+        "text_raw": text_raw,
+        "text_train": text_train,
+        "text_eval": text_eval,
+        "sentence_key": sentence_key(language, text_eval),
+        "source": source,
+        "source_version": source_version,
+        "license": license,
+        "provenance": provenance,
+    }
+    if validation is not None:
+        row["validation"] = validation
+    if extra:
+        row.update(extra)
+    return row
+
+
+def validate_row(row: dict) -> list[str]:
+    errors = [f"missing field {k!r}" for k in REQUIRED_FIELDS if k not in row]
+    if errors:
+        return errors
+    if row["provenance"] not in PROVENANCES:
+        errors.append(f"bad provenance {row['provenance']!r}")
+    if "split" in row and row["split"] not in SPLITS and not str(row["split"]).startswith("fold"):
+        errors.append(f"bad split {row['split']!r}")
+    if not isinstance(row["duration"], (int, float)) or row["duration"] <= 0:
+        errors.append("non-positive duration")
+    return errors
+
+
+def load_manifest(paths: str | os.PathLike | Iterable[str | os.PathLike], *, languages: Iterable[str] | None = None,
+                  splits: Iterable[str] | None = None, provenances: Iterable[str] | None = ("human",),
+                  include_unvalidated: bool = False,
+                  min_duration: float = 0.0, max_duration: float | None = None) -> list[dict]:
+    """Load and filter one or more manifests.
+
+    ``provenances`` defaults to human-labelled data only; pseudo-labelled or
+    synthetic rows must be requested explicitly. Rows whose audio failed or
+    lacks community validation (``validation == "unvalidated"``) are excluded
+    unless ``include_unvalidated`` is set.
+    """
+    if isinstance(paths, (str, os.PathLike)):
+        paths = [paths]
+    langs = set(languages) if languages else None
+    spl = set(splits) if splits else None
+    prov = set(provenances) if provenances else None
+    rows = []
+    for p in paths:
+        for row in read_jsonl(p):
+            if langs is not None and row["language"] not in langs:
+                continue
+            if spl is not None and row.get("split") not in spl:
+                continue
+            if prov is not None and row.get("provenance") not in prov:
+                continue
+            if not include_unvalidated and row.get("validation") == "unvalidated":
+                continue
+            d = row["duration"]
+            if d < min_duration or (max_duration is not None and d > max_duration):
+                continue
+            rows.append(row)
+    return rows
+
+
+def summarize(rows: Iterable[dict]) -> dict:
+    out: dict = {}
+    for r in rows:
+        key = (r["language"], r.get("split", "unsplit"))
+        s = out.setdefault(key, {"utterances": 0, "hours": 0.0, "speakers": set(), "sentences": set()})
+        s["utterances"] += 1
+        s["hours"] += r["duration"] / 3600
+        s["speakers"].add(r["speaker_id"])
+        s["sentences"].add(r["sentence_key"])
+    result: dict = {}
+    for (lang, split), s in sorted(out.items()):
+        result.setdefault(lang, {})[split] = {
+            "utterances": s["utterances"],
+            "hours": round(s["hours"], 4),
+            "speakers": len(s["speakers"]),
+            "sentences": len(s["sentences"]),
+        }
+    return result

canopy/data/sampler.py

@@ -0,0 +1,128 @@
+"""Duration-bucketed batching and temperature-balanced multilingual sampling.
+
+Language ``l`` is drawn with probability ``p_l ∝ min(h_l, cap_l) ** (1/T)``
+where ``h_l`` is its training hours. T=1 is proportional sampling, T→∞ is
+uniform. ``cap_l`` bounds the effective size of dominant languages (Pashto has
+>3,000 validated hours against ~10 h for most CV-Pakistan languages).
+Each virtual epoch draws ``epoch_hours`` of audio in total.
+"""
+
+from __future__ import annotations
+
+import math
+import random
+from collections import defaultdict
+from collections.abc import Iterator, Sequence
+
+from torch.utils.data import Sampler
+
+
+def language_probabilities(hours: dict[str, float], temperature: float, caps: dict[str, float] | None = None,
+                           default_cap: float | None = None) -> dict[str, float]:
+    caps = caps or {}
+    eff = {}
+    for lang, h in hours.items():
+        cap = caps.get(lang, default_cap)
+        eff[lang] = min(h, cap) if cap is not None else h
+    if math.isinf(temperature):
+        w = {lang: 1.0 for lang in eff}
+    else:
+        w = {lang: h ** (1.0 / temperature) for lang, h in eff.items() if h > 0}
+    z = sum(w.values())
+    return {lang: v / z for lang, v in w.items()}
+
+
+def _pack(indices: list[int], durations: Sequence[float], max_batch_seconds: float, max_batch_size: int) -> list[list[int]]:
+    batches, cur, cur_max = [], [], 0.0
+    for i in indices:
+        d = durations[i]
+        new_max = max(cur_max, d)
+        if cur and (new_max * (len(cur) + 1) > max_batch_seconds or len(cur) >= max_batch_size):
+            batches.append(cur)
+            cur, new_max = [], d
+        cur.append(i)
+        cur_max = new_max
+    if cur:
+        batches.append(cur)
+    return batches
+
+
+class DurationBatchSampler(Sampler[list[int]]):
+    """Deterministic, sorted batches for evaluation (padding-efficient)."""
+
+    def __init__(self, durations: Sequence[float], max_batch_seconds: float, max_batch_size: int = 256):
+        order = sorted(range(len(durations)), key=lambda i: durations[i])
+        self.batches = _pack(order, durations, max_batch_seconds, max_batch_size)
+
+    def __iter__(self) -> Iterator[list[int]]:
+        return iter(self.batches)
+
+    def __len__(self) -> int:
+        return len(self.batches)
+
+
+class BalancedBatchSampler(Sampler[list[int]]):
+    def __init__(self, languages: Sequence[str], durations: Sequence[float], *, max_batch_seconds: float,
+                 max_batch_size: int = 128, temperature: float = 5.0, caps_hours: dict[str, float] | None = None,
+                 default_cap_hours: float | None = None, epoch_hours: float | None = None, bucket_size: int = 64,
+                 seed: int = 0):
+        self.languages = list(languages)
+        self.durations = list(durations)
+        self.max_batch_seconds = max_batch_seconds
+        self.max_batch_size = max_batch_size
+        self.bucket_size = bucket_size
+        self.seed = seed
+        self.by_lang: dict[str, list[int]] = defaultdict(list)
+        for i, lang in enumerate(self.languages):
+            self.by_lang[lang].append(i)
+        self.hours = {lang: sum(self.durations[i] for i in idx) / 3600 for lang, idx in self.by_lang.items()}
+        self.probs = language_probabilities(self.hours, temperature, caps_hours, default_cap_hours)
+        total_h = sum(self.hours.values())
+        self.epoch_hours = epoch_hours if epoch_hours is not None else total_h
+        self.epoch = 0
+        self.start_batch = 0
+        self._batches: list[list[int]] = []
+        self._audit: dict = {}
+        self.set_epoch(0)
+
+    def set_epoch(self, epoch: int) -> None:
+        self.epoch = epoch
+        rng = random.Random(f"{self.seed}:{epoch}")
+        drawn: list[int] = []
+        audit = {}
+        for lang in sorted(self.by_lang):
+            p = self.probs.get(lang, 0.0)
+            budget = p * self.epoch_hours * 3600
+            pool = self.by_lang[lang]
+            got, secs, cycles = [], 0.0, 0
+            order: list[int] = []
+            while secs < budget and pool:
+                if not order:
+                    order = pool[:]
+                    rng.shuffle(order)
+                    cycles += 1
+                i = order.pop()
+                got.append(i)
+                secs += self.durations[i]
+            drawn.extend(got)
+            audit[lang] = {"probability": round(p, 6), "utterances": len(got), "seconds": round(secs, 2),
+                           "passes_over_data": round(secs / max(1e-9, self.hours[lang] * 3600), 3)}
+        rng.shuffle(drawn)
+        batches: list[list[int]] = []
+        chunk = self.bucket_size * max(1, self.max_batch_size // 4)
+        for s in range(0, len(drawn), chunk):
+            part = sorted(drawn[s : s + chunk], key=lambda i: self.durations[i])
+            batches.extend(_pack(part, self.durations, self.max_batch_seconds, self.max_batch_size))
+        rng.shuffle(batches)
+        self._batches = batches
+        self._audit = {"epoch": epoch, "languages": audit, "batches": len(batches)}
+
+    def audit(self) -> dict:
+        return self._audit
+
+    def __iter__(self) -> Iterator[list[int]]:
+        start, self.start_batch = self.start_batch, 0
+        return iter(self._batches[start:])
+
+    def __len__(self) -> int:
+        return len(self._batches)

canopy/data/splits.py

@@ -0,0 +1,188 @@
+"""Speaker- AND sentence-disjoint split construction.
+
+Common Voice re-records each sentence by several speakers, so speaker-disjoint
+splits alone leak transcripts from train into test. The procedure, per
+language:
+
+1. Assign whole speakers to ``test``, then ``dev``, until each reaches its
+   target share of hours and minimum speaker count; the rest go to ``train``.
+2. Drop dev rows whose sentence occurs in test, and train rows whose sentence
+   occurs in dev or test. Dropped rows are returned and counted, never
+   silently discarded.
+
+Languages with too few speakers for a fixed held-out set get k speaker folds
+instead (``fold`` field); evaluation then cross-validates over folds.
+"""
+
+from __future__ import annotations
+
+import random
+from collections import defaultdict
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+
+
+@dataclass
+class SplitConfig:
+    test_fraction: float = 0.10
+    dev_fraction: float = 0.10
+    min_test_speakers: int = 3
+    min_dev_speakers: int = 2
+    # One speaker may provide at most this share of a held-out split's target hours
+    # (unless no other speaker is left).
+    max_speaker_share: float = 0.34
+    # Below this many speakers a language gets speaker folds instead of fixed splits.
+    min_speakers_for_fixed: int = 10
+    n_folds: int = 5
+    seed: int = 1234
+
+
+@dataclass
+class SplitResult:
+    rows: list[dict]
+    dropped: list[dict]
+    report: dict = field(default_factory=dict)
+
+
+def _by(rows: Iterable[dict], key: str) -> dict[str, list[dict]]:
+    out: dict[str, list[dict]] = defaultdict(list)
+    for r in rows:
+        out[r[key]].append(r)
+    return out
+
+
+def _hours(rows: Iterable[dict]) -> float:
+    return sum(r["duration"] for r in rows) / 3600
+
+
+def _pick_speakers(speakers: list[str], spk_hours: dict[str, float], target_h: float, min_spk: int,
+                   max_share: float) -> list[str]:
+    chosen: list[str] = []
+    total = 0.0
+    cap = max_share * target_h
+    for s in list(speakers):
+        if total >= target_h and len(chosen) >= min_spk:
+            break
+        if spk_hours[s] > cap and len(speakers) - len(chosen) > min_spk:
+            continue
+        chosen.append(s)
+        total += spk_hours[s]
+    if len(chosen) < min_spk:
+        for s in speakers:
+            if s not in chosen:
+                chosen.append(s)
+                if len(chosen) >= min_spk:
+                    break
+    return chosen
+
+
+def _fixed_split(rows: list[dict], cfg: SplitConfig, rng: random.Random) -> tuple[list[dict], list[dict], dict]:
+    spk_rows = _by(rows, "speaker_id")
+    spk_hours = {s: _hours(rs) for s, rs in spk_rows.items()}
+    speakers = sorted(spk_rows)
+    rng.shuffle(speakers)
+    total_h = _hours(rows)
+
+    test_spk = _pick_speakers(speakers, spk_hours, cfg.test_fraction * total_h, cfg.min_test_speakers, cfg.max_speaker_share)
+    rest = [s for s in speakers if s not in set(test_spk)]
+    dev_spk = _pick_speakers(rest, spk_hours, cfg.dev_fraction * total_h, cfg.min_dev_speakers, cfg.max_speaker_share)
+    assign = {s: "train" for s in speakers}
+    assign.update({s: "test" for s in test_spk})
+    assign.update({s: "dev" for s in dev_spk})
+
+    test_sent = {r["sentence_key"] for r in rows if assign[r["speaker_id"]] == "test"}
+    dev_sent = {r["sentence_key"] for r in rows if assign[r["speaker_id"]] == "dev"}
+    kept, dropped = [], []
+    for r in rows:
+        split = assign[r["speaker_id"]]
+        leak = (split == "dev" and r["sentence_key"] in test_sent) or (
+            split == "train" and (r["sentence_key"] in test_sent or r["sentence_key"] in dev_sent))
+        out = dict(r, split=split)
+        if leak:
+            out["drop_reason"] = "sentence_in_heldout"
+            dropped.append(out)
+        else:
+            kept.append(out)
+    return kept, dropped, {"mode": "fixed", "test_speakers": len(test_spk), "dev_speakers": len(dev_spk)}
+
+
+def _fold_split(rows: list[dict], cfg: SplitConfig, rng: random.Random) -> tuple[list[dict], list[dict], dict]:
+    spk_rows = _by(rows, "speaker_id")
+    speakers = sorted(spk_rows)
+    rng.shuffle(speakers)
+    speakers.sort(key=lambda s: -_hours(spk_rows[s]))
+    k = max(2, min(cfg.n_folds, len(speakers)))
+    load = [0.0] * k
+    fold_of: dict[str, int] = {}
+    for s in speakers:
+        i = min(range(k), key=lambda j: load[j])
+        fold_of[s] = i
+        load[i] += _hours(spk_rows[s])
+    kept = [dict(r, split=f"fold{fold_of[r['speaker_id']]}", fold=fold_of[r["speaker_id"]]) for r in rows]
+    return kept, [], {"mode": "speaker_folds", "n_folds": k, "fold_hours": [round(x, 3) for x in load]}
+
+
+def make_splits(rows: list[dict], cfg: SplitConfig | None = None, respect_existing: bool = True) -> SplitResult:
+    """Assign ``split`` to every row. Rows that already carry a split (e.g. FLEURS
+    official dev/test) keep it when ``respect_existing`` is set."""
+    cfg = cfg or SplitConfig()
+    fixed_rows = [r for r in rows if respect_existing and r.get("split")]
+    todo = [r for r in rows if not (respect_existing and r.get("split"))]
+    out_rows, dropped = list(fixed_rows), []
+    report: dict = {"config": cfg.__dict__, "languages": {}}
+    for lang, lrows in sorted(_by(todo, "language").items()):
+        rng = random.Random(f"{cfg.seed}:{lang}")
+        n_spk = len({r["speaker_id"] for r in lrows})
+        if n_spk < cfg.min_speakers_for_fixed:
+            kept, drop, info = _fold_split(lrows, cfg, rng)
+        else:
+            kept, drop, info = _fixed_split(lrows, cfg, rng)
+        out_rows.extend(kept)
+        dropped.extend(drop)
+        info.update({
+            "speakers": n_spk,
+            "input_hours": round(_hours(lrows), 4),
+            "dropped_rows": len(drop),
+            "dropped_hours": round(_hours(drop), 4),
+            "splits": {s: {"rows": len(rs), "hours": round(_hours(rs), 4), "speakers": len({r['speaker_id'] for r in rs})}
+                       for s, rs in sorted(_by(kept, "split").items())},
+        })
+        report["languages"][lang] = info
+    report["leakage"] = check_disjoint(out_rows)
+    return SplitResult(rows=out_rows, dropped=dropped, report=report)
+
+
+def check_disjoint(rows: Iterable[dict]) -> dict:
+    """Count speaker and sentence overlaps between splits, per language. All counts must be 0."""
+    per: dict[str, dict[str, dict[str, set]]] = defaultdict(lambda: defaultdict(lambda: {"spk": set(), "sent": set()}))
+    for r in rows:
+        split = r.get("split")
+        if split is None:
+            continue
+        group = "fold" if str(split).startswith("fold") else "fixed"
+        per[(r["language"], group)][split]["spk"].add(r["speaker_id"])
+        per[(r["language"], group)][split]["sent"].add(r["sentence_key"])
+    out: dict = {"speaker_overlaps": 0, "sentence_overlaps": 0, "details": []}
+    for (lang, group), splits in per.items():
+        names = sorted(splits)
+        for i, a in enumerate(names):
+            for b in names[i + 1:]:
+                s = len(splits[a]["spk"] & splits[b]["spk"])
+                t = len(splits[a]["sent"] & splits[b]["sent"])
+                # Folds share sentences by design; the train view drops them (fold_view).
+                if group == "fold":
+                    t = 0
+                if s or t:
+                    out["details"].append({"language": lang, "a": a, "b": b, "speakers": s, "sentences": t})
+                out["speaker_overlaps"] += s
+                out["sentence_overlaps"] += t
+    return out
+
+
+def fold_view(rows: list[dict], language: str, fold: int) -> tuple[list[dict], list[dict]]:
+    """(train, test) for one speaker fold, with test sentences removed from train."""
+    lrows = [r for r in rows if r["language"] == language and "fold" in r]
+    test = [r for r in lrows if r["fold"] == fold]
+    test_sent = {r["sentence_key"] for r in test}
+    train = [r for r in lrows if r["fold"] != fold and r["sentence_key"] not in test_sent]
+    return train, test

canopy/data/validation.py

@@ -0,0 +1,193 @@
+"""Automated data-quality checks producing ``data_quality_report.json``.
+
+Examples in the report are utterance IDs only (no transcripts), so the report
+can be shared without redistributing corpus text.
+
+Speaker near-duplicate detection is deliberately NOT implemented for Common
+Voice: its terms forbid attempts to determine speaker identity.
+"""
+
+from __future__ import annotations
+
+import unicodedata
+from collections import Counter, defaultdict
+from collections.abc import Callable
+
+from canopy.data.manifest import validate_row
+from canopy.data.splits import check_disjoint
+from canopy.languages import LanguageRegistry, get_registry
+
+MAX_EXAMPLES = 20
+
+
+def char_script(ch: str) -> str:
+    if ch.isspace():
+        return "space"
+    cat = unicodedata.category(ch)
+    if cat == "Nd":
+        return "digit"
+    if cat[0] == "M":
+        return "mark"
+    name = unicodedata.name(ch, "")
+    if name.startswith("ARABIC"):
+        return "Arab"
+    if name.startswith("LATIN"):
+        return "Latn"
+    if name.startswith(("DEVANAGARI",)):
+        return "Deva"
+    if name.startswith(("GURMUKHI",)):
+        return "Guru"
+    if name.startswith(("TIBETAN",)):
+        return "Tibt"
+    if ch in "'\u2019\u02bc\u200c":
+        return "joiner_or_apostrophe"
+    return "other"
+
+
+class _Issues:
+    def __init__(self):
+        self.counts: dict[str, Counter] = defaultdict(Counter)
+        self.examples: dict[str, dict[str, list[str]]] = defaultdict(lambda: defaultdict(list))
+
+    def add(self, lang: str, issue: str, utt_id: str):
+        self.counts[lang][issue] += 1
+        ex = self.examples[lang][issue]
+        if len(ex) < MAX_EXAMPLES:
+            ex.append(utt_id)
+
+
+def validate_manifest(rows: list[dict], *, registry: LanguageRegistry | None = None, check_audio: bool = False,
+                      expected_sr: int = 16000, min_duration: float = 0.3, max_duration: float = 30.0,
+                      max_chars_per_second: float = 30.0, min_chars_per_second: float = 0.3,
+                      token_counter: Callable[[str], tuple[int, int]] | None = None, frames_per_second: float = 25.0) -> dict:
+    """Return the data-quality report dict.
+
+    ``token_counter(text) -> (n_tokens, n_adjacent_repeats)`` enables the CTC
+    feasibility check at ``frames_per_second`` (encoder output rate).
+    """
+    registry = registry or get_registry()
+    issues = _Issues()
+    ids = Counter(r.get("id") for r in rows)
+    audio_hashes: dict[str, str] = {}
+    sent_counts: dict[str, Counter] = defaultdict(Counter)
+    char_scripts: dict[str, Counter] = defaultdict(Counter)
+
+    for r in rows:
+        lang = r.get("language", "<none>")
+        uid = r.get("id", "<no-id>")
+        for e in validate_row(r):
+            issues.add(lang, f"schema:{e}", uid)
+        if ids[uid] > 1:
+            issues.add(lang, "duplicate_id", uid)
+        if lang not in registry:
+            issues.add(lang, "unknown_language_code", uid)
+            allowed_scripts = set()
+        else:
+            L = registry[lang]
+            allowed_scripts = {L.primary_script, *L.alternate_scripts}
+        text = r.get("text_train", "")
+        raw = r.get("text_raw", "")
+        dur = float(r.get("duration", 0) or 0)
+        if not text.strip():
+            issues.add(lang, "empty_transcript", uid)
+        if dur < min_duration:
+            issues.add(lang, "too_short", uid)
+        if dur > max_duration:
+            issues.add(lang, "too_long", uid)
+        if dur > 0 and text:
+            cps = len(text) / dur
+            if cps > max_chars_per_second:
+                issues.add(lang, "impossible_ratio_too_many_chars", uid)
+            elif cps < min_chars_per_second:
+                issues.add(lang, "impossible_ratio_too_few_chars", uid)
+        if token_counter is not None and text and dur > 0:
+            n_tok, n_rep = token_counter(text)
+            if n_tok + n_rep > int(dur * frames_per_second):
+                issues.add(lang, "ctc_infeasible", uid)
+        sent_counts[lang][r.get("sentence_key")] += 1
+
+        for ch in raw:
+            cat = unicodedata.category(ch)
+            o = ord(ch)
+            if cat in ("Co", "Cn", "Cs"):
+                issues.add(lang, "unicode_unassigned_or_private", uid)
+            elif 0xFB50 <= o <= 0xFDFF or 0xFE70 <= o <= 0xFEFE:
+                issues.add(lang, "unicode_presentation_form_in_raw", uid)
+            elif cat == "Cf" and ch != "\u200c":
+                issues.add(lang, "unicode_invisible_format_char_in_raw", uid)
+        word_scripts_mixed = False
+        for word in text.split():
+            s = {char_script(c) for c in word} - {"mark", "digit", "joiner_or_apostrophe"}
+            if len(s) > 1:
+                word_scripts_mixed = True
+        if word_scripts_mixed:
+            issues.add(lang, "mixed_script_word", uid)
+        for ch in text:
+            sc = char_script(ch)
+            char_scripts[lang][sc] += 1
+            if sc in ("Arab", "Latn", "Deva", "Guru", "Tibt") and allowed_scripts and sc not in allowed_scripts:
+                issues.add(lang, f"char_outside_registered_scripts:{sc}", uid)
+            if sc == "other":
+                issues.add(lang, "unsupported_character", uid)
+
+        if check_audio:
+            _check_audio(r, lang, uid, issues, expected_sr, audio_hashes)
+
+    per_lang: dict[str, dict] = {}
+    for lang in sorted({r.get("language", "<none>") for r in rows}):
+        lrows = [r for r in rows if r.get("language") == lang]
+        sc = sent_counts[lang]
+        per_lang[lang] = {
+            "utterances": len(lrows),
+            "hours": round(sum(r["duration"] for r in lrows) / 3600, 4),
+            "speakers": len({r["speaker_id"] for r in lrows}),
+            "unique_sentences": len(sc),
+            "max_recordings_per_sentence": max(sc.values()) if sc else 0,
+            "mean_recordings_per_sentence": round(len(lrows) / max(1, len(sc)), 3),
+            "duplicate_transcript_rows": sum(c - 1 for c in sc.values() if c > 1),
+            "char_scripts": dict(char_scripts[lang]),
+            "issues": dict(issues.counts.get(lang, {})),
+            "issue_examples": {k: v for k, v in issues.examples.get(lang, {}).items()},
+        }
+    leakage = check_disjoint(rows)
+    total_issues = sum(sum(c.values()) for c in issues.counts.values())
+    return {
+        "n_utterances": len(rows),
+        "checks": {
+            "audio_checked": check_audio,
+            "near_duplicate_speakers": "not performed: Common Voice terms prohibit attempts to identify speakers",
+            "thresholds": {"min_duration": min_duration, "max_duration": max_duration,
+                           "max_chars_per_second": max_chars_per_second, "min_chars_per_second": min_chars_per_second,
+                           "ctc_frames_per_second": frames_per_second if token_counter else None},
+        },
+        "speaker_leakage": {"speaker_overlaps": leakage["speaker_overlaps"], "details": leakage["details"]},
+        "sentence_leakage": {"sentence_overlaps": leakage["sentence_overlaps"]},
+        "total_issue_count": total_issues,
+        "languages": per_lang,
+    }
+
+
+def _check_audio(r: dict, lang: str, uid: str, issues: _Issues, expected_sr: int, audio_hashes: dict[str, str]):
+    from canopy.frontend.audio import audio_info
+    from canopy.utils.io import sha256_file
+
+    path = r.get("audio_path")
+    try:
+        info = audio_info(path)
+    except Exception:
+        issues.add(lang, "broken_audio", uid)
+        return
+    if info.frames == 0:
+        issues.add(lang, "broken_audio", uid)
+        return
+    if info.sample_rate != expected_sr:
+        issues.add(lang, "sample_rate_not_16k", uid)
+    if info.channels != 1:
+        issues.add(lang, "not_mono", uid)
+    if abs(info.duration - float(r["duration"])) > 0.1:
+        issues.add(lang, "duration_mismatch", uid)
+    h = sha256_file(path)
+    if h in audio_hashes:
+        issues.add(lang, "duplicate_audio", uid)
+    else:
+        audio_hashes[h] = uid

canopy/decoding/__init__.py

@@ -0,0 +1,4 @@
+from canopy.decoding.beam import prefix_beam_search
+from canopy.decoding.greedy import ctc_collapse, greedy_decode, language_mask_bias
+
+__all__ = ["ctc_collapse", "greedy_decode", "language_mask_bias", "prefix_beam_search"]

canopy/decoding/beam.py

@@ -0,0 +1,132 @@
+"""CTC prefix beam search with optional word-level n-gram shallow fusion.
+
+score(prefix) = log P_ctc(prefix) + alpha * log P_lm(words) + beta * |words|
+
+A word is scored by the LM when it is completed, i.e. when the next emitted
+token string contains a space (character tokenizers emit " "; SentencePiece
+pieces that start a word begin with "▁", rendered as " ").
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+import numpy as np
+
+NEG_INF = -math.inf
+
+
+def _lse(a: float, b: float) -> float:
+    if a == NEG_INF:
+        return b
+    if b == NEG_INF:
+        return a
+    m = max(a, b)
+    return m + math.log(math.exp(a - m) + math.exp(b - m))
+
+
+class _LMState:
+    __slots__ = ("raw", "history", "lm", "words")
+
+    def __init__(self, raw: str, history: tuple[str, ...], lm: float, words: int):
+        self.raw, self.history, self.lm, self.words = raw, history, lm, words
+
+
+def _pending_word(raw: str) -> str | None:
+    if not raw or raw.endswith(" "):
+        return None
+    parts = raw.split()
+    return parts[-1] if parts else None
+
+
+def prefix_beam_search(log_probs: np.ndarray, token_strings: Sequence[str], *, beam: int = 10, lm=None,
+                       alpha: float = 0.5, beta: float = 1.0, token_prune: int = 20, prune_logp: float = -12.0,
+                       allowed: np.ndarray | None = None, blank: int = 0) -> list[int]:
+    """Return the best token-id sequence for one utterance.
+
+    ``log_probs``: (T, V) natural-log posteriors. ``token_strings[i]``: surface
+    string of token i ("" for blank). ``lm``: object with
+    ``score(history, word) -> ln prob``, ``score_end(history) -> ln prob`` and
+    ``order``. ``allowed``: optional bool (V,) output mask.
+    """
+    T, V = log_probs.shape
+    lp = log_probs.copy()
+    if allowed is not None:
+        lp[:, ~allowed] = NEG_INF
+    use_lm = lm is not None and alpha != 0.0
+    hist_len = max(0, (lm.order - 1)) if use_lm else 0
+
+    beams: dict[tuple[int, ...], tuple[float, float]] = {(): (0.0, NEG_INF)}
+    lm_states: dict[tuple[int, ...], _LMState] = {(): _LMState("", (), 0.0, 0)}
+
+    def lm_state(prefix: tuple[int, ...]) -> _LMState:
+        st = lm_states.get(prefix)
+        if st is not None:
+            return st
+        parent = lm_state(prefix[:-1])
+        s = token_strings[prefix[-1]]
+        raw = parent.raw + s
+        history, lm_score, words = parent.history, parent.lm, parent.words
+        if " " in s:
+            w = _pending_word(parent.raw)
+            if w is not None:
+                words += 1
+                if use_lm:
+                    lm_score += lm.score(history, w)
+                    history = (history + (w,))[-hist_len:] if hist_len else ()
+        st = _LMState(raw, history, lm_score, words)
+        lm_states[prefix] = st
+        return st
+
+    def total(prefix, pb, pnb):
+        st = lm_state(prefix)
+        return _lse(pb, pnb) + alpha * st.lm + beta * st.words
+
+    for t in range(T):
+        row = lp[t]
+        cand = np.argsort(row)[::-1][:token_prune]
+        cand = [int(c) for c in cand if row[c] > prune_logp or c == blank]
+        nxt: dict[tuple[int, ...], list[float]] = {}
+
+        def add(p, b, nb, nxt=nxt):
+            cur = nxt.get(p)
+            if cur is None:
+                nxt[p] = [b, nb]
+            else:
+                cur[0] = _lse(cur[0], b)
+                cur[1] = _lse(cur[1], nb)
+
+        for prefix, (pb, pnb) in beams.items():
+            for c in cand:
+                p = float(row[c])
+                if c == blank:
+                    add(prefix, _lse(pb, pnb) + p, NEG_INF)
+                    continue
+                last = prefix[-1] if prefix else None
+                ext = prefix + (c,)
+                if c == last:
+                    add(ext, NEG_INF, pb + p)
+                    add(prefix, NEG_INF, pnb + p)
+                else:
+                    add(ext, NEG_INF, _lse(pb, pnb) + p)
+        scored = sorted(nxt.items(), key=lambda kv: total(kv[0], kv[1][0], kv[1][1]), reverse=True)[:beam]
+        beams = {k: (v[0], v[1]) for k, v in scored}
+
+    best, best_score = (), NEG_INF
+    for prefix, (pb, pnb) in beams.items():
+        st = lm_state(prefix)
+        lm_score, words = st.lm, st.words
+        w = _pending_word(st.raw)
+        history = st.history
+        if w is not None:
+            words += 1
+            if use_lm:
+                lm_score += lm.score(history, w)
+                history = (history + (w,))[-hist_len:] if hist_len else ()
+        if use_lm:
+            lm_score += lm.score_end(history)
+        s = _lse(pb, pnb) + alpha * lm_score + beta * words
+        if s > best_score:
+            best, best_score = prefix, s
+    return list(best)

canopy/decoding/greedy.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import torch
+
+from canopy.tokenizer.base import BLANK_ID
+
+
+def ctc_collapse(ids: Sequence[int], blank: int = BLANK_ID) -> list[int]:
+    """Merge repeats, then remove blanks."""
+    out, prev = [], None
+    for i in ids:
+        if i != prev and i != blank:
+            out.append(int(i))
+        prev = i
+    return out
+
+
+def language_mask_bias(tokenizer, languages: Sequence[str | None], vocab_size: int, device=None) -> torch.Tensor | None:
+    """(B, V) additive bias: 0 for allowed tokens, -inf for tokens never seen in that
+    language's training text (output design O2). Rows for ``None`` or unknown
+    languages are all zeros (unrestricted)."""
+    if not any(languages):
+        return None
+    bias = torch.zeros(len(languages), vocab_size, device=device)
+    for b, lang in enumerate(languages):
+        if not lang:
+            continue
+        mask = tokenizer.language_mask(lang)
+        if mask is None:
+            continue
+        m = torch.tensor(mask, dtype=torch.bool, device=device)
+        bias[b, ~m] = float("-inf")
+    return bias
+
+
+@torch.no_grad()
+def greedy_decode(log_probs: torch.Tensor, lengths: torch.Tensor, mask_bias: torch.Tensor | None = None) -> list[list[int]]:
+    if mask_bias is not None:
+        log_probs = log_probs + mask_bias.unsqueeze(1)
+    best = log_probs.argmax(dim=-1).cpu()
+    return [ctc_collapse(best[b, : int(lengths[b])].tolist()) for b in range(best.shape[0])]

canopy/evaluation/__init__.py

@@ -0,0 +1,3 @@
+from canopy.evaluation.metrics import aggregate, corpus_error_rate, edit_counts, score_utterances
+
+__all__ = ["aggregate", "corpus_error_rate", "edit_counts", "score_utterances"]

canopy/evaluation/metrics.py

@@ -0,0 +1,135 @@
+"""WER / CER with per-language, macro, weighted and worst-k aggregates and
+speaker-cluster bootstrap confidence intervals.
+
+Both reference and hypothesis are scored in the language's ``eval_norm`` view.
+CER counts characters of the eval_norm string including single spaces between
+words.
+"""
+
+from __future__ import annotations
+
+import random
+from collections import defaultdict
+from collections.abc import Sequence
+from dataclasses import dataclass
+
+
+@dataclass
+class EditCounts:
+    sub: int = 0
+    dele: int = 0
+    ins: int = 0
+    ref_len: int = 0
+
+    @property
+    def errors(self) -> int:
+        return self.sub + self.dele + self.ins
+
+    def __add__(self, o: EditCounts) -> EditCounts:
+        return EditCounts(self.sub + o.sub, self.dele + o.dele, self.ins + o.ins, self.ref_len + o.ref_len)
+
+
+def edit_counts(ref: Sequence, hyp: Sequence) -> EditCounts:
+    """Levenshtein alignment counts (substitutions, deletions, insertions)."""
+    n, m = len(ref), len(hyp)
+    if n == 0:
+        return EditCounts(0, 0, m, 0)
+    if m == 0:
+        return EditCounts(0, n, 0, n)
+    # Each cell holds (cost, sub, del, ins); ties prefer substitution, then deletion.
+    prev = [(j, 0, 0, j) for j in range(m + 1)]
+    for i in range(1, n + 1):
+        cur = [(i, 0, i, 0)]
+        ri = ref[i - 1]
+        for j in range(1, m + 1):
+            d = prev[j - 1]
+            if ri == hyp[j - 1]:
+                best = d
+            else:
+                best = (d[0] + 1, d[1] + 1, d[2], d[3])
+            u = prev[j]
+            if u[0] + 1 < best[0]:
+                best = (u[0] + 1, u[1], u[2] + 1, u[3])
+            left = cur[j - 1]
+            if left[0] + 1 < best[0]:
+                best = (left[0] + 1, left[1], left[2], left[3] + 1)
+            cur.append(best)
+        prev = cur
+    c = prev[m]
+    return EditCounts(c[1], c[2], c[3], n)
+
+
+def corpus_error_rate(counts: Sequence[EditCounts]) -> float:
+    total = EditCounts()
+    for c in counts:
+        total = total + c
+    return total.errors / total.ref_len if total.ref_len else 0.0
+
+
+def score_utterances(items: Sequence[dict]) -> list[dict]:
+    """``items``: dicts with language, speaker_id, ref (eval_norm), hyp (eval_norm)."""
+    out = []
+    for it in items:
+        w = edit_counts(it["ref"].split(), it["hyp"].split())
+        c = edit_counts(list(it["ref"]), list(it["hyp"]))
+        out.append({**it, "word": w, "char": c})
+    return out
+
+
+def _bootstrap(scored: list[dict], unit: str, n_boot: int, seed: int) -> tuple[float, float] | None:
+    by_spk: dict[str, list[EditCounts]] = defaultdict(list)
+    for s in scored:
+        by_spk[s["speaker_id"]].append(s[unit])
+    speakers = list(by_spk)
+    if len(speakers) < 2:
+        return None
+    totals = {k: sum(v, EditCounts()) for k, v in by_spk.items()}
+    rng = random.Random(seed)
+    vals = []
+    for _ in range(n_boot):
+        e = r = 0
+        for _k in range(len(speakers)):
+            t = totals[speakers[rng.randrange(len(speakers))]]
+            e += t.errors
+            r += t.ref_len
+        vals.append(e / r if r else 0.0)
+    vals.sort()
+    return vals[int(0.025 * (n_boot - 1))], vals[int(0.975 * (n_boot - 1))]
+
+
+def aggregate(scored: list[dict], n_boot: int = 1000, seed: int = 0, worst_k: int = 5) -> dict:
+    per_lang: dict[str, list[dict]] = defaultdict(list)
+    for s in scored:
+        per_lang[s["language"]].append(s)
+    langs = {}
+    for lang, rows in sorted(per_lang.items()):
+        w = sum((r["word"] for r in rows), EditCounts())
+        c = sum((r["char"] for r in rows), EditCounts())
+        entry = {
+            "utterances": len(rows),
+            "speakers": len({r["speaker_id"] for r in rows}),
+            "ref_words": w.ref_len,
+            "ref_chars": c.ref_len,
+            "wer": w.errors / w.ref_len if w.ref_len else 0.0,
+            "cer": c.errors / c.ref_len if c.ref_len else 0.0,
+            "word_edits": {"sub": w.sub, "del": w.dele, "ins": w.ins},
+            "char_edits": {"sub": c.sub, "del": c.dele, "ins": c.ins},
+        }
+        if n_boot:
+            entry["wer_ci95"] = _bootstrap(rows, "word", n_boot, seed)
+            entry["cer_ci95"] = _bootstrap(rows, "char", n_boot, seed)
+        entry["insufficient_evidence"] = entry["speakers"] < 5
+        langs[lang] = entry
+    all_w = sum((s["word"] for s in scored), EditCounts())
+    all_c = sum((s["char"] for s in scored), EditCounts())
+    macro_cer = sum(v["cer"] for v in langs.values()) / len(langs) if langs else 0.0
+    macro_wer = sum(v["wer"] for v in langs.values()) / len(langs) if langs else 0.0
+    worst = sorted(langs.items(), key=lambda kv: -kv[1]["cer"])[:worst_k]
+    return {
+        "languages": langs,
+        "macro": {"cer": macro_cer, "wer": macro_wer, "n_languages": len(langs)},
+        "weighted": {"cer": all_c.errors / all_c.ref_len if all_c.ref_len else 0.0,
+                     "wer": all_w.errors / all_w.ref_len if all_w.ref_len else 0.0},
+        f"worst_{worst_k}_by_cer": [{"language": k, "cer": v["cer"], "wer": v["wer"]} for k, v in worst],
+        "bootstrap": {"unit": "speaker", "n": n_boot, "seed": seed} if n_boot else None,
+    }

canopy/evaluation/runner.py

@@ -0,0 +1,98 @@
+"""Run a model over manifest rows and score the output."""
+
+from __future__ import annotations
+
+import time
+from contextlib import nullcontext
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader
+
+from canopy.data.dataset import ASRDataset, collate
+from canopy.data.sampler import DurationBatchSampler
+from canopy.decoding.beam import prefix_beam_search
+from canopy.decoding.greedy import greedy_decode, language_mask_bias
+from canopy.evaluation.metrics import aggregate, score_utterances
+from canopy.text.normalize import get_normalizer
+
+
+@torch.no_grad()
+def transcribe_rows(model, tokenizer, rows: list[dict], *, lang2id: dict[str, int], device: torch.device,
+                    max_batch_seconds: float = 300.0, num_workers: int = 0, use_language: bool = True,
+                    use_masks: bool = True, decoder: str = "greedy", beam: int = 10, lms: dict | None = None,
+                    alpha: float = 0.5, beta: float = 1.0, amp_dtype: torch.dtype | None = None) -> list[dict]:
+    """Returns one dict per row with id, language, speaker_id, hyp_train (model output text).
+
+    ``use_language``: pass the reference language to the model (conditioning)
+    and to the output mask. With ``use_language=False`` the model runs
+    language-agnostic, as for ``language=None`` at inference.
+    """
+    model.eval()
+    ds = ASRDataset(rows, tokenizer, {**lang2id, **{r["language"]: 0 for r in rows if r["language"] not in lang2id}},
+                    filter_targets=False)
+    sampler = DurationBatchSampler([r["duration"] for r in ds.rows], max_batch_seconds)
+    dl = DataLoader(ds, batch_sampler=sampler, collate_fn=collate, num_workers=num_workers)
+    token_strings = None
+    if decoder == "beam":
+        labels = tokenizer.labels_for_pyctcdecode()
+        token_strings = [""] + labels[1:]
+    out = []
+    for batch in dl:
+        wav = batch["wav"].to(device)
+        wl = batch["wav_lengths"].to(device)
+        lids = batch["lang_ids"].to(device) if use_language else None
+        ctx = torch.autocast(device.type, dtype=amp_dtype) if amp_dtype is not None else nullcontext()
+        with ctx:
+            res = model(wav, wl, lids)
+        lp = res["log_probs"]
+        langs = batch["languages"] if (use_language and use_masks) else [None] * len(batch["ids"])
+        bias = language_mask_bias(tokenizer, langs, lp.shape[-1], device=lp.device)
+        if decoder == "greedy":
+            seqs = greedy_decode(lp, res["lengths"], bias)
+        elif decoder == "beam":
+            seqs = []
+            lpn = lp.float().cpu().numpy()
+            for b in range(lp.shape[0]):
+                allowed = None
+                if bias is not None:
+                    allowed = np.isfinite(bias[b].cpu().numpy())
+                lm = (lms or {}).get(batch["languages"][b]) if use_language else (lms or {}).get("*")
+                seqs.append(prefix_beam_search(lpn[b, : int(res["lengths"][b])], token_strings, beam=beam, lm=lm,
+                                               alpha=alpha, beta=beta, allowed=allowed))
+        else:
+            raise ValueError(f"unknown decoder {decoder!r}")
+        blank_rate = (lp.argmax(-1) == 0).float()
+        for b, seq in enumerate(seqs):
+            L = int(res["lengths"][b])
+            out.append({
+                "id": batch["ids"][b],
+                "language": batch["languages"][b],
+                "speaker_id": batch["speaker_ids"][b],
+                "hyp_train": tokenizer.decode(seq),
+                "frames": L,
+                "blank_rate": float(blank_rate[b, :L].mean()),
+            })
+    return out
+
+
+def evaluate_rows(model, tokenizer, rows: list[dict], *, n_boot: int = 1000, **kw) -> tuple[dict, list[dict]]:
+    by_id = {r["id"]: r for r in rows}
+    t0 = time.perf_counter()
+    hyps = transcribe_rows(model, tokenizer, rows, **kw)
+    elapsed = time.perf_counter() - t0
+    items = []
+    for h in hyps:
+        r = by_id[h["id"]]
+        norm = get_normalizer(r["language"])
+        items.append({"id": h["id"], "language": r["language"], "speaker_id": r["speaker_id"],
+                      "ref": r["text_eval"], "hyp": norm.eval_norm(h["hyp_train"])})
+    scored = score_utterances(items)
+    metrics = aggregate(scored, n_boot=n_boot)
+    audio_s = sum(r["duration"] for r in rows)
+    metrics["decode"] = {"seconds": elapsed, "audio_seconds": audio_s, "rtf": elapsed / max(audio_s, 1e-9),
+                         "mean_blank_rate": float(np.mean([h["blank_rate"] for h in hyps])) if hyps else None}
+    for h, it in zip(hyps, items):
+        h["ref_eval"] = it["ref"]
+        h["hyp_eval"] = it["hyp"]
+    return metrics, hyps

canopy/evaluation/status.py

@@ -0,0 +1,79 @@
+"""Language support status from evidence (RESEARCH_PLAN.md section 12.3).
+
+Statuses are computed, never assigned by hand. Thresholds are the proposed
+values and must be ratified before release 0.2.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Thresholds:
+    prod_train_hours: float = 50.0
+    prod_test_speakers: int = 20
+    prod_test_hours: float = 2.0
+    prod_cer: float = 0.10
+    prod_wer: float = 0.25
+    prod_min_buckets: int = 2
+    prod_ci_halfwidth: float = 0.02
+    prod_int8_shift: float = 0.01
+    sup_train_hours: float = 10.0
+    sup_test_speakers: int = 5
+    sup_test_hours: float = 0.5
+    sup_cer: float = 0.15
+    sup_wer: float = 0.40
+    limited_train_hours: float = 5.0
+    limited_speakers: int = 10
+
+
+def compute_status(*, has_clean_data: bool, evaluated: bool, train_hours: float, total_speakers: int,
+                   test_speakers: int = 0, test_hours: float = 0.0, cer: float | None = None, wer: float | None = None,
+                   cer_ci95: tuple[float, float] | None = None, buckets_passing: int = 0,
+                   int8_cer_shift: float | None = None, native_review_done: bool = False,
+                   second_seed_within_ci: bool = False, auxiliary: bool = False, th: Thresholds = Thresholds()) -> dict:
+    reasons: list[str] = []
+    if auxiliary:
+        return {"status": "auxiliary", "reasons": ["auxiliary training language"]}
+    if not has_clean_data:
+        return {"status": "not-currently-supported", "reasons": ["no licence-clean transcribed data"]}
+    if train_hours < th.limited_train_hours or total_speakers < th.limited_speakers:
+        return {"status": "data-limited",
+                "reasons": [f"train_hours={train_hours:.2f} (<{th.limited_train_hours}) or "
+                            f"speakers={total_speakers} (<{th.limited_speakers})"]}
+    if not evaluated or cer is None or wer is None:
+        return {"status": "evaluation-pending", "reasons": ["not yet evaluated under the Canopy protocol"]}
+
+    ci_half = (cer_ci95[1] - cer_ci95[0]) / 2 if cer_ci95 else None
+    prod_checks = {
+        "train_hours": train_hours >= th.prod_train_hours,
+        "test_speakers": test_speakers >= th.prod_test_speakers,
+        "test_hours": test_hours >= th.prod_test_hours,
+        "cer": cer <= th.prod_cer,
+        "wer": wer <= th.prod_wer,
+        "buckets": buckets_passing >= th.prod_min_buckets,
+        "ci": ci_half is not None and ci_half <= th.prod_ci_halfwidth,
+        "int8": int8_cer_shift is not None and int8_cer_shift <= th.prod_int8_shift,
+        "native_review": native_review_done,
+        "second_seed": second_seed_within_ci,
+    }
+    if all(prod_checks.values()):
+        return {"status": "production", "reasons": ["all production criteria met"], "checks": prod_checks}
+    sup_checks = {
+        "train_hours": train_hours >= th.sup_train_hours,
+        "test_speakers": test_speakers >= th.sup_test_speakers,
+        "test_hours": test_hours >= th.sup_test_hours,
+        "cer": cer <= th.sup_cer,
+        "wer": wer <= th.sup_wer,
+        "ci_reported": cer_ci95 is not None,
+        "second_seed": second_seed_within_ci,
+    }
+    if all(sup_checks.values()):
+        failed = [k for k, v in prod_checks.items() if not v]
+        return {"status": "supported", "reasons": [f"production criteria not met: {failed}"], "checks": sup_checks}
+    failed = [k for k, v in sup_checks.items() if not v]
+    reasons.append(f"supported criteria not met: {failed}")
+    if test_speakers < th.sup_test_speakers:
+        reasons.append("insufficient evidence: fewer than 5 test speakers")
+    return {"status": "experimental", "reasons": reasons, "checks": sup_checks}

canopy/frontend/__init__.py

@@ -0,0 +1,5 @@
+from canopy.frontend.audio import load_audio, resample
+from canopy.frontend.augment import SpecAugment, speed_perturb
+from canopy.frontend.features import LogMelFrontend, feature_lengths, mel_filterbank
+
+__all__ = ["LogMelFrontend", "SpecAugment", "feature_lengths", "load_audio", "mel_filterbank", "resample", "speed_perturb"]

canopy/frontend/audio.py

@@ -0,0 +1,69 @@
+"""Audio loading and resampling. All model inputs are 16 kHz mono float32."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from canopy import SAMPLE_RATE
+
+
+@dataclass
+class AudioInfo:
+    sample_rate: int
+    channels: int
+    frames: int
+
+    @property
+    def duration(self) -> float:
+        return self.frames / self.sample_rate if self.sample_rate else 0.0
+
+
+def audio_info(path: str | os.PathLike) -> AudioInfo:
+    import soundfile as sf
+
+    i = sf.info(str(path))
+    return AudioInfo(sample_rate=int(i.samplerate), channels=int(i.channels), frames=int(i.frames))
+
+
+def _read(path: str) -> tuple[np.ndarray, int]:
+    try:
+        import soundfile as sf
+
+        data, sr = sf.read(path, dtype="float32", always_2d=True)
+        return data, int(sr)
+    except Exception as sf_err:
+        try:
+            import torchaudio
+
+            wav, sr = torchaudio.load(path)
+            return wav.numpy().T.astype(np.float32), int(sr)
+        except Exception as ta_err:
+            raise RuntimeError(
+                f"could not decode {path!r} (soundfile: {sf_err}; torchaudio: {ta_err}). "
+                "For MP3 input use libsndfile >= 1.1 (soundfile >= 0.12) or convert with scripts/prepare_data.py --convert."
+            ) from ta_err
+
+
+def resample(wav: torch.Tensor, orig_sr: int, target_sr: int = SAMPLE_RATE) -> torch.Tensor:
+    if orig_sr == target_sr:
+        return wav
+    import torchaudio.functional as taf
+
+    return taf.resample(wav, orig_sr, target_sr)
+
+
+def load_audio(path: str | os.PathLike, target_sr: int = SAMPLE_RATE) -> torch.Tensor:
+    """Load an audio file as a 1-D float32 tensor at ``target_sr``, downmixed to mono."""
+    data, sr = _read(str(path))
+    wav = torch.from_numpy(np.ascontiguousarray(data.mean(axis=1)))
+    return resample(wav, sr, target_sr)
+
+
+def save_audio(path: str | os.PathLike, wav: torch.Tensor, sr: int = SAMPLE_RATE) -> None:
+    import soundfile as sf
+
+    sf.write(str(path), wav.detach().cpu().numpy().astype(np.float32), sr)

canopy/frontend/augment.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+import random
+
+import torch
+from torch import nn
+
+from canopy import SAMPLE_RATE
+
+
+def speed_perturb(wav: torch.Tensor, factor: float, sample_rate: int = SAMPLE_RATE) -> torch.Tensor:
+    """SoX-style speed change (tempo and pitch together) by resampling."""
+    if factor == 1.0:
+        return wav
+    import torchaudio.functional as taf
+
+    src = int(round(sample_rate * factor))
+    return taf.resample(wav, src, sample_rate)
+
+
+class SpecAugment(nn.Module):
+    """Batched SpecAugment on (B, T, F) features.
+
+    Time-mask width is adaptive: at most ``time_mask_ratio`` of each
+    utterance's valid length, because Canopy clips average ~4 s (400 frames at
+    100 fps), where a fixed 50-100 frame mask would erase whole words.
+    """
+
+    def __init__(self, freq_masks: int = 2, freq_width: int = 27, time_masks: int = 2,
+                 time_mask_ratio: float = 0.05, min_time_width: int = 1):
+        super().__init__()
+        self.freq_masks = freq_masks
+        self.freq_width = freq_width
+        self.time_masks = time_masks
+        self.time_mask_ratio = time_mask_ratio
+        self.min_time_width = min_time_width
+
+    def forward(self, feats: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
+        if not self.training:
+            return feats
+        B, T, F = feats.shape
+        mask = torch.zeros(B, T, F, dtype=torch.bool, device=feats.device)
+        for b in range(B):
+            L = int(lengths[b])
+            for _ in range(self.freq_masks):
+                w = random.randint(0, min(self.freq_width, F))
+                f0 = random.randint(0, F - w)
+                mask[b, :, f0 : f0 + w] = True
+            max_w = max(self.min_time_width, int(L * self.time_mask_ratio))
+            for _ in range(self.time_masks):
+                w = random.randint(0, max_w)
+                if L - w <= 0:
+                    continue
+                t0 = random.randint(0, L - w)
+                mask[b, t0 : t0 + w, :] = True
+        return feats.masked_fill(mask, 0.0)

canopy/frontend/features.py

@@ -0,0 +1,111 @@
+"""Log-mel frontend implemented in plain PyTorch (ONNX-exportable, no torchaudio needed).
+
+Defaults: 25 ms Hann window, 10 ms hop, 512-point FFT, 80 Slaney mel bins, 100 fps.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import nn
+
+from canopy import SAMPLE_RATE
+
+
+def _hz_to_mel_slaney(f: torch.Tensor) -> torch.Tensor:
+    f_sp = 200.0 / 3
+    min_log_hz = 1000.0
+    min_log_mel = min_log_hz / f_sp
+    logstep = math.log(6.4) / 27.0
+    mel = f / f_sp
+    return torch.where(f >= min_log_hz, min_log_mel + torch.log(f.clamp(min=1e-10) / min_log_hz) / logstep, mel)
+
+
+def _mel_to_hz_slaney(m: torch.Tensor) -> torch.Tensor:
+    f_sp = 200.0 / 3
+    min_log_hz = 1000.0
+    min_log_mel = min_log_hz / f_sp
+    logstep = math.log(6.4) / 27.0
+    return torch.where(m >= min_log_mel, min_log_hz * torch.exp(logstep * (m - min_log_mel)), f_sp * m)
+
+
+def mel_filterbank(n_freqs: int, n_mels: int, sample_rate: int, f_min: float = 0.0, f_max: float | None = None) -> torch.Tensor:
+    """Slaney-scale, Slaney-normalised triangular filterbank of shape (n_freqs, n_mels)."""
+    f_max = f_max if f_max is not None else sample_rate / 2
+    all_freqs = torch.linspace(0, sample_rate // 2, n_freqs, dtype=torch.float64)
+    m_pts = torch.linspace(_hz_to_mel_slaney(torch.tensor(f_min, dtype=torch.float64)).item(),
+                           _hz_to_mel_slaney(torch.tensor(f_max, dtype=torch.float64)).item(), n_mels + 2,
+                           dtype=torch.float64)
+    f_pts = _mel_to_hz_slaney(m_pts)
+    f_diff = f_pts[1:] - f_pts[:-1]
+    slopes = f_pts.unsqueeze(0) - all_freqs.unsqueeze(1)
+    down = -slopes[:, :-2] / f_diff[:-1]
+    up = slopes[:, 2:] / f_diff[1:]
+    fb = torch.clamp(torch.minimum(down, up), min=0.0)
+    enorm = 2.0 / (f_pts[2 : n_mels + 2] - f_pts[:n_mels])
+    return (fb * enorm.unsqueeze(0)).float()
+
+
+def feature_lengths(num_samples: torch.Tensor, hop_length: int = 160) -> torch.Tensor:
+    """Frame count for a centred STFT."""
+    return torch.div(num_samples, hop_length, rounding_mode="floor") + 1
+
+
+class LogMelFrontend(nn.Module):
+    def __init__(
+        self,
+        sample_rate: int = SAMPLE_RATE,
+        n_fft: int = 512,
+        win_length: int = 400,
+        hop_length: int = 160,
+        n_mels: int = 80,
+        normalize: str = "utterance",
+        eps: float = 1e-10,
+    ):
+        super().__init__()
+        if normalize not in ("utterance", "per_bin", "none"):
+            raise ValueError(f"normalize must be utterance|per_bin|none, got {normalize!r}")
+        self.sample_rate = sample_rate
+        self.n_fft = n_fft
+        self.win_length = win_length
+        self.hop_length = hop_length
+        self.n_mels = n_mels
+        self.normalize = normalize
+        self.eps = eps
+        self.register_buffer("window", torch.hann_window(win_length), persistent=False)
+        self.register_buffer("fbank", mel_filterbank(n_fft // 2 + 1, n_mels, sample_rate), persistent=False)
+
+    @torch.no_grad()
+    def forward(self, wav: torch.Tensor, lengths: torch.Tensor | None = None) -> tuple[torch.Tensor, torch.Tensor]:
+        """wav (B, S) float -> features (B, T, n_mels), lengths (B,)."""
+        with torch.autocast(device_type=wav.device.type, enabled=False):
+            return self._forward(wav, lengths)
+
+    def _forward(self, wav: torch.Tensor, lengths: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]:
+        if wav.dim() == 1:
+            wav = wav.unsqueeze(0)
+        wav = wav.float()
+        B, S = wav.shape
+        if lengths is None:
+            lengths = torch.full((B,), S, dtype=torch.long, device=wav.device)
+        spec = torch.stft(wav, n_fft=self.n_fft, hop_length=self.hop_length, win_length=self.win_length,
+                          window=self.window, center=True, pad_mode="reflect", return_complex=True)
+        power = spec.real.pow(2) + spec.imag.pow(2)  # (B, F, T)
+        mel = torch.matmul(power.transpose(1, 2), self.fbank)  # (B, T, M)
+        logmel = torch.log(mel.clamp(min=self.eps))
+        feat_len = feature_lengths(lengths, self.hop_length).clamp(max=logmel.shape[1])
+        mask = torch.arange(logmel.shape[1], device=wav.device).unsqueeze(0) < feat_len.unsqueeze(1)
+        m = mask.unsqueeze(-1).float()
+        if self.normalize == "utterance":
+            n = (feat_len.float() * self.n_mels).view(B, 1, 1)
+            mean = (logmel * m).sum(dim=(1, 2), keepdim=True) / n
+            var = (((logmel - mean) * m) ** 2).sum(dim=(1, 2), keepdim=True) / n
+            logmel = (logmel - mean) / (var.sqrt() + 1e-6)
+        elif self.normalize == "per_bin":
+            n = feat_len.float().view(B, 1, 1)
+            mean = (logmel * m).sum(dim=1, keepdim=True) / n
+            var = (((logmel - mean) * m) ** 2).sum(dim=1, keepdim=True) / n
+            logmel = (logmel - mean) / (var.sqrt() + 1e-6)
+        logmel = logmel * m
+        return logmel, feat_len

canopy/inference/__init__.py

@@ -0,0 +1,3 @@
+from canopy.inference.transcriber import Transcriber
+
+__all__ = ["Transcriber"]

canopy/inference/benchmark.py

@@ -0,0 +1,64 @@
+"""CPU real-time-factor protocol (docs/BENCHMARK_PLAN.md).
+
+Fixed durations, fixed thread count, warm-up runs excluded, median and p90
+over repeats. Synthetic noise input is used for latency only (never for
+accuracy), so the benchmark needs no dataset.
+"""
+
+from __future__ import annotations
+
+import os
+import platform
+import statistics
+import time
+
+import torch
+
+from canopy import SAMPLE_RATE
+
+
+def _rss_mb() -> float | None:
+    try:
+        import psutil
+
+        return psutil.Process(os.getpid()).memory_info().rss / 2**20
+    except Exception:
+        return None
+
+
+@torch.no_grad()
+def measure_rtf(model, durations=(1.0, 5.0, 10.0, 30.0), threads: int = 4, warmup: int = 3, repeats: int = 15,
+                device: str = "cpu", seed: int = 0) -> dict:
+    torch.set_num_threads(threads)
+    model = model.to(device).eval()
+    g = torch.Generator().manual_seed(seed)
+    results = {}
+    for d in durations:
+        n = int(d * SAMPLE_RATE)
+        wav = (0.05 * torch.randn(1, n, generator=g)).to(device)
+        lens = torch.tensor([n], device=device)
+        for _ in range(warmup):
+            model(wav, lens)
+        times = []
+        for _ in range(repeats):
+            t0 = time.perf_counter()
+            model(wav, lens)
+            if device.startswith("cuda"):
+                torch.cuda.synchronize()
+            times.append(time.perf_counter() - t0)
+        times.sort()
+        med = statistics.median(times)
+        results[f"{d:g}s"] = {"median_s": med, "p90_s": times[int(0.9 * (len(times) - 1))], "rtf_median": med / d,
+                              "rtf_p90": times[int(0.9 * (len(times) - 1))] / d}
+    return {
+        "device": device,
+        "threads": threads,
+        "warmup": warmup,
+        "repeats": repeats,
+        "torch": torch.__version__,
+        "cpu": platform.processor() or platform.machine(),
+        "platform": platform.platform(),
+        "rss_mb": _rss_mb(),
+        "parameters": sum(p.numel() for p in model.parameters()),
+        "results": results,
+    }

canopy/inference/quantize.py

@@ -0,0 +1,36 @@
+"""Dynamic INT8 quantization of Linear layers for CPU inference.
+
+INT8 must pass the per-language regression gate (docs/BENCHMARK_PLAN.md)
+before any INT8 artifact is released: evaluate FP32 and INT8 on the same dev
+set and compare per-language CER.
+"""
+
+from __future__ import annotations
+
+import copy
+
+import torch
+from torch import nn
+
+
+def quantize_dynamic_int8(model: nn.Module) -> nn.Module:
+    m = copy.deepcopy(model).cpu().eval()
+    try:
+        from torch.ao.quantization import quantize_dynamic
+    except ImportError:  # pragma: no cover
+        from torch.quantization import quantize_dynamic
+    return quantize_dynamic(m, {nn.Linear}, dtype=torch.qint8)
+
+
+def int8_regression(fp32_metrics: dict, int8_metrics: dict, max_abs_cer_shift: float = 0.01) -> dict:
+    rows = {}
+    worst = 0.0
+    for lang, m in fp32_metrics["languages"].items():
+        q = int8_metrics["languages"].get(lang)
+        if q is None:
+            continue
+        shift = q["cer"] - m["cer"]
+        worst = max(worst, shift)
+        rows[lang] = {"fp32_cer": m["cer"], "int8_cer": q["cer"], "shift": shift, "pass": shift <= max_abs_cer_shift}
+    return {"max_abs_cer_shift": max_abs_cer_shift, "worst_shift": worst,
+            "passed": all(r["pass"] for r in rows.values()), "languages": rows}

canopy/inference/transcriber.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+import json
+import tempfile
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from canopy import SAMPLE_RATE
+from canopy.decoding.beam import prefix_beam_search
+from canopy.decoding.greedy import greedy_decode, language_mask_bias
+from canopy.frontend.audio import load_audio
+from canopy.tokenizer.base import load_tokenizer
+from canopy.training.checkpoint import load_checkpoint, model_from_checkpoint
+
+
+def _tokenizer_from_checkpoint(ck: dict, fallback_dir: str | None = None):
+    if "tokenizer_payload" in ck:
+        tmp = Path(tempfile.mkdtemp(prefix="canopy_tok_"))
+        (tmp / "tokenizer.json").write_text(json.dumps(ck["tokenizer_payload"], ensure_ascii=False), encoding="utf-8")
+        if ck.get("spm_model"):
+            (tmp / "spm.model").write_bytes(ck["spm_model"])
+        return load_tokenizer(tmp)
+    return load_tokenizer(fallback_dir or ck["tokenizer_dir"])
+
+
+class Transcriber:
+    """Load a checkpoint and transcribe audio.
+
+    ``language``: ISO 639-3 code. When given, the model is conditioned on it
+    (if trained with a language embedding), the output is restricted to that
+    language's tokens, and that language's LM is used. ``None`` runs
+    language-agnostic decoding.
+    """
+
+    def __init__(self, model, tokenizer, device: str = "cpu", max_chunk_seconds: float = 30.0):
+        self.model = model.to(device).eval()
+        self.tokenizer = tokenizer
+        self.device = torch.device(device)
+        self.languages = list(model.cfg.languages)
+        self.lang2id = {lang: i + 1 for i, lang in enumerate(self.languages)}
+        self.max_chunk = int(max_chunk_seconds * SAMPLE_RATE)
+        self.lms: dict[str, object] = {}
+
+    @classmethod
+    def from_checkpoint(cls, path: str | Path, device: str = "cpu", tokenizer_dir: str | None = None, **kw) -> Transcriber:
+        ck = load_checkpoint(path)
+        model = model_from_checkpoint(ck)
+        return cls(model, _tokenizer_from_checkpoint(ck, tokenizer_dir), device=device, **kw)
+
+    @classmethod
+    def from_pretrained(cls, repo_id: str = "ProximaAI/Canopy-M", *, filename: str = "canopy-m.pt",
+                        device: str = "cpu", revision: str | None = None, token: str | None = None, **kw) -> Transcriber:
+        """Download an inference checkpoint from the Hugging Face Hub and load it."""
+        try:
+            from huggingface_hub import hf_hub_download
+        except ImportError as exc:
+            raise ImportError("pip install huggingface_hub") from exc
+        path = hf_hub_download(repo_id=repo_id, filename=filename, revision=revision, token=token)
+        return cls.from_checkpoint(path, device=device, **kw)
+
+    def attach_lm(self, language: str, lm) -> None:
+        """``lm``: ArpaLM or KenLMScorer. Use language "*" for a language-agnostic LM."""
+        self.lms[language] = lm
+
+    @torch.no_grad()
+    def log_probs(self, wav: torch.Tensor, language: str | None = None) -> torch.Tensor:
+        lid = None
+        if language is not None and self.model.lang_embed is not None:
+            lid = torch.tensor([self.lang2id.get(language, 0)], device=self.device)
+        chunks = [wav[i : i + self.max_chunk] for i in range(0, max(1, wav.numel()), self.max_chunk)]
+        outs = []
+        for c in chunks:
+            if c.numel() < 400:
+                continue
+            x = c.unsqueeze(0).to(self.device)
+            res = self.model(x, torch.tensor([c.numel()], device=self.device), lid)
+            outs.append(res["log_probs"][0, : int(res["lengths"][0])])
+        if not outs:
+            return torch.zeros(0, self.tokenizer.vocab_size)
+        return torch.cat(outs, dim=0)
+
+    def transcribe(self, audio, language: str | None = None, decoder: str = "greedy", beam: int = 10,
+                   alpha: float = 0.5, beta: float = 1.0) -> str:
+        wav = load_audio(audio) if isinstance(audio, (str, Path)) else torch.as_tensor(audio, dtype=torch.float32)
+        lp = self.log_probs(wav, language)
+        if lp.shape[0] == 0:
+            return ""
+        bias = language_mask_bias(self.tokenizer, [language], lp.shape[-1], device=lp.device) if language else None
+        if decoder == "greedy":
+            ids = greedy_decode(lp.unsqueeze(0), torch.tensor([lp.shape[0]]), bias)[0]
+        elif decoder == "beam":
+            labels = self.tokenizer.labels_for_pyctcdecode()
+            allowed = np.isfinite(bias[0].cpu().numpy()) if bias is not None else None
+            lm = self.lms.get(language or "*") or self.lms.get("*")
+            ids = prefix_beam_search(lp.float().cpu().numpy(), [""] + labels[1:], beam=beam, lm=lm, alpha=alpha,
+                                     beta=beta, allowed=allowed)
+        else:
+            raise ValueError(f"unknown decoder {decoder!r}")
+        return self.tokenizer.decode(ids)

canopy/language_model/__init__.py

@@ -0,0 +1,4 @@
+from canopy.language_model.arpa import ArpaLM, build_arpa, write_arpa
+from canopy.language_model.text import filter_heldout, lm_sentences
+
+__all__ = ["ArpaLM", "build_arpa", "filter_heldout", "lm_sentences", "write_arpa"]

canopy/language_model/arpa.py

@@ -0,0 +1,188 @@
+"""Word n-gram LMs in ARPA format: a pure-Python reader/scorer and a small
+builder (interpolated absolute discounting, written as a back-off ARPA).
+
+For large corpora build with KenLM's ``lmplz`` (scripts/train_lm.py
+--backend kenlm); the pure-Python builder is intended for the small
+per-language CV-Pakistan corpora and for platforms without KenLM binaries.
+"""
+
+from __future__ import annotations
+
+import gzip
+import math
+from collections import Counter, defaultdict
+from collections.abc import Iterable
+from pathlib import Path
+
+LN10 = math.log(10.0)
+BOS, EOS, UNK = "<s>", "</s>", "<unk>"
+
+
+class ArpaLM:
+    def __init__(self, order: int, probs: dict[tuple[str, ...], float], backoffs: dict[tuple[str, ...], float],
+                 unk_logp10: float = -7.0):
+        self.order = order
+        self.probs = probs
+        self.backoffs = backoffs
+        self.unk_logp10 = probs.get((UNK,), unk_logp10)
+
+    @classmethod
+    def load(cls, path: str | Path) -> ArpaLM:
+        path = Path(path)
+        opener = gzip.open if path.suffix == ".gz" else open
+        probs: dict[tuple[str, ...], float] = {}
+        backoffs: dict[tuple[str, ...], float] = {}
+        order = 0
+        section = None
+        with opener(path, "rt", encoding="utf-8") as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                if line == "\\data\\":
+                    section = "data"
+                    continue
+                if line == "\\end\\":
+                    break
+                if line.startswith("\\") and line.endswith("-grams:"):
+                    section = int(line[1:].split("-")[0])
+                    order = max(order, section)
+                    continue
+                if section == "data" or section is None:
+                    continue
+                parts = line.split("\t")
+                if len(parts) == 1:
+                    parts = line.split()
+                    n = section
+                    logp = float(parts[0])
+                    words = tuple(parts[1 : 1 + n])
+                    bo = float(parts[1 + n]) if len(parts) > 1 + n else None
+                else:
+                    logp = float(parts[0])
+                    words = tuple(parts[1].split())
+                    bo = float(parts[2]) if len(parts) > 2 else None
+                probs[words] = logp
+                if bo is not None:
+                    backoffs[words] = bo
+        return cls(order, probs, backoffs)
+
+    def _logp10(self, context: tuple[str, ...], word: str) -> float:
+        context = context[-(self.order - 1):] if self.order > 1 else ()
+        bo_total = 0.0
+        while True:
+            key = context + (word,)
+            if key in self.probs:
+                return bo_total + self.probs[key]
+            if not context:
+                return bo_total + self.unk_logp10
+            bo_total += self.backoffs.get(context, 0.0)
+            context = context[1:]
+
+    def score(self, history: tuple[str, ...], word: str) -> float:
+        """Natural-log P(word | <s> history)."""
+        return self._logp10((BOS,) + tuple(history), word) * LN10
+
+    def score_end(self, history: tuple[str, ...]) -> float:
+        return self._logp10((BOS,) + tuple(history), EOS) * LN10
+
+    def sentence_logprob(self, words: list[str]) -> float:
+        h: list[str] = []
+        total = 0.0
+        for w in words:
+            total += self.score(tuple(h), w)
+            h.append(w)
+        return total + self.score_end(tuple(h))
+
+
+def build_arpa(sentences: Iterable[str], order: int = 3, discount: float = 0.75, unk_mass: float = 1e-3) -> tuple[dict, dict, dict]:
+    """Interpolated absolute-discounting n-gram model.
+
+    Returns (probs_log10, backoffs_log10, counts_per_order) for ``write_arpa``.
+    """
+    counts: list[Counter] = [Counter() for _ in range(order + 1)]
+    for s in sentences:
+        toks = [BOS] + s.split() + [EOS]
+        for n in range(1, order + 1):
+            for i in range(len(toks) - n + 1):
+                ng = tuple(toks[i : i + n])
+                if n == 1 and ng == (BOS,):
+                    continue
+                counts[n][ng] += 1
+
+    # Unigrams: maximum likelihood, with a fixed mass reserved for <unk>.
+    total1 = sum(counts[1].values())
+    p: dict[tuple[str, ...], float] = {}
+    for ng, c in counts[1].items():
+        p[ng] = (1 - unk_mass) * c / total1
+    p[(UNK,)] = unk_mass
+    p[(BOS,)] = 0.0
+
+    ctx_counts: list[dict] = [dict() for _ in range(order + 1)]
+    followers: list[dict] = [defaultdict(int) for _ in range(order + 1)]
+    for n in range(2, order + 1):
+        cc: dict[tuple[str, ...], int] = defaultdict(int)
+        for ng, c in counts[n].items():
+            cc[ng[:-1]] += c
+            followers[n][ng[:-1]] += 1
+        ctx_counts[n] = cc
+
+    def lower(ng: tuple[str, ...]) -> float:
+        if len(ng) == 1:
+            return p.get(ng, p[(UNK,)])
+        return p.get(ng) if ng in p else lower_interp(ng)
+
+    def lower_interp(ng: tuple[str, ...]) -> float:
+        # Probability of an unseen n-gram via back-off to its suffix.
+        ctx = ng[:-1]
+        n = len(ng)
+        if ctx not in ctx_counts[n]:
+            return lower(ng[1:])
+        lam = discount * followers[n][ctx] / ctx_counts[n][ctx]
+        return lam * lower(ng[1:])
+
+    for n in range(2, order + 1):
+        for ng, c in counts[n].items():
+            ctx = ng[:-1]
+            cc = ctx_counts[n][ctx]
+            lam = discount * followers[n][ctx] / cc
+            p[ng] = max(c - discount, 0.0) / cc + lam * lower(ng[1:])
+
+    # Back-off weights so that unseen continuations of each context sum correctly.
+    bows: dict[tuple[str, ...], float] = {}
+    for n in range(2, order + 1):
+        seen_by_ctx: dict[tuple[str, ...], list[tuple[str, ...]]] = defaultdict(list)
+        for ng in counts[n]:
+            seen_by_ctx[ng[:-1]].append(ng)
+        for ctx, ngs in seen_by_ctx.items():
+            num = 1.0 - sum(p[ng] for ng in ngs)
+            den = 1.0 - sum(lower(ng[1:]) for ng in ngs)
+            if num <= 0 or den <= 0:
+                bows[ctx] = 0.0
+            else:
+                bows[ctx] = math.log10(num / den)
+
+    probs10 = {ng: (math.log10(v) if v > 0 else -99.0) for ng, v in p.items()}
+    return probs10, bows, {n: len(counts[n]) + (2 if n == 1 else 0) for n in range(1, order + 1)}
+
+
+def write_arpa(path: str | Path, probs10: dict, bows: dict, order: int) -> None:
+    by_n: dict[int, list[tuple[str, ...]]] = defaultdict(list)
+    for ng in probs10:
+        by_n[len(ng)].append(ng)
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    opener = gzip.open if path.suffix == ".gz" else open
+    with opener(path, "wt", encoding="utf-8") as f:
+        f.write("\\data\\\n")
+        for n in range(1, order + 1):
+            f.write(f"ngram {n}={len(by_n[n])}\n")
+        for n in range(1, order + 1):
+            f.write(f"\n\\{n}-grams:\n")
+            for ng in sorted(by_n[n]):
+                line = f"{probs10[ng]:.6f}\t{' '.join(ng)}"
+                if n < order and ng in bows:
+                    line += f"\t{bows[ng]:.6f}"
+                elif n < order and ng == (BOS,):
+                    line += "\t0.0"
+                f.write(line + "\n")
+        f.write("\n\\end\\\n")

canopy/language_model/kenlm_backend.py

@@ -0,0 +1,62 @@
+"""Optional KenLM integration: building with ``lmplz``/``build_binary`` and
+scoring through the ``kenlm`` Python module (both LGPL-2.1, not vendored)."""
+
+from __future__ import annotations
+
+import shutil
+import subprocess
+from pathlib import Path
+
+LN10 = 2.302585092994046
+
+
+def build_with_lmplz(text_file: str | Path, out_arpa: str | Path, order: int = 4, prune: list[int] | None = None,
+                     memory: str = "40%", lmplz: str | None = None, discount_fallback: bool = True) -> Path:
+    exe = lmplz or shutil.which("lmplz")
+    if not exe:
+        raise FileNotFoundError("KenLM `lmplz` not found on PATH; build KenLM (docs/TRAINING.md) or use --backend python")
+    cmd = [exe, "-o", str(order), "-S", memory, "--text", str(text_file), "--arpa", str(out_arpa)]
+    if prune:
+        cmd += ["--prune", *map(str, prune)]
+    if discount_fallback:
+        cmd.append("--discount_fallback")
+    subprocess.run(cmd, check=True)
+    return Path(out_arpa)
+
+
+def build_binary(arpa: str | Path, out_bin: str | Path, exe: str | None = None) -> Path:
+    exe = exe or shutil.which("build_binary")
+    if not exe:
+        raise FileNotFoundError("KenLM `build_binary` not found on PATH")
+    subprocess.run([exe, str(arpa), str(out_bin)], check=True)
+    return Path(out_bin)
+
+
+class KenLMScorer:
+    """Same interface as ArpaLM (natural-log scores) on top of ``kenlm.Model``."""
+
+    def __init__(self, path: str | Path):
+        import kenlm
+
+        self.model = kenlm.Model(str(path))
+        self.order = self.model.order
+        self._kenlm = kenlm
+
+    def _state_for(self, history: tuple[str, ...]):
+        st = self._kenlm.State()
+        self.model.BeginSentenceWrite(st)
+        for w in history[-(self.order - 1):]:
+            out = self._kenlm.State()
+            self.model.BaseScore(st, w, out)
+            st = out
+        return st
+
+    def score(self, history: tuple[str, ...], word: str) -> float:
+        st = self._state_for(history)
+        out = self._kenlm.State()
+        return self.model.BaseScore(st, word, out) * LN10
+
+    def score_end(self, history: tuple[str, ...]) -> float:
+        st = self._state_for(history)
+        out = self._kenlm.State()
+        return self.model.BaseScore(st, "</s>", out) * LN10

canopy/language_model/text.py

@@ -0,0 +1,38 @@
+"""LM training text preparation with held-out filtering.
+
+LM text is in ``text_train`` space (what the acoustic model emits). Any
+sentence whose ``sentence_key`` occurs in a dev or test manifest is removed,
+because CV-Pakistan sentences are re-recorded by several speakers and would
+otherwise leak test transcripts into the LM.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+
+from canopy.data.manifest import sentence_key
+from canopy.text.normalize import get_normalizer
+
+
+def filter_heldout(sentences: Iterable[str], language: str, heldout_keys: set[str]) -> tuple[list[str], int]:
+    norm = get_normalizer(language)
+    kept, dropped = [], 0
+    for s in sentences:
+        if sentence_key(language, norm.eval_norm(s)) in heldout_keys:
+            dropped += 1
+            continue
+        t = norm.train_norm(s)
+        if t:
+            kept.append(t)
+    return kept, dropped
+
+
+def lm_sentences(train_rows: Iterable[dict], language: str, heldout_rows: Iterable[dict], extra_text: Iterable[str] = (),
+                 dedupe: bool = True) -> tuple[list[str], dict]:
+    heldout = {r["sentence_key"] for r in heldout_rows if r["language"] == language}
+    base = [r["text_raw"] for r in train_rows if r["language"] == language]
+    kept, dropped = filter_heldout(list(base) + list(extra_text), language, heldout)
+    if dedupe:
+        kept = list(dict.fromkeys(kept))
+    return kept, {"language": language, "sentences": len(kept), "dropped_heldout_matches": dropped,
+                  "heldout_sentence_keys": len(heldout)}

canopy/languages.py

@@ -0,0 +1,92 @@
+"""Access to the master language registry (configs/languages.yaml)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from functools import lru_cache
+from pathlib import Path
+
+from canopy.utils.io import load_yaml
+
+DEFAULT_REGISTRY = Path(__file__).resolve().parents[1] / "configs" / "languages.yaml"
+
+STATUSES = (
+    "production",
+    "supported",
+    "experimental",
+    "data-limited",
+    "evaluation-pending",
+    "not-currently-supported",
+    "auxiliary",
+)
+# Statuses whose data may be used for training. `not-currently-supported`
+# means there is no licence-clean transcribed data at all.
+TRAINABLE_STATUSES = frozenset(STATUSES) - {"not-currently-supported"}
+
+
+@dataclass(frozen=True)
+class Language:
+    code: str
+    name: str
+    status: str
+    primary_script: str
+    alternate_scripts: tuple[str, ...] = ()
+    family: str | None = None
+    omnilingual_seen: bool | None = None
+    raw: dict = field(default_factory=dict, compare=False, repr=False)
+
+    @property
+    def trainable(self) -> bool:
+        return self.status in TRAINABLE_STATUSES
+
+
+class LanguageRegistry:
+    def __init__(self, path: str | Path = DEFAULT_REGISTRY):
+        self.path = Path(path)
+        data = load_yaml(self.path)
+        self.schema_version = data.get("schema_version")
+        self._langs: dict[str, Language] = {}
+        for entry in data["languages"]:
+            scripts = entry.get("scripts") or {}
+            status = entry.get("status")
+            if status not in STATUSES:
+                raise ValueError(f"{self.path}: language {entry.get('code')!r} has unknown status {status!r}")
+            lang = Language(
+                code=entry["code"],
+                name=entry["name"],
+                status=status,
+                primary_script=scripts.get("primary", "Arab"),
+                alternate_scripts=tuple(scripts.get("alternates") or ()),
+                family=entry.get("family"),
+                omnilingual_seen=entry.get("omnilingual_seen"),
+                raw=entry,
+            )
+            if lang.code in self._langs:
+                raise ValueError(f"{self.path}: duplicate language code {lang.code!r}")
+            self._langs[lang.code] = lang
+
+    def __contains__(self, code: str) -> bool:
+        return code in self._langs
+
+    def __getitem__(self, code: str) -> Language:
+        try:
+            return self._langs[code]
+        except KeyError:
+            raise KeyError(f"unknown language code {code!r}; add it to {self.path} first") from None
+
+    def __iter__(self):
+        return iter(self._langs.values())
+
+    def __len__(self) -> int:
+        return len(self._langs)
+
+    def codes(self, trainable_only: bool = False) -> list[str]:
+        return [c for c, lang in self._langs.items() if lang.trainable or not trainable_only]
+
+    def by_status(self, status: str) -> list[Language]:
+        return [lang for lang in self._langs.values() if lang.status == status]
+
+
+@lru_cache(maxsize=4)
+def get_registry(path: str | None = None) -> LanguageRegistry:
+    return LanguageRegistry(path or DEFAULT_REGISTRY)

canopy/models/__init__.py

@@ -0,0 +1,4 @@
+from canopy.models.canopy_ctc import CanopyCTC
+from canopy.models.config import PRESETS, ModelConfig
+
+__all__ = ["PRESETS", "CanopyCTC", "ModelConfig"]

canopy/models/canopy_ctc.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.utils.checkpoint import checkpoint
+
+from canopy.frontend.features import LogMelFrontend
+from canopy.models.config import ModelConfig
+from canopy.models.modules import ConvSubsampling4x, MacaronBlock
+
+
+def lengths_to_mask(lengths: torch.Tensor, T: int) -> torch.Tensor:
+    return torch.arange(T, device=lengths.device).unsqueeze(0) < lengths.unsqueeze(1)
+
+
+class CanopyCTC(nn.Module):
+    """Conformer-style CTC encoder with intermediate CTC, optional language
+    conditioning, self-conditioning and a language-ID head."""
+
+    def __init__(self, cfg: ModelConfig):
+        super().__init__()
+        if cfg.vocab_size <= 1:
+            raise ValueError("cfg.vocab_size must be set from the tokenizer")
+        for k in cfg.aux_ctc_layers:
+            if not 1 <= k < cfg.n_layers:
+                raise ValueError(f"aux_ctc_layers entries must be in [1, n_layers), got {k}")
+        self.cfg = cfg
+        d = cfg.d_model
+        self.frontend = LogMelFrontend(n_mels=cfg.n_mels, normalize=cfg.feature_normalize)
+        self.subsampling = ConvSubsampling4x(cfg.n_mels, d, cfg.dropout)
+        self.lang_embed = nn.Embedding(cfg.num_languages + 1, d) if cfg.language_embedding else None
+        if self.lang_embed is not None:
+            nn.init.zeros_(self.lang_embed.weight)
+        n = cfg.n_layers
+        self.blocks = nn.ModuleList([
+            MacaronBlock(d, cfg.n_heads, cfg.ff_multiplier, cfg.conv_kernel, cfg.conv_expansion, cfg.se_ratio,
+                         cfg.dropout, cfg.stochastic_depth * i / max(n - 1, 1), cfg.conv_norm)
+            for i in range(n)
+        ])
+        self.ctc_head = nn.Linear(d, cfg.vocab_size)
+        self.aux_heads = nn.ModuleDict({str(k): nn.Linear(d, cfg.vocab_size) for k in cfg.aux_ctc_layers})
+        self.self_cond = (nn.ModuleDict({str(k): nn.Linear(cfg.vocab_size, d, bias=False) for k in cfg.aux_ctc_layers})
+                          if cfg.self_conditioning else None)
+        self.lid_head = nn.Linear(d, cfg.num_languages + 1) if cfg.lid_head else None
+        with torch.no_grad():
+            for head in [self.ctc_head, *self.aux_heads.values()]:
+                head.bias[0] = cfg.blank_bias_init
+
+    # ------------------------------------------------------------------
+    def encode(self, feats: torch.Tensor, feat_lengths: torch.Tensor, lang_ids: torch.Tensor | None = None) -> dict:
+        x, lengths = self.subsampling(feats, feat_lengths)
+        mask = lengths_to_mask(lengths, x.shape[1])
+        if self.lang_embed is not None:
+            if lang_ids is None:
+                lang_ids = torch.zeros(x.shape[0], dtype=torch.long, device=x.device)
+            x = x + self.lang_embed(lang_ids).unsqueeze(1)
+        aux_logits: dict[int, torch.Tensor] = {}
+        for i, block in enumerate(self.blocks, start=1):
+            if self.cfg.gradient_checkpointing and self.training:
+                x = checkpoint(block, x, mask, use_reentrant=False)
+            else:
+                x = block(x, mask)
+            key = str(i)
+            if key in self.aux_heads:
+                logits = self.aux_heads[key](x)
+                aux_logits[i] = logits
+                if self.self_cond is not None:
+                    x = x + self.self_cond[key](logits.float().softmax(-1).to(x.dtype))
+        out = {"hidden": x, "lengths": lengths, "mask": mask, "logits": self.ctc_head(x), "aux_logits": aux_logits}
+        if self.lid_head is not None:
+            m = mask.unsqueeze(-1).to(x.dtype)
+            pooled = (x * m).sum(1) / m.sum(1).clamp(min=1.0)
+            out["lid_logits"] = self.lid_head(pooled)
+        return out
+
+    def forward(self, wav: torch.Tensor, wav_lengths: torch.Tensor, lang_ids: torch.Tensor | None = None,
+                spec_augment: nn.Module | None = None) -> dict:
+        feats, feat_lengths = self.frontend(wav, wav_lengths)
+        if spec_augment is not None and self.training:
+            feats = spec_augment(feats, feat_lengths)
+        out = self.encode(feats, feat_lengths, lang_ids)
+        out["log_probs"] = F.log_softmax(out["logits"].float(), dim=-1)
+        out["aux_log_probs"] = {k: F.log_softmax(v.float(), dim=-1) for k, v in out["aux_logits"].items()}
+        return out
+
+    # ------------------------------------------------------------------
+    def num_parameters(self, include_aux: bool = True) -> int:
+        n = sum(p.numel() for p in self.parameters())
+        if not include_aux:
+            n -= sum(p.numel() for p in self.aux_heads.parameters())
+        return n
+
+    def parameter_breakdown(self) -> dict[str, int]:
+        def count(m):
+            return sum(p.numel() for p in m.parameters()) if m is not None else 0
+
+        return {
+            "subsampling": count(self.subsampling),
+            "blocks": count(self.blocks),
+            "ctc_head": count(self.ctc_head),
+            "aux_heads": count(self.aux_heads),
+            "self_conditioning": count(self.self_cond),
+            "language_embedding": count(self.lang_embed),
+            "lid_head": count(self.lid_head),
+            "total": count(self),
+        }

canopy/models/config.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field, fields
+
+
+@dataclass
+class ModelConfig:
+    # Frontend
+    n_mels: int = 80
+    feature_normalize: str = "utterance"  # utterance | per_bin | none
+    # Encoder
+    d_model: int = 384
+    n_layers: int = 12
+    n_heads: int = 6
+    ff_multiplier: float = 3.0
+    conv_kernel: int = 31
+    conv_expansion: int = 2
+    conv_norm: str = "batch"  # batch | layer
+    se_ratio: int = 4
+    dropout: float = 0.1
+    stochastic_depth: float = 0.1
+    # Output
+    vocab_size: int = 0
+    blank_bias_init: float = 0.0
+    # Intermediate CTC heads after these blocks (1-indexed: 6 = output of the 6th block).
+    aux_ctc_layers: list[int] = field(default_factory=lambda: [6])
+    self_conditioning: bool = False
+    # Language conditioning. `languages` fixes the id order; id 0 means "no language".
+    languages: list[str] = field(default_factory=list)
+    language_embedding: bool = False
+    lid_head: bool = False
+    # Memory
+    gradient_checkpointing: bool = False
+
+    @property
+    def num_languages(self) -> int:
+        return len(self.languages)
+
+    def to_dict(self) -> dict:
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, d: dict) -> ModelConfig:
+        known = {f.name for f in fields(cls)}
+        unknown = set(d) - known
+        if unknown:
+            raise ValueError(f"unknown model config keys: {sorted(unknown)}")
+        return cls(**d)
+
+
+PRESETS: dict[str, dict] = {
+    "canopy-s": {"d_model": 320, "n_layers": 12, "n_heads": 5},
+    "canopy-m": {"d_model": 384, "n_layers": 12, "n_heads": 6},
+    "canopy-m16": {"d_model": 384, "n_layers": 16, "n_heads": 6},
+    "canopy-l": {"d_model": 512, "n_layers": 16, "n_heads": 8, "aux_ctc_layers": [8]},
+    # For tests and smoke runs only.
+    "canopy-tiny": {"d_model": 64, "n_layers": 2, "n_heads": 2, "aux_ctc_layers": [1], "conv_kernel": 7},
+}

canopy/models/modules.py

@@ -0,0 +1,183 @@
+"""Encoder building blocks: RMSNorm, rotary MHSA, SwiGLU, SE convolution, Macaron block."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        xf = x.float()
+        xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
+        return xf.type_as(x) * self.weight
+
+
+class Rotary(nn.Module):
+    def __init__(self, head_dim: int, base: float = 10000.0):
+        super().__init__()
+        inv = 1.0 / (base ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
+        self.register_buffer("inv_freq", inv, persistent=False)
+        self._len = 0
+        self._cos: torch.Tensor | None = None
+        self._sin: torch.Tensor | None = None
+
+    def forward(self, T: int, device, dtype) -> tuple[torch.Tensor, torch.Tensor]:
+        if self._cos is None or T > self._len or self._cos.device != device:
+            n = max(T, 2 * self._len, 1024)
+            t = torch.arange(n, device=device, dtype=torch.float32)
+            freqs = torch.outer(t, self.inv_freq.to(device))
+            emb = torch.cat([freqs, freqs], dim=-1)
+            self._cos, self._sin, self._len = emb.cos(), emb.sin(), n
+        return self._cos[:T].to(dtype), self._sin[:T].to(dtype)
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    a, b = x.chunk(2, dim=-1)
+    return torch.cat([-b, a], dim=-1)
+
+
+class RotaryMHSA(nn.Module):
+    def __init__(self, d_model: int, n_heads: int, dropout: float):
+        super().__init__()
+        if d_model % n_heads:
+            raise ValueError("d_model must be divisible by n_heads")
+        self.h = n_heads
+        self.dh = d_model // n_heads
+        self.q = nn.Linear(d_model, d_model, bias=False)
+        self.k = nn.Linear(d_model, d_model, bias=False)
+        self.v = nn.Linear(d_model, d_model, bias=False)
+        self.o = nn.Linear(d_model, d_model, bias=False)
+        self.rotary = Rotary(self.dh)
+        self.p = dropout
+
+    def forward(self, x: torch.Tensor, key_mask: torch.Tensor | None) -> torch.Tensor:
+        B, T, D = x.shape
+        q = self.q(x).view(B, T, self.h, self.dh).transpose(1, 2)
+        k = self.k(x).view(B, T, self.h, self.dh).transpose(1, 2)
+        v = self.v(x).view(B, T, self.h, self.dh).transpose(1, 2)
+        cos, sin = self.rotary(T, x.device, q.dtype)
+        q = q * cos + _rotate_half(q) * sin
+        k = k * cos + _rotate_half(k) * sin
+        attn_mask = key_mask[:, None, None, :] if key_mask is not None else None
+        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=self.p if self.training else 0.0)
+        return self.o(out.transpose(1, 2).reshape(B, T, D))
+
+
+class SwiGLUFFN(nn.Module):
+    def __init__(self, d_model: int, multiplier: float, dropout: float):
+        super().__init__()
+        inner = int(d_model * multiplier)
+        self.norm = RMSNorm(d_model)
+        self.gate = nn.Linear(d_model, inner, bias=False)
+        self.up = nn.Linear(d_model, inner, bias=False)
+        self.down = nn.Linear(inner, d_model, bias=False)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        h = self.down(self.drop(F.silu(self.gate(h)) * self.up(h)))
+        return self.drop(h)
+
+
+class SqueezeExcite(nn.Module):
+    def __init__(self, channels: int, ratio: int):
+        super().__init__()
+        mid = max(1, channels // ratio)
+        self.fc1 = nn.Linear(channels, mid)
+        self.fc2 = nn.Linear(mid, channels)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+        # x: (B, T, C). Mean over valid frames only.
+        if mask is None:
+            s = x.mean(dim=1, keepdim=True)
+        else:
+            m = mask.unsqueeze(-1).to(x.dtype)
+            s = (x * m).sum(dim=1, keepdim=True) / m.sum(dim=1, keepdim=True).clamp(min=1.0)
+        return x * torch.sigmoid(self.fc2(F.silu(self.fc1(s))))
+
+
+class ConvModule(nn.Module):
+    def __init__(self, d_model: int, kernel: int, expansion: int, se_ratio: int, dropout: float, norm: str = "batch"):
+        super().__init__()
+        inner = d_model * expansion
+        self.norm = RMSNorm(d_model)
+        self.pw1 = nn.Linear(d_model, 2 * inner)
+        self.dw = nn.Conv1d(inner, inner, kernel, padding=kernel // 2, groups=inner)
+        if norm == "batch":
+            self.conv_norm: nn.Module = nn.BatchNorm1d(inner)
+        elif norm == "layer":
+            self.conv_norm = nn.GroupNorm(1, inner)
+        else:
+            raise ValueError(f"conv_norm must be batch|layer, got {norm!r}")
+        self.se = SqueezeExcite(inner, se_ratio)
+        self.pw2 = nn.Linear(inner, d_model)
+        self.drop = nn.Dropout(dropout)
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+        h = F.glu(self.pw1(self.norm(x)), dim=-1)
+        if mask is not None:
+            h = h.masked_fill(~mask.unsqueeze(-1), 0.0)
+        h = self.dw(h.transpose(1, 2))
+        h = F.silu(self.conv_norm(h)).transpose(1, 2)
+        h = self.se(h, mask)
+        return self.drop(self.pw2(h))
+
+
+class MacaronBlock(nn.Module):
+    """½FFN → MHSA → Conv → ½FFN → RMSNorm, each with a residual connection."""
+
+    def __init__(self, d_model: int, n_heads: int, ff_multiplier: float, conv_kernel: int, conv_expansion: int,
+                 se_ratio: int, dropout: float, drop_rate: float, conv_norm: str = "batch"):
+        super().__init__()
+        self.ff1 = SwiGLUFFN(d_model, ff_multiplier, dropout)
+        self.attn_norm = RMSNorm(d_model)
+        self.attn = RotaryMHSA(d_model, n_heads, dropout)
+        self.attn_drop = nn.Dropout(dropout)
+        self.conv = ConvModule(d_model, conv_kernel, conv_expansion, se_ratio, dropout, conv_norm)
+        self.ff2 = SwiGLUFFN(d_model, ff_multiplier, dropout)
+        self.final_norm = RMSNorm(d_model)
+        self.drop_rate = drop_rate
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor | None) -> torch.Tensor:
+        if self.training and self.drop_rate > 0 and torch.rand(()) < self.drop_rate:
+            return x
+        x = x + 0.5 * self.ff1(x)
+        x = x + self.attn_drop(self.attn(self.attn_norm(x), mask))
+        x = x + self.conv(x, mask)
+        x = x + 0.5 * self.ff2(x)
+        return self.final_norm(x)
+
+
+class ConvSubsampling4x(nn.Module):
+    """Two 3×3 stride-2 Conv2d layers (4× in time and frequency) → Linear → RMSNorm."""
+
+    def __init__(self, n_mels: int, d_model: int, dropout: float):
+        super().__init__()
+        self.conv1 = nn.Conv2d(1, d_model, 3, stride=2, padding=1)
+        self.conv2 = nn.Conv2d(d_model, d_model, 3, stride=2, padding=1)
+        f = (n_mels + 1) // 2
+        f = (f + 1) // 2
+        self.linear = nn.Linear(d_model * f, d_model)
+        self.norm = RMSNorm(d_model)
+        self.drop = nn.Dropout(dropout)
+
+    @staticmethod
+    def output_lengths(lengths: torch.Tensor) -> torch.Tensor:
+        lengths = torch.div(lengths - 1, 2, rounding_mode="floor") + 1
+        lengths = torch.div(lengths - 1, 2, rounding_mode="floor") + 1
+        return lengths.clamp(min=1)
+
+    def forward(self, x: torch.Tensor, lengths: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        x = F.silu(self.conv1(x.unsqueeze(1)))
+        x = F.silu(self.conv2(x))
+        B, C, T, Fq = x.shape
+        x = x.permute(0, 2, 1, 3).reshape(B, T, C * Fq)
+        x = self.drop(self.norm(self.linear(x)))
+        return x, self.output_lengths(lengths).clamp(max=T)

canopy/text/__init__.py

@@ -0,0 +1,3 @@
+from canopy.text.normalize import Normalizer, get_normalizer
+
+__all__ = ["Normalizer", "get_normalizer"]

canopy/text/normalize.py

@@ -0,0 +1,110 @@
+"""Text normalization with three separate views per utterance.
+
+``raw``        the transcript exactly as distributed (never modified)
+``train_norm`` the tokenizer/model target
+``eval_norm``  the scoring view used for WER/CER
+
+Both derived views are deterministic functions of ``raw`` and the language code.
+"""
+
+from __future__ import annotations
+
+import re
+import unicodedata
+from collections import Counter
+from collections.abc import Iterable
+from functools import lru_cache
+
+from canopy.text.rules import HARAKAT, LanguageRules, rules_for
+
+ZWNJ = "\u200c"
+TATWEEL = "\u0640"
+# Invisible format characters that never carry orthographic meaning here.
+_INVISIBLE = frozenset(
+    ["\u200b", "\u200d", "\u200e", "\u200f", "\u2060", "\ufeff", "\u00ad"]
+    + [chr(c) for c in range(0x202A, 0x202F)]
+    + [chr(c) for c in range(0x2066, 0x206A)]
+)
+_WS = re.compile(r"\s+")
+
+_EASTERN_DIGITS = {chr(0x0660 + i): str(i) for i in range(10)}
+_EASTERN_DIGITS.update({chr(0x06F0 + i): str(i) for i in range(10)})
+
+
+def _is_presentation_form(ch: str) -> bool:
+    o = ord(ch)
+    return 0xFB50 <= o <= 0xFDFF or 0xFE70 <= o <= 0xFEFE
+
+
+class Normalizer:
+    def __init__(self, rules: LanguageRules):
+        self.rules = rules
+        self.code = rules.code
+
+    def _base(self, text: str, keep_zwnj: bool) -> str:
+        text = unicodedata.normalize("NFC", text)
+        out = []
+        for ch in text:
+            if _is_presentation_form(ch):
+                out.append(unicodedata.normalize("NFKC", ch))
+            elif ch in _INVISIBLE or ch == TATWEEL:
+                continue
+            elif ch == ZWNJ:
+                if keep_zwnj:
+                    out.append(ch)
+            else:
+                out.append(ch)
+        text = "".join(out)
+        if self.rules.char_map:
+            text = "".join(self.rules.char_map.get(ch, ch) for ch in text)
+        buf = []
+        for ch in text:
+            cat = unicodedata.category(ch)
+            if ch in self.rules.keep_chars or ch == ZWNJ:
+                buf.append(ch)
+            elif cat[0] in ("P", "S") or cat in ("Cc", "Cf"):
+                buf.append(" ")
+            else:
+                buf.append(ch)
+        text = "".join(buf)
+        if self.rules.lowercase:
+            text = text.lower()
+        text = unicodedata.normalize("NFC", text)
+        return _WS.sub(" ", text).strip()
+
+    def train_norm(self, raw: str) -> str:
+        return self._base(raw, keep_zwnj=self.rules.keep_zwnj_train)
+
+    def eval_norm(self, raw: str) -> str:
+        text = self._base(raw, keep_zwnj=False)
+        text = "".join(_EASTERN_DIGITS.get(ch, ch) for ch in text)
+        if self.rules.eval_char_map:
+            text = "".join(self.rules.eval_char_map.get(ch, ch) for ch in text)
+        if self.rules.strip_harakat_eval:
+            text = "".join(ch for ch in text if ch not in HARAKAT)
+        return _WS.sub(" ", text).strip()
+
+    def eval_norm_hypothesis(self, hyp: str) -> str:
+        """Hypotheses are already in train_norm space; apply the same scoring view."""
+        return self.eval_norm(hyp)
+
+
+@lru_cache(maxsize=128)
+def get_normalizer(code: str, script: str | None = None) -> Normalizer:
+    return Normalizer(rules_for(code, script))
+
+
+def char_inventory(texts: Iterable[str]) -> Counter:
+    c: Counter = Counter()
+    for t in texts:
+        c.update(t)
+    return c
+
+
+def describe_char(ch: str) -> dict:
+    return {
+        "char": ch,
+        "codepoint": f"U+{ord(ch):04X}",
+        "name": unicodedata.name(ch, "<unnamed>"),
+        "category": unicodedata.category(ch),
+    }

canopy/text/rules.py

@@ -0,0 +1,88 @@
+"""Per-language normalization rules.
+
+Policy (docs/LANGUAGE_MATRIX.md, section 4):
+
+* ``raw_reference`` is never modified.
+* A letter-level mapping applies ONLY to languages listed here with that
+  mapping. Nothing is inherited from Urdu (or any other language) by default.
+* Every language not listed explicitly gets the script-generic rules only,
+  which are limited to changes that carry no orthographic meaning
+  (Unicode NFC, presentation forms, invisible format characters, tatweel,
+  punctuation, whitespace).
+* ``reviewed`` stays False until a native-speaker reviewer signs off the
+  rule set; the evaluation report prints this flag next to each language.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+ARABIC_KAF = "\u0643"  # ك
+KEHEH = "\u06a9"  # ک
+ARABIC_YEH = "\u064a"  # ي
+ALEF_MAKSURA = "\u0649"  # ى
+FARSI_YEH = "\u06cc"  # ی
+ARABIC_HEH = "\u0647"  # ه
+HEH_GOAL = "\u06c1"  # ہ
+
+# Optional vowel signs (harakat), superscript alef and small high marks.
+HARAKAT = frozenset(chr(c) for c in list(range(0x064B, 0x0653)) + [0x0670])
+
+
+@dataclass(frozen=True)
+class LanguageRules:
+    code: str
+    script: str = "Arab"
+    # Applied in both train_norm and eval_norm.
+    char_map: dict[str, str] = field(default_factory=dict)
+    # Applied additionally in eval_norm only.
+    eval_char_map: dict[str, str] = field(default_factory=dict)
+    lowercase: bool = False
+    keep_zwnj_train: bool = False
+    strip_harakat_eval: bool = False
+    # Characters that must survive punctuation removal (e.g. apostrophes used as letters).
+    keep_chars: frozenset[str] = frozenset()
+    reviewed: bool = False
+    notes: str = ""
+
+
+_LATIN_APOSTROPHES = frozenset({"'", "\u2019", "\u02bc"})
+
+RULES: dict[str, LanguageRules] = {
+    "urd": LanguageRules(
+        code="urd",
+        char_map={ARABIC_KAF: KEHEH, ARABIC_YEH: FARSI_YEH, ALEF_MAKSURA: FARSI_YEH, ARABIC_HEH: HEH_GOAL},
+        keep_zwnj_train=True,
+        strip_harakat_eval=True,
+        notes="Yeh barree U+06D2 and do-chashmi heh U+06BE are kept distinct. Harakat stripped in eval_norm only.",
+    ),
+    "pus": LanguageRules(
+        code="pus",
+        char_map={ARABIC_KAF: KEHEH},
+        notes="Pashto yeh forms (ی ي ې ۍ ئ) are distinct letters and heh U+0647 is standard; neither is mapped.",
+    ),
+    "snd": LanguageRules(
+        code="snd",
+        notes="Sindhi uses swash kaf U+06AA for /k/ and Arabic yeh/heh; the Urdu mappings must not be applied.",
+    ),
+    "kas": LanguageRules(
+        code="kas",
+        notes="Kashmiri vowel marks are orthographically required; harakat are never stripped.",
+    ),
+    "eng": LanguageRules(code="eng", script="Latn", lowercase=True, keep_chars=_LATIN_APOSTROPHES),
+}
+
+for _code in ("bsk", "kls", "wbl"):
+    RULES[_code] = LanguageRules(
+        code=_code,
+        script="Latn",
+        lowercase=True,
+        keep_chars=_LATIN_APOSTROPHES,
+        notes="Latin-script community orthography. Lowercasing and apostrophe handling await reviewer confirmation.",
+    )
+
+
+def rules_for(code: str, script: str | None = None) -> LanguageRules:
+    if code in RULES:
+        return RULES[code]
+    return LanguageRules(code=code, script=script or "Arab", notes="script-generic rules only")

canopy/tokenizer/__init__.py

@@ -0,0 +1,5 @@
+from canopy.tokenizer.base import BLANK, UNK, CTCTokenizer, load_tokenizer
+from canopy.tokenizer.char import CharTokenizer
+from canopy.tokenizer.spm import SentencePieceTokenizer, balanced_sample
+
+__all__ = ["BLANK", "UNK", "CTCTokenizer", "CharTokenizer", "SentencePieceTokenizer", "balanced_sample", "load_tokenizer"]

canopy/tokenizer/base.py

@@ -0,0 +1,96 @@
+"""Common CTC tokenizer interface. Id 0 is always the CTC blank, id 1 is <unk>."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+from canopy.utils.io import sha256_file
+
+BLANK = "<blank>"
+UNK = "<unk>"
+BLANK_ID = 0
+UNK_ID = 1
+
+
+class CTCTokenizer:
+    type: str = "base"
+
+    def __init__(self, tokens: list[str], language_token_ids: dict[str, list[int]] | None = None, meta: dict | None = None):
+        if tokens[BLANK_ID] != BLANK or tokens[UNK_ID] != UNK:
+            raise ValueError("token 0 must be <blank> and token 1 must be <unk>")
+        self.tokens = tokens
+        self.token_to_id = {t: i for i, t in enumerate(tokens)}
+        self.language_token_ids = {k: sorted(set(v)) for k, v in (language_token_ids or {}).items()}
+        self.meta = meta or {}
+        self.version: str | None = None
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.tokens)
+
+    @property
+    def blank_id(self) -> int:
+        return BLANK_ID
+
+    @property
+    def unk_id(self) -> int:
+        return UNK_ID
+
+    def encode(self, text: str, language: str | None = None) -> list[int]:
+        raise NotImplementedError
+
+    def decode(self, ids: list[int]) -> str:
+        raise NotImplementedError
+
+    def always_allowed_ids(self) -> list[int]:
+        return [BLANK_ID]
+
+    def language_mask(self, language: str) -> list[bool] | None:
+        """Allowed output ids for ``language`` (O2 output masking), or None if unknown."""
+        ids = self.language_token_ids.get(language)
+        if ids is None:
+            return None
+        mask = [False] * self.vocab_size
+        for i in ids:
+            mask[i] = True
+        for i in self.always_allowed_ids():
+            mask[i] = True
+        return mask
+
+    def labels_for_pyctcdecode(self) -> list[str]:
+        raise NotImplementedError
+
+    # ------------------------------------------------------------------ io
+    def _payload(self) -> dict:
+        return {"type": self.type, "tokens": self.tokens, "language_token_ids": self.language_token_ids, "meta": self.meta}
+
+    def save(self, directory: str | os.PathLike) -> Path:
+        d = Path(directory)
+        d.mkdir(parents=True, exist_ok=True)
+        path = d / "tokenizer.json"
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(self._payload(), f, ensure_ascii=False, indent=1)
+        self.version = sha256_file(path)[:16]
+        return path
+
+
+def load_tokenizer(directory: str | os.PathLike) -> CTCTokenizer:
+    d = Path(directory)
+    path = d / "tokenizer.json" if d.is_dir() else d
+    with open(path, encoding="utf-8") as f:
+        payload = json.load(f)
+    t = payload["type"]
+    if t == "char":
+        from canopy.tokenizer.char import CharTokenizer
+
+        tok = CharTokenizer.from_payload(payload)
+    elif t == "sentencepiece":
+        from canopy.tokenizer.spm import SentencePieceTokenizer
+
+        tok = SentencePieceTokenizer.from_payload(payload, path.parent)
+    else:
+        raise ValueError(f"unknown tokenizer type {t!r}")
+    tok.version = sha256_file(path)[:16]
+    return tok

canopy/tokenizer/char.py

@@ -0,0 +1,76 @@
+"""Character tokenizers.
+
+``unit="codepoint"``: one token per Unicode code point (NFC).
+``unit="grapheme"``: a base character plus its following combining marks form
+one token, so vowel signs and other diacritics stay attached to their letter.
+"""
+
+from __future__ import annotations
+
+import unicodedata
+from collections import Counter
+from collections.abc import Iterable
+
+from canopy.tokenizer.base import BLANK, UNK, UNK_ID, CTCTokenizer
+
+SPACE = " "
+
+
+def split_units(text: str, unit: str) -> list[str]:
+    if unit == "codepoint":
+        return list(text)
+    if unit != "grapheme":
+        raise ValueError(f"unit must be codepoint|grapheme, got {unit!r}")
+    out: list[str] = []
+    for ch in text:
+        if out and out[-1] != SPACE and unicodedata.category(ch) in ("Mn", "Me", "Mc"):
+            out[-1] += ch
+        else:
+            out.append(ch)
+    return out
+
+
+class CharTokenizer(CTCTokenizer):
+    type = "char"
+
+    def __init__(self, tokens, language_token_ids=None, meta=None):
+        super().__init__(tokens, language_token_ids, meta)
+        self.unit = self.meta.get("unit", "codepoint")
+
+    @classmethod
+    def train(cls, texts_by_language: dict[str, Iterable[str]], unit: str = "codepoint", min_count: int = 1) -> CharTokenizer:
+        counts: Counter = Counter()
+        per_lang: dict[str, Counter] = {}
+        for lang, texts in texts_by_language.items():
+            c: Counter = Counter()
+            for t in texts:
+                c.update(split_units(t, unit))
+            per_lang[lang] = c
+            counts.update(c)
+        units = sorted(u for u, n in counts.items() if n >= min_count and u != SPACE)
+        tokens = [BLANK, UNK, SPACE, *units]
+        idx = {t: i for i, t in enumerate(tokens)}
+        lang_ids = {lang: [idx[u] for u in c if u in idx] for lang, c in per_lang.items()}
+        meta = {"unit": unit, "min_count": min_count, "unit_counts": {u: counts[u] for u in units}}
+        return cls(tokens, lang_ids, meta)
+
+    @classmethod
+    def from_payload(cls, payload: dict) -> CharTokenizer:
+        return cls(payload["tokens"], payload.get("language_token_ids"), payload.get("meta"))
+
+    def always_allowed_ids(self) -> list[int]:
+        return [0, self.token_to_id[SPACE]]
+
+    def encode(self, text: str, language: str | None = None) -> list[int]:
+        return [self.token_to_id.get(u, UNK_ID) for u in split_units(text, self.unit)]
+
+    def decode(self, ids: list[int]) -> str:
+        out = []
+        for i in ids:
+            if i <= 0:
+                continue
+            out.append("\u2047" if i == UNK_ID else self.tokens[i])
+        return " ".join("".join(out).split())
+
+    def labels_for_pyctcdecode(self) -> list[str]:
+        return ["", "\u2047", *self.tokens[2:]]

canopy/tokenizer/diagnostics.py

@@ -0,0 +1,95 @@
+"""Static tokenizer diagnostics (no acoustic training): RESEARCH_PLAN.md section 7.2."""
+
+from __future__ import annotations
+
+from collections import Counter
+
+import numpy as np
+
+from canopy.data.dataset import ctc_required_frames
+from canopy.data.validation import char_script
+from canopy.tokenizer.base import UNK_ID
+
+
+def _pct(a, q):
+    return float(np.percentile(a, q)) if len(a) else 0.0
+
+
+def tokenization_report(tokenizer, rows: list[dict], *, fps: float = 25.0, d_model: int = 384,
+                        tokenizer_id: str = "unnamed") -> dict:
+    """Per-language diagnostics. ``rows`` are training-split manifest rows."""
+    by_lang: dict[str, list[dict]] = {}
+    for r in rows:
+        by_lang.setdefault(r["language"], []).append(r)
+
+    encoded = {lang: [tokenizer.encode(r["text_train"], language=lang) for r in lrows] for lang, lrows in by_lang.items()}
+    global_counts: Counter = Counter()
+    lang_types: dict[str, set] = {}
+    for lang, encs in encoded.items():
+        s = set()
+        for e in encs:
+            global_counts.update(e)
+            s.update(e)
+        lang_types[lang] = s
+
+    reports = {}
+    for lang, lrows in by_lang.items():
+        encs = encoded[lang]
+        n_tok = np.array([len(e) for e in encs], dtype=float)
+        n_chr = np.array([len(r["text_train"].replace(" ", "")) for r in lrows], dtype=float)
+        dur = np.array([r["duration"] for r in lrows], dtype=float)
+        frames = np.floor(dur * fps)
+        req = np.array([ctc_required_frames(e) for e in encs], dtype=float)
+        ratio = req / np.maximum(frames, 1)
+        repeats = req - n_tok
+        types = lang_types[lang]
+        others = set().union(*(v for k, v in lang_types.items() if k != lang)) if len(lang_types) > 1 else set()
+        n_types = max(1, len(types))
+        total_tokens = max(1.0, n_tok.sum())
+        chars = Counter(c for r in lrows for c in r["text_train"] if not c.isspace())
+        total_chars = max(1, sum(chars.values()))
+        script_cov: Counter = Counter()
+        for c, k in chars.items():
+            script_cov[char_script(c)] += k
+        cpt = n_chr / np.maximum(n_tok, 1)
+        tps = n_tok / np.maximum(dur, 1e-6)
+        reports[lang] = {
+            "tokenizer_id": tokenizer_id,
+            "language": lang,
+            "n_utterances": len(lrows),
+            "frames_per_second": fps,
+            "mean_chars_per_token": float(cpt.mean()) if len(cpt) else 0.0,
+            "p50_chars_per_token": _pct(cpt, 50),
+            "p95_chars_per_token": _pct(cpt, 95),
+            "tokens_per_second": {"mean": float(tps.mean()) if len(tps) else 0.0, "p95": _pct(tps, 95), "p99": _pct(tps, 99)},
+            "target_token_frame_ratio": {"mean": float(ratio.mean()) if len(ratio) else 0.0, "p95": _pct(ratio, 95),
+                                         "p99": _pct(ratio, 99), "max": float(ratio.max()) if len(ratio) else 0.0},
+            "required_blanks_for_repeats_per_utt": float(repeats.mean()) if len(repeats) else 0.0,
+            "frame_budget_violation_rate": float((ratio > 1.0).mean()) if len(ratio) else 0.0,
+            "estimated_blank_rate": float(1.0 - (n_tok.sum() / max(1.0, frames.sum()))),
+            "rare_token_percentage": {
+                f"lt{k}": 100.0 * sum(1 for t in types if global_counts[t] < k) / n_types for k in (10, 100, 1000)
+            },
+            "tokens_used_by_language": len(types),
+            "tokens_exclusive_to_language": len(types - others),
+            "oov_or_unk_rate": float(sum(e.count(UNK_ID) for e in encs) / total_tokens),
+            "byte_fallback_rate": 0.0,
+            "script_coverage": {k: v / total_chars for k, v in sorted(script_cov.items())},
+            "digit_handling": "split" if tokenizer.type == "sentencepiece" else "per-character",
+            "punctuation_handling": "removed_in_train_norm",
+            "ctc_head_params": (d_model + 1) * tokenizer.vocab_size,
+        }
+    return reports
+
+
+def prescreen(reports: dict, max_violation_rate: float = 0.005, max_rare_share: float = 50.0) -> dict:
+    """Apply the section 7.2 pre-screen gate. Returns {passed, failures}."""
+    failures = []
+    for lang, r in reports.items():
+        if r["frame_budget_violation_rate"] > max_violation_rate:
+            failures.append({"language": lang, "reason": "frame_budget_violation_rate",
+                             "value": r["frame_budget_violation_rate"]})
+        if r["rare_token_percentage"]["lt100"] > max_rare_share:
+            failures.append({"language": lang, "reason": "rare_token_share_lt100",
+                             "value": r["rare_token_percentage"]["lt100"]})
+    return {"passed": not failures, "failures": failures}

canopy/tokenizer/spm.py

@@ -0,0 +1,111 @@
+"""SentencePiece (unigram / BPE) tokenizer with CTC blank at id 0.
+
+SentencePiece's own ids are shifted by +1, so SentencePiece <unk> (0) becomes
+Canopy id 1. Normalization is disabled inside SentencePiece
+(``normalization_rule_name=identity``): Canopy's per-language normalizer has
+already produced ``text_train``, and SentencePiece's default NFKC would fold
+characters that are distinct in some Pakistani orthographies.
+"""
+
+from __future__ import annotations
+
+import io
+import math
+import random
+from collections.abc import Iterable
+from pathlib import Path
+
+from canopy.tokenizer.base import BLANK, UNK, CTCTokenizer
+
+SPM_SPACE = "\u2581"
+
+
+def balanced_sample(texts_by_language: dict[str, list[str]], temperature: float, seed: int = 0,
+                    total: int | None = None) -> tuple[list[str], dict[str, int]]:
+    """Resample sentences so language ``l`` gets a share ∝ n_l ** (1/temperature)."""
+    rng = random.Random(seed)
+    n = {lang: len(t) for lang, t in texts_by_language.items() if t}
+    total = total or sum(n.values())
+    if math.isinf(temperature):
+        w = {lang: 1.0 for lang in n}
+    else:
+        w = {lang: c ** (1.0 / temperature) for lang, c in n.items()}
+    z = sum(w.values())
+    out, counts = [], {}
+    for lang, texts in texts_by_language.items():
+        if not texts:
+            continue
+        k = max(1, round(total * w[lang] / z))
+        if k <= len(texts):
+            picked = rng.sample(texts, k)
+        else:
+            picked = texts * (k // len(texts)) + rng.sample(texts, k % len(texts))
+        out.extend(picked)
+        counts[lang] = k
+    rng.shuffle(out)
+    return out, counts
+
+
+class SentencePieceTokenizer(CTCTokenizer):
+    type = "sentencepiece"
+
+    def __init__(self, tokens, language_token_ids=None, meta=None, model_proto: bytes | None = None):
+        super().__init__(tokens, language_token_ids, meta)
+        import sentencepiece as spm
+
+        self._proto = model_proto
+        self.sp = spm.SentencePieceProcessor(model_proto=model_proto)
+
+    @classmethod
+    def train(cls, texts_by_language: dict[str, list[str]], vocab_size: int, model_type: str = "unigram",
+              temperature: float = 5.0, seed: int = 0, character_coverage: float = 1.0,
+              max_sentences: int | None = 2_000_000) -> SentencePieceTokenizer:
+        import sentencepiece as spm
+
+        sample, counts = balanced_sample(texts_by_language, temperature, seed,
+                                         total=min(max_sentences or 10**12, sum(len(v) for v in texts_by_language.values())))
+        buf = io.BytesIO()
+        spm.SentencePieceTrainer.train(
+            sentence_iterator=iter(sample), model_writer=buf, vocab_size=vocab_size, model_type=model_type,
+            character_coverage=character_coverage, byte_fallback=False, split_digits=True,
+            normalization_rule_name="identity", add_dummy_prefix=True, remove_extra_whitespaces=True,
+            unk_id=0, bos_id=-1, eos_id=-1, pad_id=-1, hard_vocab_limit=False,
+            input_sentence_size=0, shuffle_input_sentence=False, num_threads=4, minloglevel=2,
+        )
+        proto = buf.getvalue()
+        sp = spm.SentencePieceProcessor(model_proto=proto)
+        tokens = [BLANK, UNK] + [sp.id_to_piece(i) for i in range(1, sp.get_piece_size())]
+        meta = {"model_type": model_type, "requested_vocab_size": vocab_size, "temperature": temperature,
+                "seed": seed, "character_coverage": character_coverage, "training_sentences_per_language": counts}
+        tok = cls(tokens, None, meta, proto)
+        tok.language_token_ids = {
+            lang: sorted({i for t in texts for i in tok.encode(t)}) for lang, texts in texts_by_language.items()
+        }
+        return tok
+
+    @classmethod
+    def from_payload(cls, payload: dict, directory: Path) -> SentencePieceTokenizer:
+        proto = (directory / "spm.model").read_bytes()
+        return cls(payload["tokens"], payload.get("language_token_ids"), payload.get("meta"), proto)
+
+    def save(self, directory):
+        d = Path(directory)
+        d.mkdir(parents=True, exist_ok=True)
+        (d / "spm.model").write_bytes(self._proto)
+        return super().save(d)
+
+    def encode(self, text: str, language: str | None = None) -> list[int]:
+        return [i + 1 for i in self.sp.encode(text, out_type=int)]
+
+    def decode(self, ids: Iterable[int]) -> str:
+        return self.sp.decode([i - 1 for i in ids if i > 0])
+
+    def always_allowed_ids(self) -> list[int]:
+        ids = [0]
+        sp_space = self.token_to_id.get(SPM_SPACE)
+        if sp_space is not None:
+            ids.append(sp_space)
+        return ids
+
+    def labels_for_pyctcdecode(self) -> list[str]:
+        return ["", "\u2047", *[t.replace(SPM_SPACE, " ") if t.startswith(SPM_SPACE) else t for t in self.tokens[2:]]]

canopy/training/__init__.py

@@ -0,0 +1,3 @@
+from canopy.training.schedule import WarmupCosine
+
+__all__ = ["WarmupCosine"]

canopy/training/checkpoint.py

@@ -0,0 +1,77 @@
+"""Checkpoint format.
+
+Training checkpoints hold model, optimizer, scheduler, scaler and loop state.
+Inference checkpoints (``export_inference_checkpoint``) hold only the model
+config, weights, language list and the tokenizer payload, so one file is
+enough to transcribe.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+from canopy.models.canopy_ctc import CanopyCTC
+from canopy.models.config import ModelConfig
+
+FORMAT_VERSION = 1
+
+
+def save_training_checkpoint(path: Path, *, model, optimizer, scheduler, scaler, state: dict, model_config: dict,
+                             tokenizer_dir: str) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_suffix(".tmp")
+    torch.save({
+        "format_version": FORMAT_VERSION,
+        "kind": "training",
+        "model_config": model_config,
+        "model_state": model.state_dict(),
+        "optimizer_state": optimizer.state_dict(),
+        "scheduler_state": scheduler.state_dict(),
+        "scaler_state": scaler.state_dict() if scaler is not None else None,
+        "loop_state": state,
+        "tokenizer_dir": tokenizer_dir,
+        "rng": {"torch": torch.get_rng_state(),
+                "cuda": torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None},
+    }, tmp)
+    tmp.replace(path)
+
+
+def load_checkpoint(path: str | Path, map_location="cpu") -> dict:
+    return torch.load(str(path), map_location=map_location, weights_only=False)
+
+
+def export_inference_checkpoint(train_ckpt: str | Path, out: str | Path, tokenizer_dir: str | Path | None = None,
+                                extra_meta: dict | None = None) -> Path:
+    ck = load_checkpoint(train_ckpt)
+    tok_dir = Path(tokenizer_dir or ck["tokenizer_dir"])
+    payload = json.loads((tok_dir / "tokenizer.json").read_text(encoding="utf-8"))
+    spm_model = (tok_dir / "spm.model").read_bytes() if (tok_dir / "spm.model").exists() else None
+    state = {k: v for k, v in ck["model_state"].items() if not k.startswith("aux_heads.")}
+    cfg = dict(ck["model_config"])
+    if not cfg.get("self_conditioning"):
+        cfg["aux_ctc_layers"] = []
+    out = Path(out)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    torch.save({
+        "format_version": FORMAT_VERSION,
+        "kind": "inference",
+        "model_config": cfg,
+        "model_state": state,
+        "tokenizer_payload": payload,
+        "spm_model": spm_model,
+        "meta": {"source_checkpoint": str(train_ckpt), "loop_state": ck.get("loop_state"), **(extra_meta or {})},
+    }, out)
+    return out
+
+
+def model_from_checkpoint(ck: dict) -> CanopyCTC:
+    cfg = ModelConfig.from_dict(ck["model_config"])
+    model = CanopyCTC(cfg)
+    missing, unexpected = model.load_state_dict(ck["model_state"], strict=False)
+    unexpected = [k for k in unexpected if not k.startswith("aux_heads.")]
+    if missing or unexpected:
+        raise RuntimeError(f"checkpoint mismatch: missing={missing} unexpected={unexpected}")
+    return model

canopy/training/losses.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+def ctc_loss(log_probs: torch.Tensor, lengths: torch.Tensor, targets: torch.Tensor, target_lengths: torch.Tensor,
+             blank: int = 0) -> torch.Tensor:
+    """Mean over utterances of per-utterance CTC loss divided by target length."""
+    return F.ctc_loss(log_probs.transpose(0, 1).float(), targets, lengths, target_lengths, blank=blank,
+                      reduction="mean", zero_infinity=True)
+
+
+def canopy_loss(out: dict, batch: dict, *, aux_weight: float = 0.3, lid_weight: float = 0.0,
+                lid_targets: torch.Tensor | None = None) -> tuple[torch.Tensor, dict]:
+    """InterCTC-style objective: (1 - w) * CTC_final + w * mean(CTC_aux) [+ lid_weight * CE_lid]."""
+    main = ctc_loss(out["log_probs"], out["lengths"], batch["targets"], batch["target_lengths"])
+    parts = {"ctc": main.detach()}
+    loss = main
+    aux = out.get("aux_log_probs") or {}
+    if aux and aux_weight > 0:
+        aux_losses = [ctc_loss(lp, out["lengths"], batch["targets"], batch["target_lengths"]) for lp in aux.values()]
+        aux_mean = torch.stack(aux_losses).mean()
+        loss = (1 - aux_weight) * main + aux_weight * aux_mean
+        parts["aux_ctc"] = aux_mean.detach()
+    if lid_weight > 0 and "lid_logits" in out and lid_targets is not None:
+        lid = F.cross_entropy(out["lid_logits"].float(), lid_targets)
+        loss = loss + lid_weight * lid
+        parts["lid"] = lid.detach()
+    return loss, parts

canopy/training/run.py

@@ -0,0 +1,48 @@
+"""Run directories and reproducibility artifacts.
+
+runs/canopy_exp_<id>_<slug>_<YYYYmmdd-HHMMSS>/
+    config.yaml  git_commit.txt  environment.txt  environment.json
+    dataset_manifest.json  dataset_hashes.json  tokenizer_version  tokenizer/
+    model_config.json  training_log.jsonl  eval_history.jsonl  sampler_audit.jsonl
+    training_metrics.json  evaluation_metrics.json  checkpoints/
+Existing run directories are never overwritten.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from pathlib import Path
+
+from canopy.utils.env import environment_info, environment_text, git_commit
+from canopy.utils.io import save_json, save_yaml, sha256_file
+
+
+def create_run_dir(output_root: str | Path, exp_id: str, slug: str) -> Path:
+    stamp = time.strftime("%Y%m%d-%H%M%S")
+    d = Path(output_root) / f"canopy_exp_{exp_id}_{slug}_{stamp}"
+    if d.exists():
+        raise FileExistsError(f"{d} already exists; runs are never overwritten")
+    (d / "checkpoints").mkdir(parents=True)
+    return d
+
+
+def write_static_artifacts(run_dir: Path, *, config: dict, manifests: list[str], dataset_summary: dict,
+                           tokenizer, model_config: dict, parameter_breakdown: dict) -> None:
+    save_yaml(config, run_dir / "config.yaml")
+    g = git_commit(Path(__file__).resolve().parents[2])
+    (run_dir / "git_commit.txt").write_text(
+        f"{g['commit']}\nbranch: {g['branch']}\ndirty: {g['dirty']}\n", encoding="utf-8")
+    info = environment_info()
+    (run_dir / "environment.txt").write_text(environment_text(info), encoding="utf-8")
+    save_json(info, run_dir / "environment.json")
+    save_json({"manifests": manifests, "summary": dataset_summary}, run_dir / "dataset_manifest.json")
+    save_json({m: sha256_file(m) for m in dict.fromkeys(manifests)}, run_dir / "dataset_hashes.json")
+    tok_path = tokenizer.save(run_dir / "tokenizer")
+    (run_dir / "tokenizer_version").write_text(f"{tokenizer.type} {tokenizer.version} {tok_path.name}\n", encoding="utf-8")
+    save_json({"model": model_config, "parameters": parameter_breakdown}, run_dir / "model_config.json")
+
+
+def append_jsonl(path: Path, obj: dict) -> None:
+    with open(path, "a", encoding="utf-8") as f:
+        f.write(json.dumps(obj, ensure_ascii=False) + "\n")

canopy/training/schedule.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import math
+
+
+class WarmupCosine:
+    """Linear warmup from ``init_ratio * lr`` to ``lr``, then cosine decay to ``final_ratio * lr``."""
+
+    def __init__(self, optimizer, *, total_steps: int, warmup_steps: int, init_ratio: float = 0.02,
+                 final_ratio: float = 0.02):
+        self.opt = optimizer
+        self.total = max(1, total_steps)
+        self.warmup = max(0, warmup_steps)
+        self.init_ratio = init_ratio
+        self.final_ratio = final_ratio
+        self.base = [g["lr"] for g in optimizer.param_groups]
+        self.step_num = 0
+        self._apply()
+
+    def factor(self, step: int) -> float:
+        if self.warmup and step < self.warmup:
+            return self.init_ratio + (1 - self.init_ratio) * step / self.warmup
+        progress = min(1.0, (step - self.warmup) / max(1, self.total - self.warmup))
+        return self.final_ratio + (1 - self.final_ratio) * 0.5 * (1 + math.cos(math.pi * progress))
+
+    def _apply(self):
+        f = self.factor(self.step_num)
+        for g, b in zip(self.opt.param_groups, self.base):
+            g["lr"] = b * f
+
+    def step(self):
+        self.step_num += 1
+        self._apply()
+
+    def state_dict(self) -> dict:
+        return {"step_num": self.step_num, "base": self.base}
+
+    def load_state_dict(self, d: dict):
+        self.step_num = d["step_num"]
+        self.base = d["base"]
+        self._apply()

canopy/training/trainer.py

@@ -0,0 +1,385 @@
+"""Single-GPU trainer for CanopyCTC.
+
+Model selection uses the dev split only. The test split is never loaded here.
+"""
+
+from __future__ import annotations
+
+import random
+import shutil
+import time
+from pathlib import Path
+
+import torch
+from torch.utils.data import DataLoader
+
+from canopy.data.dataset import ASRDataset, collate
+from canopy.data.manifest import load_manifest, summarize
+from canopy.data.sampler import BalancedBatchSampler
+from canopy.evaluation.runner import evaluate_rows
+from canopy.frontend.augment import SpecAugment
+from canopy.models.canopy_ctc import CanopyCTC
+from canopy.models.config import PRESETS, ModelConfig
+from canopy.tokenizer.base import load_tokenizer
+from canopy.training.checkpoint import load_checkpoint, save_training_checkpoint
+from canopy.training.losses import canopy_loss
+from canopy.training.run import append_jsonl, create_run_dir, write_static_artifacts
+from canopy.training.schedule import WarmupCosine
+from canopy.utils.io import save_json
+from canopy.utils.seed import seed_everything
+
+
+def resolve_device(name: str) -> torch.device:
+    if name == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(name)
+
+
+def resolve_precision(name: str, device: torch.device) -> str:
+    if name != "auto":
+        if name == "bf16" and device.type == "cuda" and not torch.cuda.is_bf16_supported():
+            raise RuntimeError("bf16 requested but this GPU does not support it; use precision=fp16")
+        return name
+    if device.type != "cuda":
+        return "fp32"
+    return "bf16" if torch.cuda.is_bf16_supported() else "fp16"
+
+
+AMP_DTYPES = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": None}
+
+
+def build_model_config(cfg: dict, vocab_size: int, languages: list[str]) -> ModelConfig:
+    mcfg = cfg["model"]
+    base = dict(PRESETS[mcfg.get("preset", "canopy-m")])
+    base.update(mcfg.get("overrides") or {})
+    hw = cfg.get("hardware", {})
+    if "gradient_checkpointing" in hw:
+        base["gradient_checkpointing"] = bool(hw["gradient_checkpointing"])
+    base["vocab_size"] = vocab_size
+    base["languages"] = languages
+    return ModelConfig.from_dict(base)
+
+
+def param_groups(model: torch.nn.Module, weight_decay: float) -> list[dict]:
+    decay, no_decay = [], []
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        (no_decay if p.ndim < 2 or "lang_embed" in name else decay).append(p)
+    return [{"params": decay, "weight_decay": weight_decay}, {"params": no_decay, "weight_decay": 0.0}]
+
+
+def _subsample_dev(rows: list[dict], per_lang: int | None, seed: int) -> list[dict]:
+    if not per_lang:
+        return rows
+    by: dict[str, list[dict]] = {}
+    for r in rows:
+        by.setdefault(r["language"], []).append(r)
+    rng = random.Random(seed)
+    out = []
+    for _lang, rs in sorted(by.items()):
+        rs = sorted(rs, key=lambda r: r["id"])
+        out.extend(rs if len(rs) <= per_lang else rng.sample(rs, per_lang))
+    return out
+
+
+class Trainer:
+    @staticmethod
+    def _log_resume(run_dir: Path, cfg: dict, state: dict) -> None:
+        """Append to resume_log.jsonl; warn when the effective batch differs from the original run."""
+        import yaml
+
+        from canopy.utils.env import environment_info
+
+        original = {}
+        if (run_dir / "config.yaml").exists():
+            original = yaml.safe_load((run_dir / "config.yaml").read_text(encoding="utf-8")) or {}
+
+        def batch_key(c: dict) -> tuple:
+            h, t = c.get("hardware", {}), c.get("train", {})
+            return h.get("max_batch_seconds"), h.get("max_batch_size"), t.get("grad_accum")
+
+        changed = bool(original) and batch_key(original) != batch_key(cfg)
+        if changed:
+            print(f"[canopy] WARNING: batch settings changed on resume {batch_key(original)} -> {batch_key(cfg)} "
+                  "(max_batch_seconds, max_batch_size, grad_accum); the remainder of the run is not comparable "
+                  "to an uninterrupted run and the sampler position within the epoch is approximate")
+        append_jsonl(run_dir / "resume_log.jsonl", {
+            "time": time.strftime("%Y-%m-%dT%H:%M:%S"), "step": state["step"], "epoch": state["epoch"],
+            "sources": cfg.get("_sources"), "hardware": cfg.get("hardware"), "grad_accum": cfg.get("train", {}).get("grad_accum"),
+            "batch_settings_changed": changed, "environment": environment_info()})
+
+    def __init__(self, cfg: dict, resume: str | None = None):
+        self.cfg = cfg
+        self.resume = resume
+
+    # ------------------------------------------------------------------ setup
+    def _load_data(self):
+        d = self.cfg["data"]
+        prov = ["human"] + (["pseudo"] if d.get("include_pseudo") else []) + (["synthetic"] if d.get("include_synthetic") else [])
+        train_manifests = list(d["train_manifests"])
+        dev_manifests = list(d.get("dev_manifests") or train_manifests)
+        train = load_manifest(train_manifests, languages=d.get("languages"), splits=[d.get("train_split", "train")],
+                              provenances=prov, min_duration=d.get("min_duration", 0.3),
+                              max_duration=d.get("max_duration", 20.0))
+        dev_prov = ["human"] + (["synthetic"] if d.get("include_synthetic") else [])
+        dev = load_manifest(dev_manifests, languages=d.get("languages"), splits=[d.get("dev_split", "dev")],
+                            provenances=dev_prov, max_duration=d.get("max_duration_eval", 60.0))
+        if not train:
+            raise RuntimeError("no training rows after filtering; check data.train_manifests / languages / splits")
+        if not dev:
+            raise RuntimeError("no dev rows; model selection needs a dev split (never the test split)")
+        dev = _subsample_dev(dev, d.get("max_dev_utts_per_language"), self.cfg.get("seed", 0))
+        return train, dev, train_manifests + [m for m in dev_manifests if m not in train_manifests]
+
+    def run(self) -> Path:
+        cfg = self.cfg
+        seed = int(cfg.get("seed", 1234))
+        seed_everything(seed, deterministic=bool(cfg.get("deterministic", False)))
+        hw = cfg.get("hardware", {})
+        device = resolve_device(hw.get("device", "auto"))
+        precision = resolve_precision(hw.get("precision", "auto"), device)
+        amp_dtype = AMP_DTYPES[precision]
+        if device.type == "cuda":
+            torch.backends.cuda.matmul.allow_tf32 = True
+            torch.backends.cudnn.allow_tf32 = True
+
+        train_rows, dev_rows, manifests = self._load_data()
+        tokenizer = load_tokenizer(cfg["tokenizer"]["path"])
+        languages = sorted({r["language"] for r in train_rows})
+        lang2id = {lang: i + 1 for i, lang in enumerate(languages)}
+        mcfg = build_model_config(cfg, tokenizer.vocab_size, languages)
+        model = CanopyCTC(mcfg).to(device)
+
+        aug = cfg.get("augment", {})
+        train_ds = ASRDataset(train_rows, tokenizer, lang2id, encoder_fps=25.0,
+                              speed_factors=tuple(aug.get("speed_factors", [1.0])), train=True)
+        samp = cfg.get("sampling", {})
+        sampler = BalancedBatchSampler(
+            [r["language"] for r in train_ds.rows], [r["duration"] for r in train_ds.rows],
+            max_batch_seconds=hw.get("max_batch_seconds", 200), max_batch_size=hw.get("max_batch_size", 64),
+            temperature=float(samp.get("temperature", 5.0)), caps_hours=samp.get("caps_hours"),
+            default_cap_hours=samp.get("default_cap_hours"), epoch_hours=samp.get("epoch_hours"), seed=seed)
+        nw = int(hw.get("num_workers", 4))
+        loader = DataLoader(train_ds, batch_sampler=sampler, collate_fn=collate, num_workers=nw,
+                            pin_memory=device.type == "cuda", persistent_workers=nw > 0,
+                            prefetch_factor=hw.get("prefetch_factor", 4) if nw > 0 else None)
+        sa = aug.get("specaugment", {}) or {}
+        specaug = SpecAugment(**sa) if aug.get("specaugment_enabled", True) else None
+
+        oc = cfg["optim"]
+        tc = cfg["train"]
+        optimizer = torch.optim.AdamW(param_groups(model, oc.get("weight_decay", 0.01)), lr=float(oc["lr"]),
+                                      betas=tuple(oc.get("betas", [0.9, 0.98])), eps=float(oc.get("eps", 1e-8)),
+                                      fused=device.type == "cuda")
+        max_steps = int(tc["max_steps"])
+        scheduler = WarmupCosine(optimizer, total_steps=max_steps, warmup_steps=int(oc.get("warmup_steps", 5000)),
+                                 init_ratio=float(oc.get("warmup_init_ratio", 0.02)),
+                                 final_ratio=float(oc.get("final_lr_ratio", 0.02)))
+        scaler = torch.amp.GradScaler("cuda", enabled=precision == "fp16")
+
+        state = {"step": 0, "micro": 0, "epoch": 0, "batch_in_epoch": 0, "audio_seconds": 0.0, "best_macro_cer": None,
+                 "best_step": None, "evals_since_best": 0, "oom_skips": 0, "nan_skips": 0, "elapsed_s": 0.0}
+        if self.resume:
+            ck = load_checkpoint(self.resume, map_location=device)
+            run_dir = Path(self.resume).resolve().parent.parent
+            model.load_state_dict(ck["model_state"])
+            optimizer.load_state_dict(ck["optimizer_state"])
+            scheduler.load_state_dict(ck["scheduler_state"])
+            if ck.get("scaler_state"):
+                scaler.load_state_dict(ck["scaler_state"])
+            state.update(ck["loop_state"])
+            print(f"[canopy] resumed from {self.resume} at step {state['step']}")
+            self._log_resume(run_dir, cfg, state)
+        else:
+            ex = cfg.get("experiment", {})
+            run_dir = create_run_dir(ex.get("output_root", "runs"), str(ex.get("id", "000")), ex.get("slug", "run"))
+            write_static_artifacts(
+                run_dir, config=cfg, manifests=manifests,
+                dataset_summary={"train": summarize(train_ds.rows), "dev": summarize(dev_rows),
+                                 "train_filter": train_ds.stats.as_dict(), "languages": languages},
+                tokenizer=tokenizer, model_config=mcfg.to_dict(), parameter_breakdown=model.parameter_breakdown())
+        tok_dir = str(run_dir / "tokenizer")
+        if hw.get("compile"):
+            model = torch.compile(model)  # type: ignore[assignment]
+        core = getattr(model, "_orig_mod", model)
+
+        print(f"[canopy] run dir: {run_dir}")
+        print(f"[canopy] device={device} precision={precision} params={core.num_parameters():,} "
+              f"languages={len(languages)} train_utts={len(train_ds)} dev_utts={len(dev_rows)}")
+        print(f"[canopy] filtered: {train_ds.stats.as_dict()}")
+
+        accum = int(tc.get("grad_accum", 1))
+        lc = cfg.get("loss", {})
+        aux_w = float(lc.get("aux_ctc_weight", 0.3))
+        lid_w = float(lc.get("lid_weight", 0.0))
+        cond_drop = float(lc.get("lang_condition_dropout", 0.5))
+        clip = float(oc.get("grad_clip", 1.0))
+        log_every, eval_every = int(tc.get("log_every", 50)), int(tc.get("eval_every", 2000))
+        save_every = int(tc.get("save_every", eval_every))
+        max_hours = tc.get("max_audio_hours")
+        patience = tc.get("early_stop_patience")
+        max_nan = int(tc.get("max_consecutive_nan", 50))
+        max_oom = int(tc.get("max_oom_skips", 200))
+        ev = cfg.get("eval", {})
+
+        def save(name: str):
+            state["elapsed_s"] = elapsed_before + (time.time() - t_start)
+            save_training_checkpoint(run_dir / "checkpoints" / name, model=core, optimizer=optimizer, scheduler=scheduler,
+                                     scaler=scaler, state=dict(state), model_config=mcfg.to_dict(), tokenizer_dir=tok_dir)
+
+        def evaluate(n_boot: int = 0) -> dict:
+            metrics, _ = evaluate_rows(core, tokenizer, dev_rows, lang2id=lang2id, device=device,
+                                       max_batch_seconds=hw.get("eval_max_batch_seconds", 2 * hw.get("max_batch_seconds", 200)),
+                                       use_language=ev.get("use_language", True), use_masks=ev.get("use_masks", True),
+                                       amp_dtype=amp_dtype, n_boot=n_boot)
+            core.train()
+            return metrics
+
+        def do_eval() -> bool:
+            m = evaluate()
+            rec = {"step": state["step"], "audio_hours": state["audio_seconds"] / 3600, "macro": m["macro"],
+                   "weighted": m["weighted"], "languages": {k: {"cer": v["cer"], "wer": v["wer"]} for k, v in m["languages"].items()},
+                   "mean_blank_rate": m["decode"]["mean_blank_rate"]}
+            append_jsonl(run_dir / "eval_history.jsonl", rec)
+            cer = m["macro"]["cer"]
+            print(f"[canopy] eval step {state['step']}: dev macro CER {cer:.4f} WER {m['macro']['wer']:.4f} "
+                  f"(weighted CER {m['weighted']['cer']:.4f})")
+            if state["best_macro_cer"] is None or cer < state["best_macro_cer"]:
+                state.update(best_macro_cer=cer, best_step=state["step"], evals_since_best=0)
+                save("best.pt")
+            else:
+                state["evals_since_best"] += 1
+            return bool(patience) and state["evals_since_best"] >= int(patience)
+
+        def rotate():
+            keep = int(tc.get("keep_last", 3))
+            ckpts = sorted((run_dir / "checkpoints").glob("step_*.pt"), key=lambda p: int(p.stem.split("_")[1]))
+            for p in ckpts[:-keep] if keep > 0 else []:
+                p.unlink(missing_ok=True)
+
+        model.train()
+        optimizer.zero_grad(set_to_none=True)
+        t_start = time.time()
+        elapsed_before = float(state.get("elapsed_s", 0.0))
+        window = {"loss": 0.0, "ctc": 0.0, "aux": 0.0, "n": 0, "audio": 0.0, "t": time.time()}
+        consecutive_nan = 0
+        stop = False
+        if device.type == "cuda":
+            torch.cuda.reset_peak_memory_stats(device)
+
+        if len(sampler) == 0:
+            raise RuntimeError("sampler produced no batches")
+        while not stop and state["step"] < max_steps:
+            sampler.set_epoch(state["epoch"])
+            if state["batch_in_epoch"] == 0:
+                append_jsonl(run_dir / "sampler_audit.jsonl", sampler.audit())
+            sampler.start_batch = state["batch_in_epoch"]
+            for batch in loader:
+                state["batch_in_epoch"] += 1
+                wav = batch["wav"].to(device, non_blocking=True)
+                wl = batch["wav_lengths"].to(device, non_blocking=True)
+                b = {"targets": batch["targets"].to(device, non_blocking=True),
+                     "target_lengths": batch["target_lengths"].to(device, non_blocking=True)}
+                true_lids = batch["lang_ids"].to(device, non_blocking=True)
+                lids = true_lids
+                if core.lang_embed is not None and cond_drop > 0:
+                    lids = torch.where(torch.rand(lids.shape, device=device) < cond_drop, torch.zeros_like(lids), lids)
+                try:
+                    with torch.autocast(device.type, dtype=amp_dtype, enabled=amp_dtype is not None):
+                        out = model(wav, wl, lids, spec_augment=specaug)
+                    loss, parts = canopy_loss(out, b, aux_weight=aux_w, lid_weight=lid_w, lid_targets=true_lids)
+                    if not torch.isfinite(loss):
+                        consecutive_nan += 1
+                        state["nan_skips"] += 1
+                        optimizer.zero_grad(set_to_none=True)
+                        if consecutive_nan >= max_nan:
+                            raise RuntimeError(f"{consecutive_nan} consecutive non-finite losses; aborting")
+                        continue
+                    consecutive_nan = 0
+                    scaler.scale(loss / accum).backward()
+                except torch.cuda.OutOfMemoryError:
+                    state["oom_skips"] += 1
+                    del wav, wl, b
+                    optimizer.zero_grad(set_to_none=True)
+                    torch.cuda.empty_cache()
+                    print(f"[canopy] OOM on a batch (skipped {state['oom_skips']}); lower hardware.max_batch_seconds if frequent")
+                    if state["oom_skips"] > max_oom:
+                        raise
+                    continue
+
+                batch_audio = float(batch["wav_lengths"].sum()) / 16000
+                state["audio_seconds"] += batch_audio
+                window["audio"] += batch_audio
+                window["loss"] += float(loss.detach())
+                window["ctc"] += float(parts["ctc"])
+                window["aux"] += float(parts.get("aux_ctc", 0.0))
+                window["n"] += 1
+                state["micro"] += 1
+                if state["micro"] % accum:
+                    continue
+
+                scaler.unscale_(optimizer)
+                gnorm = torch.nn.utils.clip_grad_norm_(core.parameters(), clip)
+                scaler.step(optimizer)
+                scaler.update()
+                optimizer.zero_grad(set_to_none=True)
+                scheduler.step()
+                state["step"] += 1
+                step = state["step"]
+
+                if step % log_every == 0:
+                    dt = time.time() - window["t"]
+                    n = max(1, window["n"])
+                    rec = {"step": step, "epoch": state["epoch"], "loss": window["loss"] / n, "ctc": window["ctc"] / n,
+                           "aux_ctc": window["aux"] / n, "lr": optimizer.param_groups[0]["lr"], "grad_norm": float(gnorm),
+                           "audio_hours_seen": state["audio_seconds"] / 3600, "x_realtime": window["audio"] / max(dt, 1e-9),
+                           "oom_skips": state["oom_skips"], "nan_skips": state["nan_skips"]}
+                    if device.type == "cuda":
+                        rec["gpu_peak_mem_gb"] = torch.cuda.max_memory_allocated(device) / 2**30
+                    if device.type == "cuda":
+                        rec["max_mem_gb"] = torch.cuda.max_memory_allocated(device) / 2**30
+                    append_jsonl(run_dir / "training_log.jsonl", rec)
+                    print(f"[canopy] step {step} loss {rec['loss']:.3f} ctc {rec['ctc']:.3f} lr {rec['lr']:.2e} "
+                          f"gnorm {rec['grad_norm']:.2f} {rec['x_realtime']:.0f}xRT seen {rec['audio_hours_seen']:.1f}h"
+                          + (f" mem {rec['max_mem_gb']:.1f}GB" if "max_mem_gb" in rec else ""))
+                    window = {"loss": 0.0, "ctc": 0.0, "aux": 0.0, "n": 0, "audio": 0.0, "t": time.time()}
+
+                if step % eval_every == 0:
+                    stop = do_eval()
+                if step % save_every == 0:
+                    save(f"step_{step}.pt")
+                    save("last.pt")
+                    rotate()
+                if stop or step >= max_steps or (max_hours and state["audio_seconds"] / 3600 >= float(max_hours)):
+                    stop = True
+                    break
+            else:
+                state["epoch"] += 1
+                state["batch_in_epoch"] = 0
+
+        if state["best_step"] != state["step"]:
+            do_eval()
+        save("last.pt")
+        elapsed = elapsed_before + (time.time() - t_start)
+
+        best = load_checkpoint(run_dir / "checkpoints" / "best.pt", map_location=device)
+        core.load_state_dict(best["model_state"])
+        final = evaluate(n_boot=int(ev.get("final_bootstrap", 1000)))
+        final["checkpoint"] = "best.pt"
+        final["step"] = best["loop_state"]["step"]
+        final["split"] = "dev"
+        save_json(final, run_dir / "evaluation_metrics.json")
+        save_json({
+            "steps": state["step"], "micro_batches": state["micro"], "epochs_completed": state["epoch"],
+            "audio_hours_seen": state["audio_seconds"] / 3600, "elapsed_hours": elapsed / 3600,
+            "x_realtime_overall": state["audio_seconds"] / max(elapsed, 1e-9), "precision": precision,
+            "device": str(device), "best_step": state["best_step"], "best_dev_macro_cer": state["best_macro_cer"],
+            "oom_skips": state["oom_skips"], "nan_skips": state["nan_skips"], "train_filter": train_ds.stats.as_dict(),
+            "sampler_language_probabilities": sampler.probs, "train_hours_by_language": sampler.hours,
+            "max_memory_gb": torch.cuda.max_memory_allocated(device) / 2**30 if device.type == "cuda" else None,
+        }, run_dir / "training_metrics.json")
+        if tc.get("copy_best_to"):
+            shutil.copy2(run_dir / "checkpoints" / "best.pt", tc["copy_best_to"])
+        print(f"[canopy] done. best dev macro CER {state['best_macro_cer']:.4f} at step {state['best_step']}; run dir {run_dir}")
+        return run_dir

canopy/utils/__init__.py

@@ -0,0 +1,13 @@
+from canopy.utils.io import load_yaml, read_jsonl, save_json, save_yaml, sha256_file, sha256_text, write_jsonl
+from canopy.utils.seed import seed_everything
+
+__all__ = [
+    "load_yaml",
+    "read_jsonl",
+    "save_json",
+    "save_yaml",
+    "seed_everything",
+    "sha256_file",
+    "sha256_text",
+    "write_jsonl",
+]

canopy/utils/env.py

@@ -0,0 +1,64 @@
+"""Capture the software/hardware environment for reproducibility records."""
+
+from __future__ import annotations
+
+import importlib
+import platform
+import subprocess
+import sys
+from pathlib import Path
+
+import torch
+
+_PACKAGES = ["numpy", "torch", "torchaudio", "sentencepiece", "soundfile", "yaml", "kenlm", "pyctcdecode",
+             "transformers", "onnx", "onnxruntime"]
+
+
+def _version(name: str) -> str | None:
+    try:
+        mod = importlib.import_module(name)
+    except Exception:
+        return None
+    return getattr(mod, "__version__", "installed (no __version__)")
+
+
+def git_commit(repo: str | Path = ".") -> dict:
+    def run(*args: str) -> str | None:
+        try:
+            return subprocess.check_output(["git", *args], cwd=repo, stderr=subprocess.DEVNULL, text=True).strip()
+        except Exception:
+            return None
+
+    commit = run("rev-parse", "HEAD")
+    status = run("status", "--porcelain")
+    return {"commit": commit, "dirty": bool(status) if status is not None else None, "branch": run("rev-parse", "--abbrev-ref", "HEAD")}
+
+
+def environment_info() -> dict:
+    info = {
+        "python": sys.version.split()[0],
+        "platform": platform.platform(),
+        "processor": platform.processor(),
+        "packages": {p: _version(p) for p in _PACKAGES},
+        "torch_cuda": torch.version.cuda,
+        "cudnn": torch.backends.cudnn.version() if torch.backends.cudnn.is_available() else None,
+        "cuda_available": torch.cuda.is_available(),
+        "gpus": [],
+    }
+    if torch.cuda.is_available():
+        for i in range(torch.cuda.device_count()):
+            p = torch.cuda.get_device_properties(i)
+            info["gpus"].append({"name": p.name, "total_memory_gb": round(p.total_memory / 2**30, 2),
+                                 "capability": f"{p.major}.{p.minor}"})
+    return info
+
+
+def environment_text(info: dict | None = None) -> str:
+    info = info or environment_info()
+    lines = [f"python: {info['python']}", f"platform: {info['platform']}", f"torch_cuda: {info['torch_cuda']}",
+             f"cudnn: {info['cudnn']}", f"cuda_available: {info['cuda_available']}"]
+    for g in info["gpus"]:
+        lines.append(f"gpu: {g['name']} ({g['total_memory_gb']} GB, sm_{g['capability']})")
+    for k, v in info["packages"].items():
+        lines.append(f"{k}: {v}")
+    return "\n".join(lines) + "\n"