First pass

Author: awalker4

prepline_general/api/general.py

@@ -166,7 +166,13 @@ def partition_pdf_splits(
     pages_per_pdf = int(os.environ.get("UNSTRUCTURED_PARALLEL_MODE_SPLIT_SIZE", 1))
 
     # If it's small enough, just process locally
+    # (Some kwargs need to be renamed for local partition)
     if len(pdf_pages) <= pages_per_pdf:
+        # Get the test to fail first
+        # if "hi_res_model_name" in partition_kwargs:
+        #     partition_kwargs['model_name'] = partition_kwargs['hi_res_model_name']
+        #     del partition_kwargs['hi_res_model_name']
+
         return partition(
             file=file, file_filename=file_filename, content_type=content_type, **partition_kwargs
         )
@@ -317,6 +323,9 @@ def pipeline_api(
             )
         )
 
+        # Be careful of naming differences in api params vs partition params!
+        # These kwargs are going back into the api
+        # If needed, rename them in partition_pdf_splits
         if file_content_type == "application/pdf" and pdf_parallel_mode_enabled:
             elements = partition_pdf_splits(
                 request,
@@ -328,7 +337,7 @@ def pipeline_api(
                 # partition_kwargs
                 encoding=encoding,
                 include_page_breaks=include_page_breaks,
-                model_name=hi_res_model_name,
+                hi_res_model_name=hi_res_model_name,
                 ocr_languages=ocr_languages,
                 pdf_infer_table_structure=pdf_infer_table_structure,
                 skip_infer_table_types=skip_infer_table_types,

scripts/parallel-mode-test.sh

@@ -15,9 +15,10 @@ declare -a curl_params=(
     "-F files=@sample-docs/layout-parser-paper.pdf -F 'strategy=fast'"
     "-F files=@sample-docs/layout-parser-paper.pdf -F 'strategy=auto"
     "-F files=@sample-docs/layout-parser-paper.pdf -F 'strategy=hi_res'"
-    "-F files=@sample-docs/layout-parser-paper.pdf -F 'coordinates=true' -F 'strategy=fast'"
-    "-F files=@sample-docs/layout-parser-paper.pdf -F 'coordinates=true' -F 'strategy=fast' -F 'encoding=utf-8'"
-    "-F files=@sample-docs/layout-parser-paper.pdf -F 'coordinates=true' -F 'strategy=fast' -F 'include_page_breaks=true'"
+    "-F files=@sample-docs/layout-parser-paper.pdf -F 'coordinates=true'"
+    "-F files=@sample-docs/layout-parser-paper.pdf -F 'encoding=utf-8'"
+    "-F files=@sample-docs/layout-parser-paper.pdf -F 'include_page_breaks=true'"
+    "-F files=@sample-docs/layout-parser-paper.pdf -F 'hi_res_model_name=yolox'"
 )
 
 for params in "${curl_params[@]}"

test_general/api/test_app.py

@@ -7,8 +7,12 @@
 import pandas as pd
 from fastapi.testclient import TestClient
 from unstructured_api_tools.pipelines.api_conventions import get_pipeline_path
+from unittest.mock import patch, Mock
+
+import unstructured
 
 from prepline_general.api.app import app
+from prepline_general.api import general
 import tempfile
 
 MAIN_API_ROUTE = get_pipeline_path("general")
@@ -397,6 +401,92 @@ def json(self):
         return self.body
 
 
+@pytest.mark.only
+def test_parallel_mode_params(monkeypatch):
+    """
+    Verify that parallel mode passes all params to local partition.
+    If you add something to partition_kwargs, you need to explicitly test it here.
+    """
+    monkeypatch.setenv("UNSTRUCTURED_PARALLEL_MODE_ENABLED", "true")
+    monkeypatch.setenv("UNSTRUCTURED_PARALLEL_MODE_THREADS", "1")
+    monkeypatch.setenv("UNSTRUCTURED_PARALLEL_MODE_URL", "unused")
+
+    # Make this really big so we just call partition
+    monkeypatch.setenv("UNSTRUCTURED_PARALLEL_MODE_SPLIT_SIZE", "500")
+
+    # Verify we can pass a non-default value for everything
+    m_encoding = "foo"
+    m_hi_res_model_name = "the_big_chipper"
+    m_strategy = "highest_res"
+
+    m_ocr_languages = ["all", "of", "them"]
+    m_skip_infer_table_types = []
+
+    m_pdf_infer_table_structure = True
+    m_include_page_breaks = True
+    m_xml_keep_tags = True
+
+    # use mock to assert called with
+    def validate_local_params(
+            file,
+            file_filename,
+            content_type,
+            encoding,
+            include_page_breaks,
+            model_name,
+            ocr_languages,
+            pdf_infer_table_structure,
+            skip_infer_table_types,
+            strategy,
+            xml_keep_tags,
+    ):
+        logger.warn("validating")
+        assert encoding == m_encoding
+        assert include_page_breaks == m_include_page_breaks
+        assert model_name == m_hi_res_model_name
+        assert ocr_languages == m_ocr_languages
+        assert pdf_infer_table_structure == m_pdf_infer_table_structure
+        assert skip_infer_table_types == m_skip_infer_table_types
+        assert strategy == m_strategy
+        assert xml_keep_tags == m_xml_keep_tags
+
+        return []
+
+    mock_partition = Mock()
+
+    monkeypatch.setattr(
+        general,
+        "partition",
+        mock_partition,
+    )
+
+    client = TestClient(app)
+    test_file = Path("sample-docs") / "layout-parser-paper.pdf"
+
+    # with patch.object(general, "partition") as mock_partition:
+    # with patch.object(unstructured.partition.auto, "partition") as mock_partition:
+    response = client.post(
+        MAIN_API_ROUTE,
+        files=[("files", (str(test_file), open(test_file, "rb"), "application/pdf"))],
+        data={
+            "encoding": m_encoding,
+            "include_page_breaks": m_include_page_breaks,
+            "hi_res_model_name": m_hi_res_model_name,
+            "ocr_languages": m_ocr_languages,
+            "pdf_infer_table_structure": m_pdf_infer_table_structure,
+            "skip_infer_table_types": m_skip_infer_table_types,
+            "strategy": m_strategy,
+            # "xml_keep_tags": m_xml_keep_tags,
+        }
+    )
+
+    mock_partition.assert_called_once_with(
+        encoding="foo"
+    )
+
+    assert response.status_code == 200
+
+
 def test_parallel_mode_returns_errors(monkeypatch):
     """
     If we get an error sending a page to the api, bubble it up