enhancement: cache named models (#155)

Fixes the memory issue that was discovered in #152.

* Caches named models as they are loaded to prevent reloading or creating multiple copies.
* Adds test that fails if model is initialized more than once (fails on main).
* This should also address concerns that Chipper was being loaded multiple times.

Author: qued

CHANGELOG.md

@@ -1,4 +1,6 @@
-## 0.5.8-dev0
+## 0.5.8-dev1
+
+* Cache named models that have been lodaed
 
 ## 0.5.7
 

test_unstructured_inference/models/test_detectron2.py

@@ -17,14 +17,16 @@ def detect(self, x):
 
 def test_load_default_model(monkeypatch):
     monkeypatch.setattr(detectron2, "Detectron2LayoutModel", MockDetectron2LayoutModel)
+    monkeypatch.setattr(models, "models", {})
 
     with patch.object(detectron2, "is_detectron2_available", return_value=True):
         model = models.get_model("detectron2_lp")
 
     assert isinstance(model.model, MockDetectron2LayoutModel)
 
 
-def test_load_default_model_raises_when_not_available():
+def test_load_default_model_raises_when_not_available(monkeypatch):
+    monkeypatch.setattr(models, "models", {})
     with patch.object(detectron2, "is_detectron2_available", return_value=False):
         with pytest.raises(ImportError):
             models.get_model("detectron2_lp")

test_unstructured_inference/models/test_detectron2onnx.py

@@ -23,7 +23,8 @@ class input_thing:
         return [input_thing()]
 
 
-def test_load_default_model():
+def test_load_default_model(monkeypatch):
+    monkeypatch.setattr(models, "models", {})
     with patch.object(
         detectron2.onnxruntime,
         "InferenceSession",

test_unstructured_inference/models/test_model.py

@@ -1,15 +1,33 @@
+from typing import Any
+from unittest import mock
+
 import pytest
 
 import unstructured_inference.models.base as models
-from unstructured_inference.models.unstructuredmodel import ModelNotInitializedError
+from unstructured_inference.models.unstructuredmodel import (
+    ModelNotInitializedError,
+    UnstructuredObjectDetectionModel,
+)
+
+
+class MockModel(UnstructuredObjectDetectionModel):
+    call_count = 0
 
+    initialize = mock.MagicMock()
+
+    def __init__(self):
+        self.initializer = mock.MagicMock()
+        super().__init__()
 
-class MockModel:
     def initialize(self, *args, **kwargs):
-        pass
+        return self.initializer(self, *args, **kwargs)
+
+    def predict(self, x: Any) -> Any:
+        return []
 
 
 def test_get_model(monkeypatch):
+    monkeypatch.setattr(models, "models", {})
     monkeypatch.setattr(
         models,
         "UnstructuredDetectronModel",
@@ -36,3 +54,16 @@ def test_raises_invalid_model():
 def test_raises_uninitialized():
     with pytest.raises(ModelNotInitializedError):
         models.UnstructuredDetectronModel().predict(None)
+
+
+def test_model_initializes_once():
+    from unstructured_inference.inference import layout
+
+    with mock.patch.object(models, "UnstructuredDetectronONNXModel", MockModel), mock.patch.object(
+        models,
+        "models",
+        {},
+    ):
+        doc = layout.DocumentLayout.from_file("sample-docs/layout-parser-paper.pdf")
+
+        doc.pages[0].detection_model.initializer.assert_called_once()

unstructured_inference/__version__.py

@@ -1 +1 @@
-__version__ = "0.5.8-dev0"  # pragma: no cover
+__version__ = "0.5.8-dev1"  # pragma: no cover

unstructured_inference/models/base.py

@@ -1,4 +1,4 @@
-from typing import Optional
+from typing import Dict, Optional
 
 from unstructured_inference.logger import logger
 from unstructured_inference.models.chipper import MODEL_TYPES as CHIPPER_MODEL_TYPES
@@ -25,14 +25,22 @@
 
 DEFAULT_MODEL = "detectron2_onnx"
 
+models: Dict[str, UnstructuredModel] = {}
+
 
 def get_model(model_name: Optional[str] = None) -> UnstructuredModel:
     """Gets the model object by model name."""
     # TODO(alan): These cases are similar enough that we can probably do them all together with
     # importlib
+
+    global models
+
     if model_name is None:
         model_name = DEFAULT_MODEL
 
+    if model_name in models:
+        return models[model_name]
+
     if model_name in DETECTRON2_MODEL_TYPES:
         model: UnstructuredModel = UnstructuredDetectronModel()
         model.initialize(**DETECTRON2_MODEL_TYPES[model_name])
@@ -55,6 +63,7 @@ def get_model(model_name: Optional[str] = None) -> UnstructuredModel:
         model.initialize(**CHIPPER_MODEL_TYPES[model_name])
     else:
         raise UnknownModelException(f"Unknown model type: {model_name}")
+    models[model_name] = model
     return model