import hashlib import os from contextlib import nullcontext from pathlib import Path import sys import tempfile from types import ModuleType, SimpleNamespace import unittest from unittest.mock import patch from pi_dictate_service.model import NemoParakeetModel, load_coding_terms class FakeNemoModel: def __init__(self) -> None: self.eval_called = False self.to_device: str | None = None self.audio = None self.transcribe_options: dict[str, object] = {} self.cfg = SimpleNamespace(decoding={"strategy": "greedy_batch", "greedy": {"max_symbols": 10}}) self.decoding_config: dict[str, object] | None = None def eval(self) -> None: self.eval_called = True def to(self, device: str) -> None: self.to_device = device def change_decoding_strategy(self, config: dict[str, object], verbose: bool) -> None: self.decoding_config = config def transcribe(self, audio: list[object], **options: object) -> list[object]: self.audio = audio self.transcribe_options = options return [type("Hypothesis", (), {"text": "hello locally"})()] class FakeConcreteParakeetModel: restored_path: str | None = None restored_location: object | None = None restored_bytes: bytes | None = None replace_before_restore: tuple[Path, bytes] | None = None restored: FakeNemoModel | None = None @classmethod def restore_from(cls, restore_path: str, map_location: object) -> FakeNemoModel: cls.restored_path = restore_path cls.restored_location = map_location if cls.replace_before_restore is not None: path, replacement = cls.replace_before_restore replacement_path = path.with_suffix(".replacement") replacement_path.write_bytes(replacement) replacement_path.replace(path) cls.restored_bytes = Path(restore_path).read_bytes() self = FakeNemoModel() cls.restored = self return self class FakeNumpy(ModuleType): float32 = "float32" def asarray(self, samples: list[float], dtype: object) -> tuple[list[float], object]: return samples, dtype class NemoParakeetModelTests(unittest.TestCase): def setUp(self) -> None: FakeConcreteParakeetModel.restored_path = None FakeConcreteParakeetModel.restored_location = None FakeConcreteParakeetModel.restored_bytes = None FakeConcreteParakeetModel.replace_before_restore = None FakeConcreteParakeetModel.restored = None def test_load_is_lazy_offline_applies_coding_terms_and_transcribes_in_memory_audio(self) -> None: with tempfile.TemporaryDirectory() as directory: model_path = Path(directory) / "parakeet.nemo" model_path.write_bytes(b"model") terms_path = Path(directory) / "coding-terms.txt" terms_path.write_text("JSON\nnpm test\n", encoding="utf-8") modules = self.fake_modules() with patch.dict(sys.modules, modules, clear=False), patch.dict(os.environ, {}, clear=True): adapter = NemoParakeetModel( model_path, terms_path=terms_path, boost_alpha=0.5, expected_sha256=hashlib.sha256(b"model").hexdigest(), ) self.assertIsNone(FakeConcreteParakeetModel.restored_path) adapter.load() transcript = adapter.transcribe([0.0, 0.5]) self.assertEqual(os.environ["HF_HUB_OFFLINE"], "1") self.assertEqual(os.environ["TRANSFORMERS_OFFLINE"], "1") self.assertRegex(FakeConcreteParakeetModel.restored_path or "", r"^/proc/self/fd/\d+$") self.assertEqual(FakeConcreteParakeetModel.restored_bytes, b"model") self.assertEqual(FakeConcreteParakeetModel.restored_location, "device:cuda") self.assertTrue(FakeConcreteParakeetModel.restored.eval_called) self.assertEqual(FakeConcreteParakeetModel.restored.to_device, "cuda") self.assertEqual(transcript, "hello locally") self.assertEqual(FakeConcreteParakeetModel.restored.audio, [([0.0, 0.5], "float32")]) self.assertEqual( FakeConcreteParakeetModel.restored.decoding_config["greedy"]["boosting_tree"]["key_phrases_list"], ["JSON", "npm test"], ) self.assertEqual( FakeConcreteParakeetModel.restored.decoding_config["greedy"]["boosting_tree_alpha"], 0.5, ) self.assertEqual( FakeConcreteParakeetModel.restored.transcribe_options, {"batch_size": 1, "num_workers": 0, "verbose": False}, ) def test_loads_a_bounded_deduplicated_host_dictionary(self) -> None: with tempfile.TemporaryDirectory() as directory: terms_path = Path(directory) / "coding-terms.txt" terms_path.write_text("# local terms\n JSON \n\nnpm\nJSON\n", encoding="utf-8") self.assertEqual(load_coding_terms(terms_path), ["JSON", "npm"]) terms_path.write_text("x" * 129, encoding="utf-8") with self.assertRaisesRegex(ValueError, "128 UTF-8 bytes"): load_coding_terms(terms_path) def test_load_uses_the_verified_inode_if_the_configured_path_is_replaced(self) -> None: with tempfile.TemporaryDirectory() as directory: model_path = Path(directory) / "parakeet.nemo" model_path.write_bytes(b"verified model") FakeConcreteParakeetModel.replace_before_restore = (model_path, b"replacement") with patch.dict(sys.modules, self.fake_modules(), clear=False): adapter = NemoParakeetModel( model_path, expected_sha256=hashlib.sha256(b"verified model").hexdigest(), ) adapter.load() self.assertEqual(model_path.read_bytes(), b"replacement") self.assertEqual(FakeConcreteParakeetModel.restored_bytes, b"verified model") def test_refuses_a_model_that_no_longer_matches_the_pinned_hash(self) -> None: with tempfile.TemporaryDirectory() as directory: model_path = Path(directory) / "parakeet.nemo" model_path.write_bytes(b"replaced") adapter = NemoParakeetModel(model_path) with self.assertRaisesRegex(ValueError, "model checksum mismatch"): adapter.load() def test_requires_an_existing_local_nemo_file(self) -> None: with self.assertRaisesRegex(ValueError, "local .nemo"): NemoParakeetModel("/does/not/exist") @staticmethod def fake_modules() -> dict[str, ModuleType]: nemo = ModuleType("nemo") collections = ModuleType("nemo.collections") asr = ModuleType("nemo.collections.asr") models = ModuleType("nemo.collections.asr.models") models.EncDecRNNTBPEModel = FakeConcreteParakeetModel asr.models = models collections.asr = asr nemo.collections = collections omegaconf = ModuleType("omegaconf") omegaconf.open_dict = lambda _config: nullcontext() torch = ModuleType("torch") torch.device = lambda name: f"device:{name}" return { "nemo": nemo, "nemo.collections": collections, "nemo.collections.asr": asr, "nemo.collections.asr.models": models, "numpy": FakeNumpy("numpy"), "omegaconf": omegaconf, "torch": torch, } if __name__ == "__main__": unittest.main()