pipeline.py

"""One-call Arabic speech-to-text with forced diacritization (tashkeel).

Wraps the NeMo ASR model and the CATT diacritizer as a single object so
transcription -> diacritization happens as one process, one call.
"""
from __future__ import annotations

import nemo.collections.asr as nemo_asr

from diacritize import Diacritizer


class DiacritizedASR:
    """Arabic speech-to-text that returns fully diacritized transcripts."""

    def __init__(
        self,
        nemo_path: str = "stt_ar_fastconformer_hybrid_large_pcd_v1.0.nemo",
        catt_ckpt: str = "best_ed_mlm_ns_epoch_178.pt",
        device: str | None = None,
    ) -> None:
        self.asr_model = nemo_asr.models.EncDecHybridRNNTCTCBPEModel.restore_from(nemo_path)
        self.asr_model.eval()
        self.diacritizer = Diacritizer(ckpt=catt_ckpt, device=device)

    def transcribe(self, audio_path: str) -> str:
        """Audio file path -> diacritized Arabic text."""
        plain = self.asr_model.transcribe([audio_path])[0].text
        return self.diacritizer.diacritize_texts([plain])[0]


if __name__ == "__main__":
    import sys

    model = DiacritizedASR()
    print(model.transcribe(sys.argv[1]))