182 lines
7.6 KiB
Python
182 lines
7.6 KiB
Python
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()
|