chore: dl default model from huggingface (#30)

We've now seen cases where dropbox links are blacklisted, which prevent them from downloading the default detectron2 model. This PR changes the default model retrieval method to use huggingface (which already hosts the weights and config for the model). The code is also updated so getting the default model and a non-default model uses more of the same code.

Author: qued

CHANGELOG.md

@@ -1,5 +1,6 @@
-## 0.2.4-dev1
+## 0.2.4
 
+* Download default model from huggingface
 * Clarify error when trying to open file that doesn't exist as an image
 
 ## 0.2.3

test_unstructured_inference/inference/test_layout.py

@@ -8,7 +8,7 @@
 from PIL import Image
 
 import unstructured_inference.inference.layout as layout
-import unstructured_inference.models as models
+import unstructured_inference.models.base as models
 
 import unstructured_inference.models.detectron2 as detectron2
 import unstructured_inference.models.tesseract as tesseract
@@ -66,7 +66,9 @@ def detect(self, *args):
 
 
 def test_get_page_elements(monkeypatch, mock_page_layout):
-    monkeypatch.setattr(detectron2, "load_default_model", lambda: MockLayoutModel(mock_page_layout))
+    monkeypatch.setattr(
+        models, "load_model", lambda *args, **kwargs: MockLayoutModel(mock_page_layout)
+    )
     monkeypatch.setattr(detectron2, "is_detectron2_available", lambda *args: True)
 
     image = np.random.randint(12, 24, (40, 40))
@@ -88,7 +90,7 @@ def test_get_page_elements_with_ocr(monkeypatch):
     text_block = TextBlock(rectangle, text=None, type="Title")
     doc_layout = Layout([text_block])
 
-    monkeypatch.setattr(detectron2, "load_default_model", lambda: MockLayoutModel(doc_layout))
+    monkeypatch.setattr(models, "load_model", lambda *args, **kwargs: MockLayoutModel(doc_layout))
     monkeypatch.setattr(detectron2, "is_detectron2_available", lambda *args: True)
 
     image = np.random.randint(12, 24, (40, 40))
@@ -104,7 +106,9 @@ def test_read_pdf(monkeypatch, mock_page_layout):
 
     layouts = Layout([mock_page_layout, mock_page_layout])
 
-    monkeypatch.setattr(detectron2, "load_default_model", lambda: MockLayoutModel(mock_page_layout))
+    monkeypatch.setattr(
+        models, "load_model", lambda *args, **kwargs: MockLayoutModel(mock_page_layout)
+    )
     monkeypatch.setattr(detectron2, "is_detectron2_available", lambda *args: True)
 
     with patch.object(lp, "load_pdf", return_value=(layouts, images)):
@@ -138,6 +142,7 @@ def test_process_data_with_model(monkeypatch, mock_page_layout, model_name):
             "fake-binary-path",
             "fake-config-path",
             {0: "Unchecked", 1: "Checked"},
+            None,
         ),
     )
     with patch("builtins.open", mock_open(read_data=b"000000")):
@@ -168,6 +173,7 @@ def test_process_file_with_model(monkeypatch, mock_page_layout, model_name):
             "fake-binary-path",
             "fake-config-path",
             {0: "Unchecked", 1: "Checked"},
+            None,
         ),
     )
     filename = ""

test_unstructured_inference/models/test_detectron2.py

@@ -2,6 +2,7 @@
 from unittest.mock import patch
 
 import unstructured_inference.models.detectron2 as detectron2
+import unstructured_inference.models.base as models
 
 
 class MockDetectron2LayoutModel:
@@ -14,15 +15,15 @@ def test_load_default_model(monkeypatch):
     monkeypatch.setattr(detectron2, "Detectron2LayoutModel", MockDetectron2LayoutModel)
 
     with patch.object(detectron2, "is_detectron2_available", return_value=True):
-        model = detectron2.load_default_model()
+        model = models.get_model()
 
     assert isinstance(model, MockDetectron2LayoutModel)
 
 
 def test_load_default_model_raises_when_not_available():
     with patch.object(detectron2, "is_detectron2_available", return_value=False):
         with pytest.raises(ImportError):
-            detectron2.load_default_model()
+            models.get_model()
 
 
 @pytest.mark.parametrize("config_path, model_path", [("asdf", "diufs"), ("dfaw", "hfhfhfh")])

test_unstructured_inference/models/test_model.py

@@ -1,6 +1,6 @@
 import pytest
 
-from unstructured_inference import models
+import unstructured_inference.models.base as models
 
 
 class MockModel:
@@ -18,6 +18,7 @@ def test_get_model(monkeypatch):
             "fake-binary-path",
             "fake-config-path",
             {0: "Unchecked", 1: "Checked"},
+            None,
         ),
     )
     assert isinstance(models.get_model("checkbox"), MockModel)

test_unstructured_inference/test_api.py

@@ -4,7 +4,7 @@
 from fastapi.testclient import TestClient
 
 from unstructured_inference.api import app
-from unstructured_inference import models
+from unstructured_inference.models import base as models
 from unstructured_inference.inference.layout import DocumentLayout
 import unstructured_inference.models.detectron2 as detectron2
 
@@ -15,8 +15,18 @@ def __init__(self, *args, **kwargs):
         self.kwargs = kwargs
 
 
-@pytest.mark.parametrize("filetype, ext", [("pdf", "pdf"), ("image", "png")])
-def test_layout_parsing_api(monkeypatch, filetype, ext):
+@pytest.mark.parametrize(
+    "filetype, ext, data, response_code",
+    [
+        ("pdf", "pdf", None, 200),
+        ("pdf", "pdf", {"model": "checkbox"}, 200),
+        ("pdf", "pdf", {"model": "fake_model"}, 422),
+        ("image", "png", None, 200),
+        ("image", "png", {"model": "checkbox"}, 200),
+        ("image", "png", {"model": "fake_model"}, 422),
+    ],
+)
+def test_layout_parsing_api(monkeypatch, filetype, ext, data, response_code):
     monkeypatch.setattr(models, "load_model", lambda *args, **kwargs: MockModel(*args, **kwargs))
     monkeypatch.setattr(models, "hf_hub_download", lambda *args, **kwargs: "fake-path")
     monkeypatch.setattr(detectron2, "is_detectron2_available", lambda *args: True)
@@ -30,22 +40,10 @@ def test_layout_parsing_api(monkeypatch, filetype, ext):
     filename = os.path.join("sample-docs", f"loremipsum.{ext}")
 
     client = TestClient(app)
-    response = client.post(f"/layout/{filetype}", files={"file": (filename, open(filename, "rb"))})
-    assert response.status_code == 200
-
-    response = client.post(
-        f"/layout/{filetype}",
-        files={"file": (filename, open(filename, "rb"))},
-        data={"model": "checkbox"},
-    )
-    assert response.status_code == 200
-
     response = client.post(
-        f"/layout/{filetype}",
-        files={"file": (filename, open(filename, "rb"))},
-        data={"model": "fake_model"},
+        f"/layout/{filetype}", files={"file": (filename, open(filename, "rb"))}, data=data
     )
-    assert response.status_code == 422
+    assert response.status_code == response_code
 
 
 def test_bad_route_404():

unstructured_inference/__version__.py

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

unstructured_inference/api.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, File, status, Request, UploadFile, Form, HTTPException
 from unstructured_inference.inference.layout import process_data_with_model
-from unstructured_inference.models import UnknownModelException
+from unstructured_inference.models.base import UnknownModelException
 from typing import List
 
 app = FastAPI()

unstructured_inference/inference/layout.py

@@ -12,8 +12,7 @@
 
 from unstructured_inference.logger import logger
 import unstructured_inference.models.tesseract as tesseract
-import unstructured_inference.models.detectron2 as detectron2
-from unstructured_inference.models import get_model
+from unstructured_inference.models.base import get_model
 
 
 @dataclass
@@ -113,7 +112,7 @@ def get_elements(self, inplace=True) -> Optional[List[LayoutElement]]:
         """Uses a layoutparser model to detect the elements on the page."""
         logger.info("Detecting page elements ...")
         if self.model is None:
-            self.model = detectron2.load_default_model()
+            self.model = get_model()
 
         elements = list()
         # NOTE(mrobinson) - We'll want make this model inference step some kind of
@@ -183,7 +182,7 @@ def process_file_with_model(
 ) -> DocumentLayout:
     """Processes pdf file with name filename into a DocumentLayout by using a model identified by
     model_name."""
-    model = None if model_name is None else get_model(model_name)
+    model = get_model(model_name)
     layout = (
         DocumentLayout.from_image_file(filename, model=model)
         if is_image

unstructured_inference/models/__init__.py

@@ -1,34 +0,0 @@
-from typing import Tuple, Dict
-from huggingface_hub import hf_hub_download
-
-from unstructured_inference.models.detectron2 import load_model, Detectron2LayoutModel
-
-
-def get_model(model: str) -> Detectron2LayoutModel:
-    """Gets the model object by model name."""
-    model_path, config_path, label_map = _get_model_loading_info(model)
-    detector = load_model(config_path=config_path, model_path=model_path, label_map=label_map)
-
-    return detector
-
-
-def _get_model_loading_info(model: str) -> Tuple[str, str, Dict[int, str]]:
-    """Gets local model binary and config locations and label map, downloading if necessary."""
-    # TODO(alan): Find the right way to map model name to retrieval. It seems off that testing
-    # needs to mock hf_hub_download.
-    if model == "checkbox":
-        repo_id = "unstructuredio/oer-checkbox"
-        binary_fn = "detectron2_finetuned_oer_checkbox.pth"
-        config_fn = "detectron2_oer_checkbox.json"
-        model_path = hf_hub_download(repo_id, binary_fn)
-        config_path = hf_hub_download(repo_id, config_fn)
-        label_map = {0: "Unchecked", 1: "Checked"}
-    else:
-        raise UnknownModelException(f"Unknown model type: {model}")
-    return model_path, config_path, label_map
-
-
-class UnknownModelException(Exception):
-    """Exception for the case where a model is called for with an unrecognized identifier."""
-
-    pass

unstructured_inference/models/base.py

@@ -0,0 +1,54 @@
+from typing import Tuple, Dict, Optional, List, Any
+from huggingface_hub import hf_hub_download
+
+from unstructured_inference.models.detectron2 import (
+    load_model,
+    Detectron2LayoutModel,
+    DEFAULT_LABEL_MAP,
+    DEFAULT_EXTRA_CONFIG,
+)
+
+
+def get_model(model: Optional[str] = None) -> Detectron2LayoutModel:
+    """Gets the model object by model name."""
+    model_path, config_path, label_map, extra_config = _get_model_loading_info(model)
+    detector = load_model(
+        config_path=config_path,
+        model_path=model_path,
+        label_map=label_map,
+        extra_config=extra_config,
+    )
+    return detector
+
+
+def _get_model_loading_info(
+    model: Optional[str],
+) -> Tuple[str, str, Dict[int, str], Optional[List[Any]]]:
+    """Gets local model binary and config locations and label map, downloading if necessary."""
+    # TODO(alan): Find the right way to map model name to retrieval. It seems off that testing
+    # needs to mock hf_hub_download.
+    if model is None:
+        repo_id = "layoutparser/detectron2"
+        binary_fn = "PubLayNet/faster_rcnn_R_50_FPN_3x/model_final.pth"
+        config_fn = "PubLayNet/faster_rcnn_R_50_FPN_3x/config.yml"
+        model_path = hf_hub_download(repo_id, binary_fn)
+        config_path = hf_hub_download(repo_id, config_fn)
+        label_map = DEFAULT_LABEL_MAP
+        extra_config = DEFAULT_EXTRA_CONFIG
+    elif model == "checkbox":
+        repo_id = "unstructuredio/oer-checkbox"
+        binary_fn = "detectron2_finetuned_oer_checkbox.pth"
+        config_fn = "detectron2_oer_checkbox.json"
+        model_path = hf_hub_download(repo_id, binary_fn)
+        config_path = hf_hub_download(repo_id, config_fn)
+        label_map = {0: "Unchecked", 1: "Checked"}
+        extra_config = None
+    else:
+        raise UnknownModelException(f"Unknown model type: {model}")
+    return model_path, config_path, label_map, extra_config
+
+
+class UnknownModelException(Exception):
+    """Exception for the case where a model is called for with an unrecognized identifier."""
+
+    pass