feat(dictate): use local socket-activated Parakeet

services/dictate/src/pi_dictate_service/model.py

@@ -0,0 +1,138 @@
+"""Lazy, offline NeMo adapter for an explicit local Parakeet archive."""
+
+from __future__ import annotations
+
+from copy import deepcopy
+import hashlib
+import os
+from pathlib import Path
+from typing import Protocol, Sequence
+
+MAX_TERMS_FILE_BYTES = 64 * 1024
+MAX_CODING_TERMS = 512
+MAX_TERM_BYTES = 128
+PINNED_MODEL_SHA256 = "3cbdc85877e668ca7b82d0d56770eb1fac76691f55d6b97545e8d61ca588d10d"
+
+
+class TranscriptionModel(Protocol):
+    def load(self) -> None: ...
+    def transcribe(self, samples: Sequence[float]) -> str: ...
+
+
+def load_coding_terms(path: str | os.PathLike[str]) -> list[str]:
+    """Load a bounded host-owned phrase list, ignoring comments and duplicates."""
+
+    terms_path = Path(path).expanduser()
+    if not terms_path.is_file():
+        raise ValueError("PI_DICTATE_TERMS_PATH must name an existing local dictionary")
+    if terms_path.stat().st_size > MAX_TERMS_FILE_BYTES:
+        raise ValueError(f"coding dictionary exceeds {MAX_TERMS_FILE_BYTES} bytes")
+
+    terms: list[str] = []
+    seen: set[str] = set()
+    for raw_line in terms_path.read_text(encoding="utf-8").splitlines():
+        term = raw_line.strip()
+        if not term or term.startswith("#"):
+            continue
+        if len(term.encode("utf-8")) > MAX_TERM_BYTES:
+            raise ValueError(f"coding dictionary terms may not exceed {MAX_TERM_BYTES} UTF-8 bytes")
+        if term not in seen:
+            terms.append(term)
+            seen.add(term)
+        if len(terms) > MAX_CODING_TERMS:
+            raise ValueError(f"coding dictionary may not exceed {MAX_CODING_TERMS} terms")
+    return terms
+
+
+class NemoParakeetModel:
+    """Load NeMo only when requested so core tests require neither CUDA nor NeMo."""
+
+    def __init__(
+        self,
+        model_path: str | os.PathLike[str],
+        *,
+        terms_path: str | os.PathLike[str] | None = None,
+        boost_alpha: float = 1.0,
+        expected_sha256: str = PINNED_MODEL_SHA256,
+    ) -> None:
+        path = Path(model_path).expanduser()
+        if path.suffix != ".nemo" or not path.is_file():
+            raise ValueError("PI_DICTATE_MODEL_PATH must name an existing local .nemo file")
+        if boost_alpha <= 0:
+            raise ValueError("PI_DICTATE_BOOST_ALPHA must be positive")
+        self._model_path = path.resolve()
+        self._terms = load_coding_terms(terms_path) if terms_path is not None else []
+        self._boost_alpha = boost_alpha
+        self._expected_sha256 = expected_sha256
+        self._model: object | None = None
+        self._numpy: object | None = None
+
+    @property
+    def loaded(self) -> bool:
+        return self._model is not None
+
+    def load(self) -> None:
+        if self._model is not None:
+            return
+        with self._model_path.open("rb") as model_file:
+            actual_sha256 = hashlib.file_digest(model_file, "sha256").hexdigest()
+        if actual_sha256 != self._expected_sha256:
+            raise ValueError(
+                "model checksum mismatch: expected the pinned Parakeet revision"
+            )
+        os.environ["HF_HUB_OFFLINE"] = "1"
+        os.environ["TRANSFORMERS_OFFLINE"] = "1"
+        os.environ["HF_DATASETS_OFFLINE"] = "1"
+        os.environ["NEMO_OFFLINE"] = "1"
+
+        import numpy
+        from nemo.collections.asr.models import ASRModel
+        from omegaconf import open_dict
+
+        model = ASRModel.restore_from(
+            restore_path=str(self._model_path),
+            map_location="cuda",
+        )
+        model.to("cuda")
+        model.eval()
+        if self._terms:
+            decoding = deepcopy(model.cfg.decoding)
+            with open_dict(decoding):
+                greedy = decoding["greedy"]
+                greedy["boosting_tree"] = {
+                    "key_phrases_list": list(self._terms),
+                    "context_score": 1.0,
+                    "depth_scaling": 2.0,
+                    "unk_score": 0.0,
+                    "final_eos_score": 1.0,
+                    "score_per_phrase": 0.0,
+                    "source_lang": "en",
+                    "use_triton": True,
+                    "uniform_weights": False,
+                    "use_bpe_dropout": False,
+                }
+                greedy["boosting_tree_alpha"] = self._boost_alpha
+            model.change_decoding_strategy(decoding, verbose=False)
+        self._numpy = numpy
+        self._model = model
+
+    def transcribe(self, samples: Sequence[float]) -> str:
+        if self._model is None or self._numpy is None:
+            raise RuntimeError("model is not loaded")
+        numpy = self._numpy
+        audio = numpy.asarray(samples, dtype=numpy.float32)  # type: ignore[attr-defined]
+        output: Sequence[object] = self._model.transcribe(  # type: ignore[attr-defined]
+            audio=[audio],
+            batch_size=1,
+            num_workers=0,
+            verbose=False,
+        )
+        if not output:
+            return ""
+        hypothesis = output[0]
+        if isinstance(hypothesis, str):
+            return hypothesis
+        text = getattr(hypothesis, "text", None)
+        if isinstance(text, str):
+            return text
+        raise RuntimeError("NeMo returned an unsupported transcription result")