140 lines
5.1 KiB
Python
140 lines
5.1 KiB
Python
"""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
|
|
import torch
|
|
from nemo.collections.asr.models import EncDecRNNTBPEModel
|
|
from omegaconf import open_dict
|
|
|
|
model = EncDecRNNTBPEModel.restore_from(
|
|
restore_path=str(self._model_path),
|
|
map_location=torch.device("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")
|