Merge pull request #12 from NetherlandsForensicInstitute/serialize-models

BSZ-213: Add functionality to serialize/deserialize LR models for given mark/score types

Author: JoBraNFI

.gitignore

@@ -2,6 +2,7 @@
 # Project specific
 lrsystem_output/
 .testmondata
+tests/test_model_storage/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

lrmodule/__init__.py

@@ -6,34 +6,32 @@
 from lir.lrsystems.lrsystems import LRSystem
 
 from lrmodule import persistence
-from lrmodule.data import get_dataset_id
 from lrmodule.data_types import ModelSettings
 from lrmodule.lrsystem import get_trained_model
 
 
-def get_model(settings: ModelSettings, training_data: FeatureData, cache_dir: Path | None) -> LRSystem:
+def get_model(settings: ModelSettings, training_data: FeatureData, model_storage_path: Path | None) -> LRSystem:
     """
     Obtain a model by loading it from disk, or by fitting it from training data.
 
     :param settings: model settings
     :param training_data: training data
-    :param cache_dir: cache dir
+    :param model_storage_path: path where trained LR models are stored
     :return: a fitted LR system
     """
-    dataset_id = get_dataset_id(training_data)
-    model = None if not cache_dir else persistence.load_model(settings, dataset_id, cache_dir)
+    model = None if not model_storage_path else persistence.load_model(settings, model_storage_path)
     if not model:
         model = get_trained_model(settings, training_data)
-        if cache_dir:
-            persistence.save_model(model, settings, dataset_id, cache_dir)
+        if model_storage_path:
+            persistence.save_model(model, settings, model_storage_path)
     return model
 
 
 def calculate_llrs(
-    features: np.ndarray, settings: ModelSettings, training_data: FeatureData, cache_dir: Path | None
+    features: np.ndarray, settings: ModelSettings, training_data: FeatureData, model_storage_path: Path | None
 ) -> LLRData:
     """Calculate LLRs after fitting a model with a training set."""
-    model = get_model(settings, training_data, cache_dir)
+    model = get_model(settings, training_data, model_storage_path)
     return model.apply(FeatureData(features=features))
 
 

lrmodule/data.py

@@ -1,25 +1,46 @@
-from hashlib import sha256
+import os
+import pickle
 from pathlib import Path
 
 from lir.lrsystems.lrsystems import LRSystem
 
 from lrmodule.data_types import ModelSettings
 
 
-def _get_model_dirname(settings: ModelSettings, dataset_id: str) -> str:
-    h = sha256()
-    h.update(str(settings).encode("utf8"))
-    h.update(dataset_id.encode("utf8"))
-    return h.hexdigest()
+def _get_model_filename(settings: ModelSettings) -> str:
+    """Construct model filename based on mark and score type."""
+    mark_type = settings.mark_type.value
+    score_type = settings.score_type.value
 
+    return f"{mark_type}_{score_type}_model.pkl"
 
-def load_model(settings: ModelSettings, dataset_id: str, cache_dir: Path) -> LRSystem | None:
+
+def load_model(settings: ModelSettings, model_storage_path: Path) -> LRSystem:
     """Load previously cached model."""
-    _ = cache_dir / _get_model_dirname(settings, dataset_id) / "model.pkl"
-    raise NotImplementedError
+    model_filename = _get_model_filename(settings)
+    model_file_path = model_storage_path / model_filename
+
+    mark_type = settings.mark_type.value
+    score_type = settings.score_type.value
+
+    if not model_file_path.exists():
+        raise FileNotFoundError(f"No model found for mark type '{mark_type}', score type: '{score_type}'.")
 
+    try:
+        with open(model_file_path, "rb") as f:
+            # It is assumed exclusively `LRSystem` models will be loaded, which are considered safe
+            return pickle.load(f)  # noqa: S301
+    except Exception:
+        raise RuntimeError(
+            f"Could not load model from .pkl file for mark type '{mark_type}', score type: '{score_type}'"
+        )
 
-def save_model(model: LRSystem, settings: ModelSettings, dataset_id: str, cache_dir: Path) -> None:
+
+def save_model(model: LRSystem, settings: ModelSettings, model_storage_path: Path) -> None:
     """Save a model to disk."""
-    _ = cache_dir / _get_model_dirname(settings, dataset_id)
-    raise NotImplementedError
+    model_filename = _get_model_filename(settings)
+    model_file_path = model_storage_path / model_filename
+
+    os.makedirs(os.path.dirname(model_file_path), exist_ok=True)
+    with open(model_file_path, "wb") as f:
+        f.write(pickle.dumps(model))

lrmodule/persistence.py

@@ -0,0 +1,32 @@
+import pytest
+from lir.data.datasets.synthesized_normal_binary import SynthesizedNormalBinaryData, SynthesizedNormalDataClass
+from lir.data.models import FeatureData
+from lir.lrsystems.lrsystems import LRSystem
+
+from lrmodule.data_types import MarkType, ModelSettings, ScoreType
+from lrmodule.lrsystem import load_lrsystem
+
+
+@pytest.fixture
+def sample_feature_data() -> FeatureData:
+    """Provide FeatureData collection of synthesized normal binary data."""
+    data = SynthesizedNormalBinaryData(
+        data_classes={
+            0: SynthesizedNormalDataClass(mean=-1, std=1, size=100),
+            1: SynthesizedNormalDataClass(mean=1, std=1, size=100),
+        },
+        seed=0,
+    )
+    data = data.get_instances()
+    data = data.replace(features=data.features.flatten())
+
+    return data
+
+
+@pytest.fixture
+def trained_lr_system(sample_feature_data: FeatureData) -> LRSystem:
+    """Provide a basic trained LR system model based on specific settings and data."""
+    lrsystem = load_lrsystem(ModelSettings(MarkType.FIRING_PIN_IMPRESSION, ScoreType.ACCF))
+    lrsystem.fit(sample_feature_data)
+
+    return lrsystem

tests/conftest.py

@@ -1,27 +1,7 @@
-from lir.data.datasets.synthesized_normal_binary import SynthesizedNormalBinaryData, SynthesizedNormalDataClass
-from lrmodule.data_types import MarkType, ModelSettings, ScoreType
-from lrmodule.lrsystem import load_lrsystem
+from lir.data.models import FeatureData
+from lir.lrsystems.lrsystems import LRSystem
 
 
-def test_load_lrsystem():
-    load_lrsystem(ModelSettings(MarkType.FIRING_PIN_IMPRESSION, ScoreType.ACCF))
-
-
-def test_run_lrsystem():
-    lrsystem = load_lrsystem(ModelSettings(MarkType.FIRING_PIN_IMPRESSION, ScoreType.ACCF))
-    data = SynthesizedNormalBinaryData(
-        data_classes={
-            0: SynthesizedNormalDataClass(mean=-1, std=1, size=100),
-            1: SynthesizedNormalDataClass(mean=1, std=1, size=100),
-        },
-        seed=0,
-    )
-    data = data.get_instances()
-    data = data.replace(features=data.features.flatten())
-    llrs = lrsystem.fit(data).apply(data)
+def test_run_lrsystem(trained_lr_system: LRSystem, sample_feature_data: FeatureData):
+    llrs = trained_lr_system.apply(sample_feature_data)
     assert llrs.features.shape == (200, 3)
-
-
-if __name__ == "__main__":
-    test_load_lrsystem()
-    test_run_lrsystem()

tests/test_data.py

@@ -1,20 +1,107 @@
+import os
+import shutil
 from pathlib import Path
+from pickle import UnpicklingError
+from unittest import mock
 
 import pytest
+from lir.data.models import FeatureData
+from lir.lrsystems.lrsystems import LRSystem
+
 from lrmodule import ModelSettings
 from lrmodule.data_types import MarkType, ScoreType
-from lrmodule.lrsystem import load_lrsystem
-from lrmodule.persistence import load_model, save_model
+from lrmodule.persistence import load_model, save_model, _get_model_filename
+
+
+MODEL_STORAGE_PATH = Path(__file__).parent / "test_model_storage"
+
+
+@pytest.fixture(autouse=True)
+def clear_test_model_storage_directory():
+    """Clean up 'test_model_storage' directory before running each test.
+
+    This ensures a fresh environment for each test. The generated artifacts are not
+    cleaned up after each test to allow easy debugging of the generated pickle files.
+    """
+    if MODEL_STORAGE_PATH.exists():
+        shutil.rmtree(MODEL_STORAGE_PATH)
+    MODEL_STORAGE_PATH.mkdir(parents=True)
+
+
+def test_serialize_trained_lr_system(trained_lr_system: LRSystem):
+    """Check that a trained LR system can be serialized."""
+    # Given that we have a trained LR system
+    settings = ModelSettings(MarkType.FIRING_PIN_IMPRESSION, ScoreType.ACCF)
+    mark_type = settings.mark_type.value
+    score_type = settings.score_type.value
+
+    # When we serialize the LR system
+    save_model(trained_lr_system, settings, MODEL_STORAGE_PATH)
+
+    # There should be a file we can load
+    model_filename = _get_model_filename(settings)
+    model_file_path = MODEL_STORAGE_PATH / model_filename
+    assert model_file_path.exists()
+
+
+def test_deserialize_trained_lr_system(trained_lr_system: LRSystem, sample_feature_data: FeatureData):
+    """Check that a deserialized, trained LR system yields exactly the same results."""
+    # Given that we have a certain LR system serialized
+    settings = ModelSettings(MarkType.FIRING_PIN_IMPRESSION, ScoreType.ACCF)
+    save_model(trained_lr_system, settings, MODEL_STORAGE_PATH)
+
+    # When the model is deserialized
+    deserialized_model = load_model(settings, MODEL_STORAGE_PATH)
+
+    # The deserialized model and the model it originated from should be of the same type of LR system
+    assert type(trained_lr_system) == type(trained_lr_system)
+
+    # The calculated LLR output should be identical to the LR system output of the serialized model
+    expected_llr_data = trained_lr_system.apply(sample_feature_data)
+    deserialized_model_data = deserialized_model.apply(sample_feature_data)
+
+    assert deserialized_model_data == expected_llr_data
+
+
+@pytest.mark.parametrize('mark_type,score_type', [
+    (MarkType.FIRING_PIN_IMPRESSION, ScoreType.CMC),  # other score type
+    (MarkType.BREECH_PIN_IMPRESSION, ScoreType.ACCF),  # other mark type
+    (MarkType.BREECH_PIN_IMPRESSION, ScoreType.CMC),  # other mark and other score type
+])
+def test_deserialize_inexistent_lr_system(trained_lr_system: LRSystem, mark_type: MarkType, score_type: ScoreType):
+    """Check that an appropriate error is raised when there is no serialized model."""
+    # Given that the LR model storage directory is empty
+    assert os.listdir(MODEL_STORAGE_PATH) == []
+
+    # Given that we have a serialized model for a given type of `ModelSettings`
+    settings = ModelSettings(MarkType.FIRING_PIN_IMPRESSION, ScoreType.ACCF)
+    save_model(trained_lr_system, settings, MODEL_STORAGE_PATH)
+    assert len(os.listdir(MODEL_STORAGE_PATH)) == 1
+
+    # When we try to deserialize a model for a different type of `ModelSettings`
+    other_settings = ModelSettings(mark_type, score_type)
+
+    # An exception should be raised mentioning that we can't find that particular deserialized LR model
+    with pytest.raises(FileNotFoundError) as exception_info:
+        load_model(other_settings, MODEL_STORAGE_PATH)
+
+    # The exception should mention no models found for the requested mark/score types
+    assert "No model found for mark type" in str(exception_info.value)
+    assert other_settings.mark_type.value in str(exception_info.value)
+    assert other_settings.score_type.value in str(exception_info.value)
 
 
-def test_persistence():
+def test_deserialize_from_invalid_pickle_file(trained_lr_system: LRSystem):
+    """Check that an appropriate error is raised when unable to unpickle serialized model."""
+    # Given that we have a serialized model for a given type of `ModelSettings`
     settings = ModelSettings(MarkType.FIRING_PIN_IMPRESSION, ScoreType.ACCF)
+    save_model(trained_lr_system, settings, MODEL_STORAGE_PATH)
 
-    # not implemented
-    with pytest.raises(Exception):
-        load_model(settings, "dataset_id", Path("/"))
+    with mock.patch('pickle.load', side_effect=UnpicklingError("Some pickle error")):
+        # When pickle can't load the given file, we expect an appropriate error to be raised
+        with pytest.raises(RuntimeError) as exception_info:
+            load_model(settings, MODEL_STORAGE_PATH)
 
-    # not implemented
-    lrsystem = load_lrsystem(settings)
-    with pytest.raises(Exception):
-        save_model(lrsystem, settings, "dataset_id", Path("/"))
+        assert "Could not load model from .pkl file for mark type" in str(exception_info.value)
+        assert settings.mark_type.value in str(exception_info.value)
+        assert settings.score_type.value in str(exception_info.value)