diff --git a/README.md b/README.md
index 6ac880de60f32c18e8179a0e21873d5d85dc8a23..80995deef6567c8a5a020efe7f7ba47e67202b4c 100644
--- a/README.md
+++ b/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
-
-
-
-
-
", "", ""
+
+
+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 | 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 .
+ 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")
diff --git a/canopy/language_model/kenlm_backend.py b/canopy/language_model/kenlm_backend.py
new file mode 100644
index 0000000000000000000000000000000000000000..dded5c7e1c865d2fab659bc5742331a0bed59447
--- /dev/null
+++ b/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, "", out) * LN10
diff --git a/canopy/language_model/text.py b/canopy/language_model/text.py
new file mode 100644
index 0000000000000000000000000000000000000000..aba2554f7212e0cd71427bbc09ee46991df8ad38
--- /dev/null
+++ b/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)}
diff --git a/canopy/languages.py b/canopy/languages.py
new file mode 100644
index 0000000000000000000000000000000000000000..2eb1c8cf9cabef21036d6a658c1729ae9bb656ab
--- /dev/null
+++ b/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)
diff --git a/canopy/models/__init__.py b/canopy/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..ccff4eae15fd8e1852409467ddef21f2bedffa6a
--- /dev/null
+++ b/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"]
diff --git a/canopy/models/canopy_ctc.py b/canopy/models/canopy_ctc.py
new file mode 100644
index 0000000000000000000000000000000000000000..622fe043a0888306736a53ca8c3d7e949d82333e
--- /dev/null
+++ b/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),
+ }
diff --git a/canopy/models/config.py b/canopy/models/config.py
new file mode 100644
index 0000000000000000000000000000000000000000..bfffa759c0a18e45dcc43605e085213ad5bc10bf
--- /dev/null
+++ b/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},
+}
diff --git a/canopy/models/modules.py b/canopy/models/modules.py
new file mode 100644
index 0000000000000000000000000000000000000000..16c60de478dc4c5d28ad4c832af840dc0df3a445
--- /dev/null
+++ b/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)
diff --git a/canopy/text/__init__.py b/canopy/text/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..bbc3d81f0befa228c7ea2e7ce413ebf61e3b3b48
--- /dev/null
+++ b/canopy/text/__init__.py
@@ -0,0 +1,3 @@
+from canopy.text.normalize import Normalizer, get_normalizer
+
+__all__ = ["Normalizer", "get_normalizer"]
diff --git a/canopy/text/normalize.py b/canopy/text/normalize.py
new file mode 100644
index 0000000000000000000000000000000000000000..3694b348d1ba236580c9eb6610e7e227d38c6837
--- /dev/null
+++ b/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, ""),
+ "category": unicodedata.category(ch),
+ }
diff --git a/canopy/text/rules.py b/canopy/text/rules.py
new file mode 100644
index 0000000000000000000000000000000000000000..1d9e4b6dd59cdb5d465337c2bda303566b684cd9
--- /dev/null
+++ b/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")
diff --git a/canopy/tokenizer/__init__.py b/canopy/tokenizer/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0194708649d9dee5aeb57453b371bc80f7646ca5
--- /dev/null
+++ b/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"]
diff --git a/canopy/tokenizer/base.py b/canopy/tokenizer/base.py
new file mode 100644
index 0000000000000000000000000000000000000000..e84d9f783680bf07be8625615fb21fb6bd1c2d54
--- /dev/null
+++ b/canopy/tokenizer/base.py
@@ -0,0 +1,96 @@
+"""Common CTC tokenizer interface. Id 0 is always the CTC blank, id 1 is ."""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+
+from canopy.utils.io import sha256_file
+
+BLANK = ""
+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 and token 1 must be ")
+ 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
diff --git a/canopy/tokenizer/char.py b/canopy/tokenizer/char.py
new file mode 100644
index 0000000000000000000000000000000000000000..c43ec9ffa5ef0109654be4301e88571179a07525
--- /dev/null
+++ b/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:]]
diff --git a/canopy/tokenizer/diagnostics.py b/canopy/tokenizer/diagnostics.py
new file mode 100644
index 0000000000000000000000000000000000000000..f897725089da263cab153580140a22b1f89dec2f
--- /dev/null
+++ b/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}
diff --git a/canopy/tokenizer/spm.py b/canopy/tokenizer/spm.py
new file mode 100644
index 0000000000000000000000000000000000000000..20b8cce8500fb40c08e01d6bdefe336224d405d4
--- /dev/null
+++ b/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 (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:]]]
diff --git a/canopy/training/__init__.py b/canopy/training/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..fd3a12ba723dc0ee25bbebf8050ec7e8133b5259
--- /dev/null
+++ b/canopy/training/__init__.py
@@ -0,0 +1,3 @@
+from canopy.training.schedule import WarmupCosine
+
+__all__ = ["WarmupCosine"]
diff --git a/canopy/training/checkpoint.py b/canopy/training/checkpoint.py
new file mode 100644
index 0000000000000000000000000000000000000000..f2180df0d799e68001a01c3681505c2e8fefa508
--- /dev/null
+++ b/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
diff --git a/canopy/training/losses.py b/canopy/training/losses.py
new file mode 100644
index 0000000000000000000000000000000000000000..b620864386b8c72ef0db42d47bc11799fa55bccf
--- /dev/null
+++ b/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
diff --git a/canopy/training/run.py b/canopy/training/run.py
new file mode 100644
index 0000000000000000000000000000000000000000..16b2cf027fbf89d98e83c5dcf4ddbaecb1b603d2
--- /dev/null
+++ b/canopy/training/run.py
@@ -0,0 +1,48 @@
+"""Run directories and reproducibility artifacts.
+
+runs/canopy_exp___/
+ 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")
diff --git a/canopy/training/schedule.py b/canopy/training/schedule.py
new file mode 100644
index 0000000000000000000000000000000000000000..5f99d082ede70ff90206847b62c5ac49838d0734
--- /dev/null
+++ b/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()
diff --git a/canopy/training/trainer.py b/canopy/training/trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..52e4edafe8ee972ee6658a8a50dba1123858a41b
--- /dev/null
+++ b/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
diff --git a/canopy/utils/__init__.py b/canopy/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2663874eccef0871a175a4c92674e34332d5f3b4
--- /dev/null
+++ b/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",
+]
diff --git a/canopy/utils/env.py b/canopy/utils/env.py
new file mode 100644
index 0000000000000000000000000000000000000000..cc308222729e3c9cda52ffda5cd522aecdab90fe
--- /dev/null
+++ b/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"
diff --git a/canopy/utils/io.py b/canopy/utils/io.py
new file mode 100644
index 0000000000000000000000000000000000000000..892752c385cbbdb83d291cfcbdd755d1b70ed3d5
--- /dev/null
+++ b/canopy/utils/io.py
@@ -0,0 +1,129 @@
+from __future__ import annotations
+
+import copy
+import hashlib
+import json
+import os
+from collections.abc import Iterable, Iterator
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+def load_yaml(path: str | os.PathLike) -> Any:
+ with open(path, encoding="utf-8") as f:
+ return yaml.safe_load(f)
+
+
+def save_yaml(obj: Any, path: str | os.PathLike) -> None:
+ Path(path).parent.mkdir(parents=True, exist_ok=True)
+ with open(path, "w", encoding="utf-8") as f:
+ yaml.safe_dump(obj, f, allow_unicode=True, sort_keys=False)
+
+
+def save_json(obj: Any, path: str | os.PathLike, indent: int = 2) -> None:
+ Path(path).parent.mkdir(parents=True, exist_ok=True)
+ tmp = f"{path}.tmp"
+ with open(tmp, "w", encoding="utf-8") as f:
+ json.dump(obj, f, ensure_ascii=False, indent=indent)
+ f.write("\n")
+ os.replace(tmp, path)
+
+
+def read_jsonl(path: str | os.PathLike) -> Iterator[dict]:
+ with open(path, encoding="utf-8") as f:
+ for line_no, line in enumerate(f, 1):
+ line = line.strip()
+ if not line:
+ continue
+ try:
+ yield json.loads(line)
+ except json.JSONDecodeError as e:
+ raise ValueError(f"{path}:{line_no}: invalid JSON: {e}") from e
+
+
+def write_jsonl(rows: Iterable[dict], path: str | os.PathLike) -> int:
+ Path(path).parent.mkdir(parents=True, exist_ok=True)
+ n = 0
+ tmp = f"{path}.tmp"
+ with open(tmp, "w", encoding="utf-8") as f:
+ for row in rows:
+ f.write(json.dumps(row, ensure_ascii=False))
+ f.write("\n")
+ n += 1
+ os.replace(tmp, path)
+ return n
+
+
+def sha256_file(path: str | os.PathLike, chunk: int = 1 << 20) -> str:
+ h = hashlib.sha256()
+ with open(path, "rb") as f:
+ while True:
+ b = f.read(chunk)
+ if not b:
+ break
+ h.update(b)
+ return h.hexdigest()
+
+
+def sha256_text(text: str) -> str:
+ return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+
+def deep_merge(base: dict, override: dict) -> dict:
+ """Recursively merge ``override`` into a copy of ``base``. An empty dict in
+ ``override`` replaces the base value instead of merging into it."""
+ out = copy.deepcopy(base)
+ for k, v in override.items():
+ if isinstance(v, dict) and v and isinstance(out.get(k), dict):
+ out[k] = deep_merge(out[k], v)
+ else:
+ out[k] = copy.deepcopy(v)
+ return out
+
+
+def load_config(path: str | os.PathLike) -> dict:
+ """Load a YAML config, resolving an optional ``inherit:`` list of parent files.
+
+ Parents are resolved relative to the child file and merged in order, then
+ the child's own keys are applied on top.
+ """
+ path = Path(path)
+ cfg = load_yaml(path) or {}
+ parents = cfg.pop("inherit", []) or []
+ if isinstance(parents, str):
+ parents = [parents]
+ merged: dict = {}
+ for p in parents:
+ merged = deep_merge(merged, load_config((path.parent / p).resolve()))
+ return deep_merge(merged, cfg)
+
+
+def set_by_dotted_key(cfg: dict, dotted: str, value: Any) -> None:
+ keys = dotted.split(".")
+ node = cfg
+ for k in keys[:-1]:
+ node = node.setdefault(k, {})
+ node[keys[-1]] = value
+
+
+def apply_overrides(cfg: dict, overrides: Iterable[str]) -> dict:
+ """Apply ``a.b.c=value`` overrides; values are parsed as YAML scalars."""
+ for item in overrides:
+ if "=" not in item:
+ raise ValueError(f"override must look like key=value, got {item!r}")
+ k, v = item.split("=", 1)
+ set_by_dotted_key(cfg, k.strip(), _parse_scalar(v))
+ return cfg
+
+
+def _parse_scalar(v: str):
+ value = yaml.safe_load(v)
+ if isinstance(value, str):
+ # YAML 1.1 reads "3e-4" as a string; learning rates are commonly written that way.
+ try:
+ return float(value)
+ except ValueError:
+ return value
+ return value
diff --git a/canopy/utils/seed.py b/canopy/utils/seed.py
new file mode 100644
index 0000000000000000000000000000000000000000..b6609cf00374f8e5ddcba4ad89e13efa586ddb2f
--- /dev/null
+++ b/canopy/utils/seed.py
@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import os
+import random
+
+import numpy as np
+import torch
+
+
+def seed_everything(seed: int, deterministic: bool = False) -> None:
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ os.environ["PYTHONHASHSEED"] = str(seed)
+ if deterministic:
+ torch.backends.cudnn.deterministic = True
+ torch.backends.cudnn.benchmark = False
+ torch.use_deterministic_algorithms(True, warn_only=True)
diff --git a/canopy_m.py b/canopy_m.py
new file mode 100644
index 0000000000000000000000000000000000000000..33885b3b1f1fcacd8afd4d3e1b8133f766e4f853
--- /dev/null
+++ b/canopy_m.py
@@ -0,0 +1,88 @@
+#!/usr/bin/env python
+"""Transcribe Urdu or Sindhi with ProximaAI/Canopy-M.
+
+Install once:
+
+ pip install torch torchaudio soundfile numpy pyyaml huggingface_hub
+ hf auth login
+
+Download this repository (it contains this script and the Canopy code), then:
+
+ 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
+
+From Python, with this directory on the path:
+
+ from canopy_m import CanopyM
+ asr = CanopyM.from_pretrained("ProximaAI/Canopy-M")
+ print(asr.transcribe("clip.wav", language="urd"))
+ print(asr.transcribe("clip.wav", language="sindhi"))
+
+``language`` is ``urd`` or ``snd``. ``ur`` / ``urdu`` and ``sd`` / ``sindhi`` are accepted.
+The published scores use greedy CTC with the language id on.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+_HERE = Path(__file__).resolve().parent
+for _root in (_HERE, _HERE.parent):
+ if (_root / "canopy" / "__init__.py").is_file():
+ sys.path.insert(0, str(_root))
+ break
+
+import torch
+
+from canopy.inference.transcriber import Transcriber
+
+REPO_ID = "ProximaAI/Canopy-M"
+ALIASES = {"ur": "urd", "urdu": "urd", "urd": "urd", "sd": "snd", "sindhi": "snd", "snd": "snd"}
+
+
+def resolve_language(language: str) -> str:
+ key = language.strip().lower()
+ if key not in ALIASES:
+ known = "urd, snd (also ur, urdu, sd, sindhi)"
+ raise ValueError(f"unknown language {language!r}. Use one of: {known}")
+ return ALIASES[key]
+
+
+class CanopyM:
+ """Urdu and Sindhi CTC recogniser backed by the Canopy-M checkpoint."""
+
+ def __init__(self, transcriber: Transcriber):
+ self.transcriber = transcriber
+
+ @classmethod
+ def from_pretrained(cls, repo_id: str = REPO_ID, *, device: str | None = None, token: str | None = None) -> CanopyM:
+ if device is None:
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ local = _HERE / "canopy-m.pt"
+ if local.is_file():
+ asr = Transcriber.from_checkpoint(local, device=device)
+ else:
+ asr = Transcriber.from_pretrained(repo_id, device=device, token=token)
+ return cls(asr)
+
+ def transcribe(self, audio, language: str = "urd", decoder: str = "greedy") -> str:
+ return self.transcriber.transcribe(audio, language=resolve_language(language), decoder=decoder)
+
+
+def main() -> None:
+ ap = argparse.ArgumentParser(description="Transcribe with ProximaAI/Canopy-M")
+ ap.add_argument("audio", nargs="+", help="wav, flac, or mp3 paths")
+ ap.add_argument("--language", required=True, help="urd or snd")
+ ap.add_argument("--device", default=None, help="cpu, cuda, or omit for auto")
+ ap.add_argument("--repo", default=REPO_ID)
+ args = ap.parse_args()
+ asr = CanopyM.from_pretrained(args.repo, device=args.device)
+ for path in args.audio:
+ print(f"{path}\t{asr.transcribe(path, language=args.language)}")
+
+
+if __name__ == "__main__":
+ main()