diff --git a/README.md b/README.md index eeeaa95c64c35d2d6391ebbc4d5347cf0ef0cabf..ad6cb57de6513b75769bca45ee703298acbc2f2a 100644 --- a/README.md +++ b/README.md @@ -58,15 +58,6 @@ print(asr.transcribe("clip.wav", language="urd")) print(asr.transcribe("clip.wav", language="snd")) ``` -The snapshot in this repository also runs on its own: - -```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 --model canopy-m --language urd -``` - `language` accepts `urd`, `ur`, `urdu`, `snd`, `sd`, and `sindhi`. Decoding in the table above is greedy CTC. Audio at any sample rate is resampled to 16 kHz mono. Files longer than 30 seconds are transcribed in chunks. If a repository ships more than one `.pt` file, pass `filename=`. ## Model diff --git a/canopy/__init__.py b/canopy/__init__.py deleted file mode 100644 index 01d06341c59da2ac14ca53331fb023023a9a7c17..0000000000000000000000000000000000000000 --- a/canopy/__init__.py +++ /dev/null @@ -1,6 +0,0 @@ -"""Canopy: multilingual non-autoregressive CTC ASR for Pakistani languages.""" - -__version__ = "0.1.0.dev0" - -BLANK_ID = 0 -SAMPLE_RATE = 16000 diff --git a/canopy/data/__init__.py b/canopy/data/__init__.py deleted file mode 100644 index c5141f3ede45d680fadae137b3e21064d1aeb5b8..0000000000000000000000000000000000000000 --- a/canopy/data/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from canopy.data.manifest import PROVENANCES, REQUIRED_FIELDS, load_manifest, make_row, validate_row - -__all__ = ["PROVENANCES", "REQUIRED_FIELDS", "load_manifest", "make_row", "validate_row"] diff --git a/canopy/data/commonvoice.py b/canopy/data/commonvoice.py deleted file mode 100644 index 0afe84aa189798a276271f71b8704e519f3ae9c2..0000000000000000000000000000000000000000 --- a/canopy/data/commonvoice.py +++ /dev/null @@ -1,113 +0,0 @@ -"""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 diff --git a/canopy/data/dataset.py b/canopy/data/dataset.py deleted file mode 100644 index f76a74b0e9fe25c7fd7939603ca658f968a7129d..0000000000000000000000000000000000000000 --- a/canopy/data/dataset.py +++ /dev/null @@ -1,104 +0,0 @@ -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), - } diff --git a/canopy/data/generic.py b/canopy/data/generic.py deleted file mode 100644 index cc9f37aa9a51a58f06a2bd287eff407c603c2613..0000000000000000000000000000000000000000 --- a/canopy/data/generic.py +++ /dev/null @@ -1,75 +0,0 @@ -"""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 diff --git a/canopy/data/manifest.py b/canopy/data/manifest.py deleted file mode 100644 index f8582acea16ad7ed79e3bea2ce5d2db9872ef9d6..0000000000000000000000000000000000000000 --- a/canopy/data/manifest.py +++ /dev/null @@ -1,124 +0,0 @@ -"""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 diff --git a/canopy/data/sampler.py b/canopy/data/sampler.py deleted file mode 100644 index 650677ff1b21b0c526018b573035a7ac08eb6948..0000000000000000000000000000000000000000 --- a/canopy/data/sampler.py +++ /dev/null @@ -1,128 +0,0 @@ -"""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) diff --git a/canopy/data/splits.py b/canopy/data/splits.py deleted file mode 100644 index 60ba47476e8f48f91f132c2512683598035025b3..0000000000000000000000000000000000000000 --- a/canopy/data/splits.py +++ /dev/null @@ -1,188 +0,0 @@ -"""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 diff --git a/canopy/data/validation.py b/canopy/data/validation.py deleted file mode 100644 index ddc8c27e68d45727b0ce5318371f97acc87c0fa9..0000000000000000000000000000000000000000 --- a/canopy/data/validation.py +++ /dev/null @@ -1,193 +0,0 @@ -"""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", "") - uid = r.get("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", "") 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 diff --git a/canopy/decoding/__init__.py b/canopy/decoding/__init__.py deleted file mode 100644 index 6411aa65da2a9c0011db71605f6f45076a48ecb0..0000000000000000000000000000000000000000 --- a/canopy/decoding/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -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"] diff --git a/canopy/decoding/beam.py b/canopy/decoding/beam.py deleted file mode 100644 index e6e999b8039df6912cb5cc2aa8671e0b54645819..0000000000000000000000000000000000000000 --- a/canopy/decoding/beam.py +++ /dev/null @@ -1,132 +0,0 @@ -"""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) diff --git a/canopy/decoding/greedy.py b/canopy/decoding/greedy.py deleted file mode 100644 index 8eba7fc40790353b68f58c365dcd17a67c2a900e..0000000000000000000000000000000000000000 --- a/canopy/decoding/greedy.py +++ /dev/null @@ -1,43 +0,0 @@ -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])] diff --git a/canopy/evaluation/__init__.py b/canopy/evaluation/__init__.py deleted file mode 100644 index d7b29ced37a57dfe7b9856d5533a9bc235b22823..0000000000000000000000000000000000000000 --- a/canopy/evaluation/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from canopy.evaluation.metrics import aggregate, corpus_error_rate, edit_counts, score_utterances - -__all__ = ["aggregate", "corpus_error_rate", "edit_counts", "score_utterances"] diff --git a/canopy/evaluation/metrics.py b/canopy/evaluation/metrics.py deleted file mode 100644 index 55cbe72c1c97d9d4a0665ca1af1332145a102257..0000000000000000000000000000000000000000 --- a/canopy/evaluation/metrics.py +++ /dev/null @@ -1,135 +0,0 @@ -"""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, - } diff --git a/canopy/evaluation/runner.py b/canopy/evaluation/runner.py deleted file mode 100644 index abc666c5e48f45035a12ff0b089499aea413fa44..0000000000000000000000000000000000000000 --- a/canopy/evaluation/runner.py +++ /dev/null @@ -1,98 +0,0 @@ -"""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 diff --git a/canopy/evaluation/status.py b/canopy/evaluation/status.py deleted file mode 100644 index 98c2e0e2da9ca1923b7d02f3ce1d2957545111ee..0000000000000000000000000000000000000000 --- a/canopy/evaluation/status.py +++ /dev/null @@ -1,79 +0,0 @@ -"""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} diff --git a/canopy/frontend/__init__.py b/canopy/frontend/__init__.py deleted file mode 100644 index 5569365d2458dabd2209c057a5d99bf9520f72fd..0000000000000000000000000000000000000000 --- a/canopy/frontend/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -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"] diff --git a/canopy/frontend/audio.py b/canopy/frontend/audio.py deleted file mode 100644 index fcea7a922574ec3dbc3c2132df12a95c41abefc6..0000000000000000000000000000000000000000 --- a/canopy/frontend/audio.py +++ /dev/null @@ -1,69 +0,0 @@ -"""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) diff --git a/canopy/frontend/augment.py b/canopy/frontend/augment.py deleted file mode 100644 index 141a09d908ac06a2aac7b884654226cc46a47b51..0000000000000000000000000000000000000000 --- a/canopy/frontend/augment.py +++ /dev/null @@ -1,56 +0,0 @@ -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) diff --git a/canopy/frontend/features.py b/canopy/frontend/features.py deleted file mode 100644 index 22cf96687d96647420ec6210952941d27f731254..0000000000000000000000000000000000000000 --- a/canopy/frontend/features.py +++ /dev/null @@ -1,111 +0,0 @@ -"""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 diff --git a/canopy/inference/__init__.py b/canopy/inference/__init__.py deleted file mode 100644 index 2c13e8b80ff50d08c9f95cea5bd63ce264876510..0000000000000000000000000000000000000000 --- a/canopy/inference/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -from canopy.inference.api import Canopy -from canopy.inference.transcriber import Transcriber - -__all__ = ["Canopy", "Transcriber"] diff --git a/canopy/inference/__main__.py b/canopy/inference/__main__.py deleted file mode 100644 index 0cf4298ac99c46c7c9abf628a2953be7e8f334b4..0000000000000000000000000000000000000000 --- a/canopy/inference/__main__.py +++ /dev/null @@ -1,3 +0,0 @@ -from canopy.inference.cli import main - -main() diff --git a/canopy/inference/api.py b/canopy/inference/api.py deleted file mode 100644 index 5da542608c77b8939baf34e0f4ceca70f07260f2..0000000000000000000000000000000000000000 --- a/canopy/inference/api.py +++ /dev/null @@ -1,54 +0,0 @@ -"""User-facing loader for every Canopy checkpoint.""" - -from __future__ import annotations - -from pathlib import Path - -import torch - -from canopy.inference.hub import resolve_repo -from canopy.inference.transcriber import Transcriber - -LANGUAGE_ALIASES = {"ur": "urd", "urdu": "urd", "sd": "snd", "sindhi": "snd"} - - -def resolve_language(language: str | None) -> str | None: - if language is None: - return None - return LANGUAGE_ALIASES.get(language.strip().lower(), language.strip().lower()) - - -class Canopy: - """One interface for every Canopy variant. - - ``model`` is a short name (``canopy-m``) or a Hub id (``ProximaAI/Canopy-M``). - A path to a local ``.pt`` file loads that file and does not contact the Hub. - """ - - def __init__(self, transcriber: Transcriber, model: str): - self.transcriber = transcriber - self.model = model - self.languages = list(transcriber.languages) - - @classmethod - def from_pretrained(cls, model: str = "canopy-m", *, filename: str | None = None, - device: str | None = None, token: str | None = None, - revision: str | None = None) -> Canopy: - if device is None: - device = "cuda" if torch.cuda.is_available() else "cpu" - path = Path(model) - if path.is_file(): - asr = Transcriber.from_checkpoint(path, device=device) - return cls(asr, str(path)) - repo_id = resolve_repo(model) - asr = Transcriber.from_pretrained( - repo_id, filename=filename, device=device, token=token, revision=revision, - ) - return cls(asr, repo_id) - - def transcribe(self, audio, language: str | None = "urd", decoder: str = "greedy", **kw) -> str: - code = resolve_language(language) - if code is not None and self.languages and code not in self.languages: - known = ", ".join(self.languages) - raise ValueError(f"{self.model} was trained on {known}, not {language!r}") - return self.transcriber.transcribe(audio, language=code, decoder=decoder, **kw) diff --git a/canopy/inference/benchmark.py b/canopy/inference/benchmark.py deleted file mode 100644 index 537137eacff2c5304f4eba679d7a6aa7c82ede84..0000000000000000000000000000000000000000 --- a/canopy/inference/benchmark.py +++ /dev/null @@ -1,64 +0,0 @@ -"""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, - } diff --git a/canopy/inference/cli.py b/canopy/inference/cli.py deleted file mode 100644 index fcf95260593c7f79208c65848f9027d5fb9d998a..0000000000000000000000000000000000000000 --- a/canopy/inference/cli.py +++ /dev/null @@ -1,36 +0,0 @@ -"""Transcribe with any published Canopy checkpoint. - -GitHub install (one package, every variant): - - pip install torch torchaudio soundfile numpy pyyaml huggingface_hub - pip install "git+https://github.com/mwzkhalil/canopy.git@danish#subdirectory=canopy" - hf auth login - canopy-transcribe clip.wav --model canopy-m --language urd - canopy-transcribe clip.wav --model ProximaAI/Canopy-M --language snd -""" - -from __future__ import annotations - -import argparse - -from canopy.inference.api import LANGUAGE_ALIASES, Canopy -from canopy.inference.hub import VARIANTS - - -def main() -> None: - names = ", ".join(sorted(VARIANTS)) - ap = argparse.ArgumentParser(description="Transcribe with a Canopy checkpoint") - ap.add_argument("audio", nargs="+", help="wav, flac, or mp3 paths") - ap.add_argument("--model", default="canopy-m", help=f"variant name ({names}) or org/repo id") - ap.add_argument("--filename", default=None, help="checkpoint filename when a repo has several") - ap.add_argument("--language", required=True, help="ISO code, for example urd or snd") - ap.add_argument("--device", default=None, help="cpu, cuda, or omit for auto") - ap.add_argument("--decoder", choices=["greedy", "beam"], default="greedy") - args = ap.parse_args() - asr = Canopy.from_pretrained(args.model, filename=args.filename, device=args.device) - for path in args.audio: - text = asr.transcribe(path, language=args.language, decoder=args.decoder) - print(f"{path}\t{text}") - - -__all__ = ["LANGUAGE_ALIASES", "main"] diff --git a/canopy/inference/hub.py b/canopy/inference/hub.py deleted file mode 100644 index 515a282804872511df7eea51436f51df39b194e0..0000000000000000000000000000000000000000 --- a/canopy/inference/hub.py +++ /dev/null @@ -1,54 +0,0 @@ -"""Resolve a Canopy variant to a Hugging Face repo and checkpoint file. - -A variant is either a short name from ``VARIANTS`` or a full ``org/repo`` id. -Add a line to ``VARIANTS`` when a new checkpoint is published. Passing -``org/repo`` directly works before that line exists. -""" - -from __future__ import annotations - -from pathlib import Path - -# Short name -> Hub repo. Weights stay on the Hub; this table is only the address. -VARIANTS: dict[str, str] = { - "canopy-m": "ProximaAI/Canopy-M", -} - - -def resolve_repo(model: str) -> str: - """Return an ``org/repo`` id. Local checkpoint paths are not repos.""" - key = model.strip() - if not key: - raise ValueError("model name is empty") - if key.lower() in VARIANTS: - return VARIANTS[key.lower()] - if "/" in key and not Path(key).exists(): - return key - known = ", ".join(sorted(VARIANTS)) - raise ValueError( - f"unknown Canopy variant {model!r}. Use one of: {known}. " - "Or pass a Hub id such as ProximaAI/Canopy-M." - ) - - -def checkpoint_name(repo_id: str) -> str: - """Default weight file: the repo name, lowercased, with a .pt suffix.""" - name = repo_id.split("/")[-1].strip().lower() - return name if name.endswith(".pt") else f"{name}.pt" - - -def pick_checkpoint(repo_id: str, files: list[str], filename: str | None = None) -> str: - """Choose one ``.pt`` at the repo root. ``filename`` wins when it is present.""" - root = [f for f in files if "/" not in f and f.endswith(".pt")] - if filename: - if filename not in files and filename not in root: - raise FileNotFoundError(f"{repo_id} has no file {filename!r}. Root checkpoints: {root or 'none'}") - return filename - preferred = checkpoint_name(repo_id) - if preferred in root: - return preferred - if len(root) == 1: - return root[0] - raise FileNotFoundError( - f"{repo_id} has {len(root)} checkpoint files {root}. Pass filename= to choose one." - ) diff --git a/canopy/inference/quantize.py b/canopy/inference/quantize.py deleted file mode 100644 index 83278e08c785463a6b8a2baf1d4a576a785a7091..0000000000000000000000000000000000000000 --- a/canopy/inference/quantize.py +++ /dev/null @@ -1,36 +0,0 @@ -"""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} diff --git a/canopy/inference/transcriber.py b/canopy/inference/transcriber.py deleted file mode 100644 index f0108f95c7323be035b0ec15d2d38ee8ead41cb2..0000000000000000000000000000000000000000 --- a/canopy/inference/transcriber.py +++ /dev/null @@ -1,109 +0,0 @@ -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 | None = None, - 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. - - ``repo_id`` is an ``org/name`` id. ``filename`` defaults to the single - ``.pt`` at the repo root, preferring ``.pt``. - """ - try: - from huggingface_hub import hf_hub_download, list_repo_files - except ImportError as exc: - raise ImportError("pip install huggingface_hub") from exc - from canopy.inference.hub import pick_checkpoint - - files = list_repo_files(repo_id, revision=revision, token=token) - chosen = pick_checkpoint(repo_id, files, filename) - path = hf_hub_download(repo_id=repo_id, filename=chosen, 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) diff --git a/canopy/language_model/__init__.py b/canopy/language_model/__init__.py deleted file mode 100644 index 22ac9507d3d30ef82b0ab027a034ff6bce15674a..0000000000000000000000000000000000000000 --- a/canopy/language_model/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -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"] diff --git a/canopy/language_model/arpa.py b/canopy/language_model/arpa.py deleted file mode 100644 index cd99bf0b9b5ab9b5b075e537893a2f8c08dc6684..0000000000000000000000000000000000000000 --- a/canopy/language_model/arpa.py +++ /dev/null @@ -1,188 +0,0 @@ -"""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 = "", "", "" - - -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 deleted file mode 100644 index dded5c7e1c865d2fab659bc5742331a0bed59447..0000000000000000000000000000000000000000 --- a/canopy/language_model/kenlm_backend.py +++ /dev/null @@ -1,62 +0,0 @@ -"""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 deleted file mode 100644 index aba2554f7212e0cd71427bbc09ee46991df8ad38..0000000000000000000000000000000000000000 --- a/canopy/language_model/text.py +++ /dev/null @@ -1,38 +0,0 @@ -"""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 deleted file mode 100644 index 2eb1c8cf9cabef21036d6a658c1729ae9bb656ab..0000000000000000000000000000000000000000 --- a/canopy/languages.py +++ /dev/null @@ -1,92 +0,0 @@ -"""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 deleted file mode 100644 index ccff4eae15fd8e1852409467ddef21f2bedffa6a..0000000000000000000000000000000000000000 --- a/canopy/models/__init__.py +++ /dev/null @@ -1,4 +0,0 @@ -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 deleted file mode 100644 index 622fe043a0888306736a53ca8c3d7e949d82333e..0000000000000000000000000000000000000000 --- a/canopy/models/canopy_ctc.py +++ /dev/null @@ -1,107 +0,0 @@ -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 deleted file mode 100644 index bfffa759c0a18e45dcc43605e085213ad5bc10bf..0000000000000000000000000000000000000000 --- a/canopy/models/config.py +++ /dev/null @@ -1,58 +0,0 @@ -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 deleted file mode 100644 index 16c60de478dc4c5d28ad4c832af840dc0df3a445..0000000000000000000000000000000000000000 --- a/canopy/models/modules.py +++ /dev/null @@ -1,183 +0,0 @@ -"""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 deleted file mode 100644 index bbc3d81f0befa228c7ea2e7ce413ebf61e3b3b48..0000000000000000000000000000000000000000 --- a/canopy/text/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from canopy.text.normalize import Normalizer, get_normalizer - -__all__ = ["Normalizer", "get_normalizer"] diff --git a/canopy/text/normalize.py b/canopy/text/normalize.py deleted file mode 100644 index 3694b348d1ba236580c9eb6610e7e227d38c6837..0000000000000000000000000000000000000000 --- a/canopy/text/normalize.py +++ /dev/null @@ -1,110 +0,0 @@ -"""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 deleted file mode 100644 index 1d9e4b6dd59cdb5d465337c2bda303566b684cd9..0000000000000000000000000000000000000000 --- a/canopy/text/rules.py +++ /dev/null @@ -1,88 +0,0 @@ -"""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 deleted file mode 100644 index 0194708649d9dee5aeb57453b371bc80f7646ca5..0000000000000000000000000000000000000000 --- a/canopy/tokenizer/__init__.py +++ /dev/null @@ -1,5 +0,0 @@ -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 deleted file mode 100644 index e84d9f783680bf07be8625615fb21fb6bd1c2d54..0000000000000000000000000000000000000000 --- a/canopy/tokenizer/base.py +++ /dev/null @@ -1,96 +0,0 @@ -"""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 deleted file mode 100644 index c43ec9ffa5ef0109654be4301e88571179a07525..0000000000000000000000000000000000000000 --- a/canopy/tokenizer/char.py +++ /dev/null @@ -1,76 +0,0 @@ -"""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 deleted file mode 100644 index f897725089da263cab153580140a22b1f89dec2f..0000000000000000000000000000000000000000 --- a/canopy/tokenizer/diagnostics.py +++ /dev/null @@ -1,95 +0,0 @@ -"""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 deleted file mode 100644 index 20b8cce8500fb40c08e01d6bdefe336224d405d4..0000000000000000000000000000000000000000 --- a/canopy/tokenizer/spm.py +++ /dev/null @@ -1,111 +0,0 @@ -"""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 deleted file mode 100644 index fd3a12ba723dc0ee25bbebf8050ec7e8133b5259..0000000000000000000000000000000000000000 --- a/canopy/training/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from canopy.training.schedule import WarmupCosine - -__all__ = ["WarmupCosine"] diff --git a/canopy/training/checkpoint.py b/canopy/training/checkpoint.py deleted file mode 100644 index f2180df0d799e68001a01c3681505c2e8fefa508..0000000000000000000000000000000000000000 --- a/canopy/training/checkpoint.py +++ /dev/null @@ -1,77 +0,0 @@ -"""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 deleted file mode 100644 index b620864386b8c72ef0db42d47bc11799fa55bccf..0000000000000000000000000000000000000000 --- a/canopy/training/losses.py +++ /dev/null @@ -1,30 +0,0 @@ -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 deleted file mode 100644 index 16b2cf027fbf89d98e83c5dcf4ddbaecb1b603d2..0000000000000000000000000000000000000000 --- a/canopy/training/run.py +++ /dev/null @@ -1,48 +0,0 @@ -"""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 deleted file mode 100644 index 5f99d082ede70ff90206847b62c5ac49838d0734..0000000000000000000000000000000000000000 --- a/canopy/training/schedule.py +++ /dev/null @@ -1,41 +0,0 @@ -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 deleted file mode 100644 index 52e4edafe8ee972ee6658a8a50dba1123858a41b..0000000000000000000000000000000000000000 --- a/canopy/training/trainer.py +++ /dev/null @@ -1,385 +0,0 @@ -"""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 deleted file mode 100644 index 2663874eccef0871a175a4c92674e34332d5f3b4..0000000000000000000000000000000000000000 --- a/canopy/utils/__init__.py +++ /dev/null @@ -1,13 +0,0 @@ -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 deleted file mode 100644 index cc308222729e3c9cda52ffda5cd522aecdab90fe..0000000000000000000000000000000000000000 --- a/canopy/utils/env.py +++ /dev/null @@ -1,64 +0,0 @@ -"""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 deleted file mode 100644 index 892752c385cbbdb83d291cfcbdd755d1b70ed3d5..0000000000000000000000000000000000000000 --- a/canopy/utils/io.py +++ /dev/null @@ -1,129 +0,0 @@ -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 deleted file mode 100644 index b6609cf00374f8e5ddcba4ad89e13efa586ddb2f..0000000000000000000000000000000000000000 --- a/canopy/utils/seed.py +++ /dev/null @@ -1,19 +0,0 @@ -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 index 5423e4bde5f87e40b74279661fb3e38756a159e8..edf90a89cf8fc8b7780a469d32e771deb8bb6c1e 100644 --- a/canopy_m.py +++ b/canopy_m.py @@ -1,18 +1,8 @@ #!/usr/bin/env python """Transcribe with any Canopy variant. -The weights for each variant live in their own Hugging Face repo. This script -is the same for all of them. The code can come from this model snapshot or -from GitHub. - -Snapshot (no GitHub): - - 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 --model canopy-m --language urd - -GitHub (one install, then any Hub repo): +The weights for each variant live in their own Hugging Face repo. Install +the library from GitHub, then run this script or canopy-transcribe. pip install git+https://github.com/Proxima-AI-Co/canopy.git hf auth login @@ -21,7 +11,7 @@ GitHub (one install, then any Hub repo): Python: - from canopy.inference import Canopy + from canopy import Canopy asr = Canopy.from_pretrained("canopy-m") asr = Canopy.from_pretrained("ProximaAI/Canopy-M") print(asr.transcribe("clip.wav", language="urd")) @@ -30,25 +20,17 @@ Python: from __future__ import annotations import argparse -import sys -from pathlib import Path - -_HERE = Path(__file__).resolve().parent def _load_api(): try: from canopy.inference.api import Canopy from canopy.inference.hub import VARIANTS - return Canopy, VARIANTS - except ImportError: - for root in (_HERE, _HERE.parent): - if (root / "canopy" / "__init__.py").is_file(): - sys.path.insert(0, str(root)) - break - from canopy.inference.api import Canopy - from canopy.inference.hub import VARIANTS - return Canopy, VARIANTS + except ImportError as exc: + raise SystemExit( + "Install the library first: pip install git+https://github.com/Proxima-AI-Co/canopy.git" + ) from exc + return Canopy, VARIANTS def main() -> None: