Add optional param for model name when partitioning pdfs (#890)

* Add optional param for model name when partitioning pdfs

* Pull in latest inference changes

* Fix linting

Author: rbiseck3

CHANGELOG.md

@@ -1,7 +1,10 @@
-## 0.8.0-dev2
+## 0.8.0
 
 ### Enhancements
 
+* Allow model used for hi res pdf partition strategy to be chosen when called.
+* Updated inference package
+
 ### Features
 
 * Add metadata_filename parameter across all partition functions

requirements/local-inference.in

@@ -1,3 +1,3 @@
 -c constraints.in
 -c base.txt
-unstructured-inference==0.5.4
+unstructured-inference==0.5.5

requirements/local-inference.txt

@@ -211,7 +211,7 @@ typing-extensions==4.7.0
     #   huggingface-hub
     #   iopath
     #   torch
-unstructured-inference==0.5.4
+unstructured-inference==0.5.5
     # via -r requirements/local-inference.in
 urllib3==1.26.16
     # via

test_unstructured/partition/test_pdf.py

@@ -129,7 +129,7 @@ def test_partition_pdf_with_spooled_file(
 
 
 @mock.patch.dict(os.environ, {"UNSTRUCTURED_HI_RES_MODEL_NAME": "checkbox"})
-def test_partition_pdf_with_model_name(
+def test_partition_pdf_with_model_name_env_var(
     monkeypatch,
     filename="example-docs/layout-parser-paper-fast.pdf",
 ):
@@ -152,6 +152,26 @@ def test_partition_pdf_with_model_name(
         )
 
 
+def test_partition_pdf_with_model_name(
+    monkeypatch,
+    filename="example-docs/layout-parser-paper-fast.pdf",
+):
+    monkeypatch.setattr(
+        strategies,
+        "is_pdf_text_extractable",
+        lambda *args, **kwargs: True,
+    )
+    with mock.patch.object(layout, "process_file_with_model", mock.MagicMock()) as mock_process:
+        pdf.partition_pdf(filename=filename, strategy="hi_res", model_name="checkbox")
+        mock_process.assert_called_once_with(
+            filename,
+            is_image=False,
+            ocr_languages="eng",
+            extract_tables=False,
+            model_name="checkbox",
+        )
+
+
 def test_partition_pdf_with_auto_strategy(
     filename="example-docs/layout-parser-paper-fast.pdf",
 ):

unstructured/__version__.py

@@ -1 +1 @@
-__version__ = "0.8.0-dev2"  # pragma: no cover
+__version__ = "0.8.0"  # pragma: no cover

unstructured/ingest/doc_processor/generalized.py

@@ -1,4 +1,4 @@
-"""Process aribritrary files with the Unstructured library"""
+"""Process arbitrary files with the Unstructured library"""
 
 import os
 from typing import Any, Dict, List, Optional
@@ -12,6 +12,13 @@
 def initialize():
     """Download default model or model specified by UNSTRUCTURED_HI_RES_MODEL_NAME environment
     variable (avoids subprocesses all doing the same)"""
+
+    # If more than one model will be supported and left up to user selection
+    supported_model = os.environ.get("UNSTRUCTURED_HI_RES_SUPPORTED_MODEL", "")
+    if supported_model:
+        for model_name in supported_model.split(","):
+            get_model(model_name=model_name)
+
     get_model(os.environ.get("UNSTRUCTURED_HI_RES_MODEL_NAME"))
 
 

unstructured/partition/pdf.py

@@ -90,6 +90,7 @@ def partition_pdf(
         infer_table_structure=infer_table_structure,
         ocr_languages=ocr_languages,
         max_partition=max_partition,
+        **kwargs,
     )
 
 
@@ -102,6 +103,7 @@ def partition_pdf_or_image(
     infer_table_structure: bool = False,
     ocr_languages: str = "eng",
     max_partition: Optional[int] = 1500,
+    **kwargs,
 ) -> List[Element]:
     """Parses a pdf or image document into a list of interpreted elements."""
     # TODO(alan): Extract information about the filetype to be processed from the template
@@ -128,6 +130,7 @@ def partition_pdf_or_image(
                 infer_table_structure=infer_table_structure,
                 include_page_breaks=include_page_breaks,
                 ocr_languages=ocr_languages,
+                **kwargs,
             )
 
     elif strategy == "fast":
@@ -160,6 +163,8 @@ def _partition_pdf_or_image_local(
     infer_table_structure: bool = False,
     include_page_breaks: bool = False,
     ocr_languages: str = "eng",
+    model_name: Optional[str] = None,
+    **kwargs,
 ) -> List[Element]:
     """Partition using package installed locally."""
     try:
@@ -182,7 +187,7 @@ def _partition_pdf_or_image_local(
             "running make install-local-inference from the root directory of the repository.",
         ) from e
 
-    model_name = os.environ.get("UNSTRUCTURED_HI_RES_MODEL_NAME")
+    model_name = model_name if model_name else os.environ.get("UNSTRUCTURED_HI_RES_MODEL_NAME")
     if file is None:
         layout = process_file_with_model(
             filename,