feat: allow users to pass OCR language into partition (#509)

* pip-compile new reqs

* bump inference version

* add language to pdf and image calls

* tests for passing in language

* version bump and changelog

* update docs

* pass ocr_languages in auto

* updated test fixtures

* typo in doc string

Author: MthwRobinson

CHANGELOG.md

@@ -1,9 +1,13 @@
-## 0.5.14-dev0
+## 0.5.14-dev1
 
 ### Enhancements
 
 * Adds an `ssl_verify` kwarg to `partition` and `partition_html` to enable turning off
   SSL verification for HTTP requests. SSL verification is on by default.
+* Allows users to pass in ocr language to `partition_pdf` and `partition_image` through
+  the `ocr_language` kwarg. `ocr_language` corresponds to the code for the language pack
+  in Tesseract. You will need to install the relevant Tesseract language pack to use a
+  given language.
 
 ### Features
 

docs/source/bricks.rst

@@ -283,6 +283,10 @@ The ``strategy`` kwarg controls the method that will be used to process the PDF.
 will identify the layout of the document using ``detectron2``. The ``"fast"`` strategy will extract the
 text using ``pdfminer`` and process the raw text with ``partition_text``. If ``detectron2`` is not available,
 and the ``"hi_res"`` strategy is set, ``partition_pdf`` will fallback to the ``"fast"`` strategy.
+You can also specify what languages to use for OCR with the ``ocr_languages`` kwarg. For example,
+use ``ocr_languages="eng+deu"`` to use the English and German language packs. See the
+`Tesseract documentation <https://github.com/tesseract-ocr/tessdata>`_ for a full list of languages and
+install instructions. OCR is only applied if the text is not already available in the PDF document.
 
 Examples:
 
@@ -293,13 +297,22 @@ Examples:
   # Returns a List[Element] present in the pages of the parsed pdf document
   elements = partition_pdf("example-docs/layout-parser-paper-fast.pdf")
 
+  # Applies the English and Swedish language pack for ocr. OCR is only applied
+  # if the text is not available in the PDF.
+  elements = partition_pdf("example-docs/layout-parser-paper-fast.pdf", ocr_languages="eng+swe")
+
 
 ``partition_image``
 ---------------------
 
 The ``partition_image`` function has the same API as ``partition_pdf``, which is document above.
 The only difference is that ``partition_image`` does not need to convert a PDF to an image
 prior to processing. The ``partition_image`` function supports ``.png`` and ``.jpg`` files.
+You can also specify what languages to use for OCR with the ``ocr_languages`` kwarg. For example,
+use ``ocr_languages="eng+deu"`` to use the English and German language packs. See the
+`Tesseract documentation <https://github.com/tesseract-ocr/tessdata>`_ for a full list of languages and
+install instructions.
+
 
 Examples:
 
@@ -310,6 +323,9 @@ Examples:
   # Returns a List[Element] present in the pages of the parsed image document
   elements = partition_image("example-docs/layout-parser-paper-fast.jpg")
 
+  # Applies the English and Swedish language pack for ocr
+  elements = partition_image("example-docs/layout-parser-paper-fast.jpg", ocr_languages="eng+swe")
+
 
 
 ``partition_email``

requirements/local-inference.txt

@@ -268,7 +268,7 @@ typing-extensions==4.5.0
     #   rich
     #   starlette
     #   torch
-unstructured-inference==0.3.2
+unstructured-inference==0.4.1
     # via unstructured (setup.py)
 urllib3==1.26.15
     # via requests

setup.py

@@ -76,7 +76,7 @@
             "transformers",
         ],
         "local-inference": [
-            "unstructured-inference==0.3.2",
+            "unstructured-inference>=0.4.1",
         ],
         "s3": ["s3fs", "fsspec"],
         "azure": ["adlfs", "fsspec"],

test_unstructured/partition/test_auto.py

@@ -282,6 +282,7 @@ def test_auto_partition_pdf_with_fast_strategy():
         include_page_breaks=False,
         encoding="utf-8",
         strategy="fast",
+        ocr_languages="eng",
     )
 
 

test_unstructured/partition/test_image.py

@@ -2,6 +2,7 @@
 
 import pytest
 import requests
+from pytesseract import TesseractError
 from unstructured_inference.inference import layout
 
 from unstructured.partition import image, pdf
@@ -157,3 +158,23 @@ def test_partition_image(url, api_called, local_called):
         image.partition_image(filename="fake.pdf", url=url)
         assert pdf._partition_via_api.called == api_called
         assert pdf._partition_pdf_or_image_local.called == local_called
+
+
+def test_partition_image_with_language_passed(filename="example-docs/example.jpg"):
+    with mock.patch.object(layout, "process_file_with_model", mock.MagicMock()) as mock_partition:
+        image.partition_image(filename=filename, ocr_languages="eng+swe")
+
+    assert mock_partition.call_args.kwargs.get("ocr_languages") == "eng+swe"
+
+
+def test_partition_image_from_file_with_language_passed(filename="example-docs/example.jpg"):
+    with mock.patch.object(layout, "process_data_with_model", mock.MagicMock()) as mock_partition:
+        with open(filename, "rb") as f:
+            image.partition_image(file=f, ocr_languages="eng+swe")
+
+    assert mock_partition.call_args.kwargs.get("ocr_languages") == "eng+swe"
+
+
+def test_partition_image_raises_with_invalid_language(filename="example-docs/example.jpg"):
+    with pytest.raises(TesseractError):
+        image.partition_image(filename=filename, ocr_languages="fakeroo")

test_unstructured_ingest/expected-structured-output/s3-small-batch/small-pdf-set/2023-Jan-economic-outlook.pdf.json

@@ -224,8 +224,8 @@
     }
   },
   {
-    "element_id": "0953470500eb215048fd49263b8829a4",
-    "text": "Forces Shaping the Outlook",
+    "element_id": "e3b0c44298fc1c149afbf4c8996fb924",
+    "text": "",
     "type": "Title",
     "metadata": {
       "page_number": 2

test_unstructured_ingest/expected-structured-output/s3-small-batch/small-pdf-set/Silent-Giant-(1).pdf.json

@@ -1,7 +1,7 @@
 [
   {
-    "element_id": "ea216492b46010685b4a036fe66de211",
-    "text": "WORLD NUCLEARASSOCIATION",
+    "element_id": "e3b0c44298fc1c149afbf4c8996fb924",
+    "text": "",
     "type": "Title",
     "metadata": {
       "page_number": 1
@@ -24,8 +24,8 @@
     }
   },
   {
-    "element_id": "53d548aa01fc3eb72da15a5be7f235e2",
-    "text": "Executive Summary",
+    "element_id": "8e76a94ac8320d515375e625bef18292",
+    "text": "Summary",
     "type": "Title",
     "metadata": {
       "page_number": 3
@@ -248,8 +248,8 @@
     }
   },
   {
-    "element_id": "3655eec20e80973efc46cc09db7a04ba",
-    "text": "Moving to a sustainable future",
+    "element_id": "ff4a9b34d6cdebbc9b8afbf9767f6e1c",
+    "text": "to a sustainable future",
     "type": "Title",
     "metadata": {
       "page_number": 6

test_unstructured_ingest/expected-structured-output/s3-small-batch/small-pdf-set/recalibrating-risk-report.pdf.json

@@ -1,7 +1,7 @@
 [
   {
-    "element_id": "ea216492b46010685b4a036fe66de211",
-    "text": "WORLD NUCLEARASSOCIATION",
+    "element_id": "e3b0c44298fc1c149afbf4c8996fb924",
+    "text": "",
     "type": "Title",
     "metadata": {
       "page_number": 1
@@ -24,8 +24,8 @@
     }
   },
   {
-    "element_id": "53d548aa01fc3eb72da15a5be7f235e2",
-    "text": "Executive Summary",
+    "element_id": "8e76a94ac8320d515375e625bef18292",
+    "text": "Summary",
     "type": "Title",
     "metadata": {
       "page_number": 3

unstructured/__version__.py

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