Load any Canopy variant from the Hub or from GitHub

README.md

@@ -49,27 +49,38 @@ One block is half SwiGLU, rotary attention, a squeeze-excite depthwise convoluti
 
 ## Use
 
-The model repository is private. Log in once with a token that can read `ProximaAI/Canopy-M`.
+Each Canopy variant is its own Hub repository. This one is `canopy-m` (`ProximaAI/Canopy-M`). The same loader accepts a short name or any `org/repo` id, so the next checkpoint does not need a new API. Log in once with a token that can read the repo.
+
+From this snapshot, without GitHub:
 
 ```bash
 pip install torch torchaudio soundfile numpy pyyaml huggingface_hub
 hf auth login
 hf download ProximaAI/Canopy-M --local-dir Canopy-M
-python Canopy-M/canopy_m.py clip.wav --language urd
-python Canopy-M/canopy_m.py clip.wav --language snd --device cpu
+python Canopy-M/canopy_m.py clip.wav --model canopy-m --language urd
+python Canopy-M/canopy_m.py clip.wav --model ProximaAI/Canopy-M --language snd --device cpu
 ```
 
-`--language` accepts `urd`, `ur`, `urdu`, `snd`, `sd`, or `sindhi`. That selects the language embedding and keeps the CTC output inside that language’s letters. Decoding is greedy CTC, the same setup as the table above.
+From GitHub, one install for every variant:
+
+```bash
+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
+```
 
 ```python
-from canopy_m import CanopyM
+from canopy.inference import Canopy
 
-asr = CanopyM.from_pretrained("ProximaAI/Canopy-M")  # cuda when available, otherwise cpu
+asr = Canopy.from_pretrained("canopy-m")                 # ProximaAI/Canopy-M
+asr = Canopy.from_pretrained("ProximaAI/Canopy-M")       # same checkpoint
 print(asr.transcribe("clip.wav", language="urd"))
 print(asr.transcribe("clip.wav", language="snd"))
 ```
 
-Run that from the `Canopy-M` directory, or add it to `PYTHONPATH`. `canopy_m.py` and the `canopy` package are both in this repository, next to `canopy-m.pt`. A 16 kHz mono file is the expected input. Other sample rates are resampled. Files longer than 30 seconds are transcribed in chunks.
+`language` accepts `urd` / `ur` / `urdu` and `snd` / `sd` / `sindhi` for this checkpoint. A future repo is loaded with its own id, for example `Canopy.from_pretrained("ProximaAI/Canopy-S")`, once that repository exists. If a repository contains more than one `.pt` file, pass `filename=`. Decoding here is greedy CTC. Audio may be any sample rate; it is resampled to 16 kHz. Files longer than 30 seconds are transcribed in chunks.
 
 ## Training data
 

canopy/inference/__init__.py

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

canopy/inference/__main__.py

@@ -0,0 +1,3 @@
+from canopy.inference.cli import main
+
+main()

canopy/inference/api.py

@@ -0,0 +1,54 @@
+"""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)

canopy/inference/cli.py

@@ -0,0 +1,36 @@
+"""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"]

canopy/inference/hub.py

@@ -0,0 +1,54 @@
+"""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."
+    )

canopy/inference/transcriber.py

@@ -50,14 +50,22 @@ class Transcriber:
         return cls(model, _tokenizer_from_checkpoint(ck, tokenizer_dir), device=device, **kw)
 
     @classmethod
-    def from_pretrained(cls, repo_id: str = "ProximaAI/Canopy-M", *, filename: str = "canopy-m.pt",
+    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."""
+        """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 ``<repo-name>.pt``.
+        """
         try:
-            from huggingface_hub import hf_hub_download
+            from huggingface_hub import hf_hub_download, list_repo_files
         except ImportError as exc:
             raise ImportError("pip install huggingface_hub") from exc
-        path = hf_hub_download(repo_id=repo_id, filename=filename, revision=revision, token=token)
+        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:

canopy_m.py

@@ -1,26 +1,31 @@
 #!/usr/bin/env python
-"""Transcribe Urdu or Sindhi with ProximaAI/Canopy-M.
+"""Transcribe with any Canopy variant.
 
-Install once:
+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
 
-Download this repository (it contains this script and the Canopy code), then:
+GitHub (one install, then any Hub repo):
 
-    hf download ProximaAI/Canopy-M --local-dir Canopy-M
-    python Canopy-M/canopy_m.py clip.wav --language urd
-    python Canopy-M/canopy_m.py clip.wav --language snd --device cpu
+    pip install "git+https://github.com/mwzkhalil/canopy.git@danish#subdirectory=canopy"
+    pip install huggingface_hub
+    hf auth login
+    canopy-transcribe clip.wav --model canopy-m --language urd
+    canopy-transcribe clip.wav --model ProximaAI/Canopy-M --language snd
 
-From Python, with this directory on the path:
+Python:
 
-    from canopy_m import CanopyM
-    asr = CanopyM.from_pretrained("ProximaAI/Canopy-M")
+    from canopy.inference import Canopy
+    asr = Canopy.from_pretrained("canopy-m")
+    asr = Canopy.from_pretrained("ProximaAI/Canopy-M")
     print(asr.transcribe("clip.wav", language="urd"))
-    print(asr.transcribe("clip.wav", language="sindhi"))
-
-``language`` is ``urd`` or ``snd``. ``ur`` / ``urdu`` and ``sd`` / ``sindhi`` are accepted.
-The published scores use greedy CTC with the language id on.
 """
 
 from __future__ import annotations
@@ -30,58 +35,36 @@ import sys
 from pathlib import Path
 
 _HERE = Path(__file__).resolve().parent
-for _root in (_HERE, _HERE.parent):
-    if (_root / "canopy" / "__init__.py").is_file():
-        sys.path.insert(0, str(_root))
-        break
-
-import torch
-
-from canopy.inference.transcriber import Transcriber
-
-REPO_ID = "ProximaAI/Canopy-M"
-ALIASES = {"ur": "urd", "urdu": "urd", "urd": "urd", "sd": "snd", "sindhi": "snd", "snd": "snd"}
-
-
-def resolve_language(language: str) -> str:
-    key = language.strip().lower()
-    if key not in ALIASES:
-        known = "urd, snd (also ur, urdu, sd, sindhi)"
-        raise ValueError(f"unknown language {language!r}. Use one of: {known}")
-    return ALIASES[key]
-
-
-class CanopyM:
-    """Urdu and Sindhi CTC recogniser backed by the Canopy-M checkpoint."""
-
-    def __init__(self, transcriber: Transcriber):
-        self.transcriber = transcriber
 
-    @classmethod
-    def from_pretrained(cls, repo_id: str = REPO_ID, *, device: str | None = None, token: str | None = None) -> CanopyM:
-        if device is None:
-            device = "cuda" if torch.cuda.is_available() else "cpu"
-        local = _HERE / "canopy-m.pt"
-        if local.is_file():
-            asr = Transcriber.from_checkpoint(local, device=device)
-        else:
-            asr = Transcriber.from_pretrained(repo_id, device=device, token=token)
-        return cls(asr)
 
-    def transcribe(self, audio, language: str = "urd", decoder: str = "greedy") -> str:
-        return self.transcriber.transcribe(audio, language=resolve_language(language), decoder=decoder)
+def _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
 
 
 def main() -> None:
-    ap = argparse.ArgumentParser(description="Transcribe with ProximaAI/Canopy-M")
+    Canopy, variants = _load_api()
+    ap = argparse.ArgumentParser(description="Transcribe with a Canopy checkpoint")
     ap.add_argument("audio", nargs="+", help="wav, flac, or mp3 paths")
-    ap.add_argument("--language", required=True, help="urd or snd")
+    ap.add_argument("--model", default="canopy-m", help="variant name or org/repo. Known: " + ", ".join(sorted(variants)))
+    ap.add_argument("--filename", default=None, help="checkpoint file when a repo contains several")
+    ap.add_argument("--language", required=True, help="language code, for example urd or snd")
     ap.add_argument("--device", default=None, help="cpu, cuda, or omit for auto")
-    ap.add_argument("--repo", default=REPO_ID)
+    ap.add_argument("--decoder", choices=["greedy", "beam"], default="greedy")
     args = ap.parse_args()
-    asr = CanopyM.from_pretrained(args.repo, device=args.device)
+    asr = Canopy.from_pretrained(args.model, filename=args.filename, device=args.device)
     for path in args.audio:
-        print(f"{path}\t{asr.transcribe(path, language=args.language)}")
+        print(f"{path}\t{asr.transcribe(path, language=args.language, decoder=args.decoder)}")
 
 
 if __name__ == "__main__":