#!/usr/bin/env python
"""Transcribe with any Canopy variant.

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
    canopy-transcribe clip.wav --model canopy-m --language urd
    canopy-transcribe clip.wav --model ProximaAI/Canopy-M --language snd

Python:

    from canopy import Canopy
    asr = Canopy.from_pretrained("canopy-m")
    asr = Canopy.from_pretrained("ProximaAI/Canopy-M")
    print(asr.transcribe("clip.wav", language="urd"))
"""

from __future__ import annotations

import argparse


def _load_api():
    try:
        from canopy.inference.api import Canopy
        from canopy.inference.hub import 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:
    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("--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("--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:
        print(f"{path}\t{asr.transcribe(path, language=args.language, decoder=args.decoder)}")


if __name__ == "__main__":
    main()