feat: add --fields-include to unstructured-ingest

CHANGELOG.md

@@ -1,9 +1,11 @@
-## 0.5.5-dev0
+## 0.5.5-dev2
 
 ### Enhancements
 
 ### Features
 
+* Add `--fields-include` parameter to `unstructured-ingest`
+* Add `--metadata-include` and `--metadata-exclude` parameters to `unstructured-ingest`
 * Add `clean_non_ascii_chars` to remove non-ascii characters from unicode string
 
 ### Fixes

test_unstructured_ingest/test-ingest-azure.sh

@@ -4,6 +4,7 @@ SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
 cd "$SCRIPT_DIR"/.. || exit 1
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
     --remote-url abfs://container1/ \
     --azure-account-name azureunstructured1 \
     --structured-output-dir azure-ingest-output \

test_unstructured_ingest/test-ingest-biomed-api.sh

@@ -10,6 +10,7 @@ if [[ "$(find test_unstructured_ingest/expected-structured-output/biomed-ingest-
 fi
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+   --metadata-exclude filename \
    --biomed-api-from "2019-01-02" \
    --biomed-api-until "2019-01-02+00:03:10" \
    --structured-output-dir biomed-ingest-output-api  \

test_unstructured_ingest/test-ingest-biomed-path.sh

@@ -10,6 +10,7 @@ if [[ "$(find test_unstructured_ingest/expected-structured-output/biomed-ingest-
 fi
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
     --biomed-path "oa_pdf/07/07/sbaa031.073.PMC7234218.pdf" \
     --structured-output-dir biomed-ingest-output-path \
     --num-processes 2 \

test_unstructured_ingest/test-ingest-github.sh

@@ -12,7 +12,12 @@ if [[ "$CI" == "true" ]]; then
 fi
 
 
-PYTHONPATH=. ./unstructured/ingest/main.py --github-url dcneiner/Downloadify --git-file-glob '*.html,*.txt' --structured-output-dir github-downloadify-output --verbose
+PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
+    --github-url dcneiner/Downloadify \
+    --git-file-glob '*.html,*.txt' \
+    --structured-output-dir github-downloadify-output \
+    --verbose
 
 if ! diff -ru test_unstructured_ingest/expected-structured-output/github-downloadify github-downloadify-output ; then
    echo

test_unstructured_ingest/test-ingest-gitlab.sh

@@ -4,6 +4,7 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 cd "$SCRIPT_DIR"/.. || exit 1
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
     --gitlab-url https://gitlab.com/gitlab-com/content-sites/docsy-gitlab \
     --git-file-glob '*.md,*.txt' \
     --structured-output-dir gitlab-ingest-output \

test_unstructured_ingest/test-ingest-s3.sh

@@ -9,7 +9,11 @@ if [[ "$(find test_unstructured_ingest/expected-structured-output/s3-small-batch
     exit 1
 fi
 
-PYTHONPATH=. ./unstructured/ingest/main.py --s3-url s3://utic-dev-tech-fixtures/small-pdf-set/ --s3-anonymous --structured-output-dir s3-small-batch-output
+PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
+    --s3-url s3://utic-dev-tech-fixtures/small-pdf-set/ \
+    --s3-anonymous \
+    --structured-output-dir s3-small-batch-output \
 
 if ! diff -ru test_unstructured_ingest/expected-structured-output/s3-small-batch s3-small-batch-output ; then
     echo

test_unstructured_ingest/test-ingest-wikipedia.sh

@@ -4,6 +4,7 @@ SCRIPT_DIR=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 cd "$SCRIPT_DIR"/.. || exit 1
 
 PYTHONPATH=. ./unstructured/ingest/main.py \
+    --metadata-exclude filename \
     --wikipedia-page-title "Open Source Software" \
     --structured-output-dir wikipedia-ingest-output \
     --num-processes 2 \

test_unstructured_ingest/test_interfaces.py

@@ -0,0 +1,78 @@
+import os
+import pathlib
+
+import pytest
+
+from unstructured.ingest.connector.git import GitIngestDoc, SimpleGitConfig
+
+DIRECTORY = pathlib.Path(__file__).parent.resolve()
+EXAMPLE_DOCS_DIRECTORY = os.path.join(DIRECTORY, "..", "example-docs")
+
+test_files = [
+    "layout-parser-paper-fast.jpg",
+    "layout-parser-paper-fast.pdf",
+]
+
+
+@pytest.mark.parametrize("filename", test_files)
+def test_process_file_include_filename(filename: str):
+    ingest_doc = GitIngestDoc(
+        path=filename,
+        config=SimpleGitConfig(
+          download_dir=EXAMPLE_DOCS_DIRECTORY,
+          metadata_include="filename",
+        ),
+    )
+    isd_elems = ingest_doc.process_file()
+
+    for elem in isd_elems:
+        for k in elem["metadata"]:
+            assert k == "filename"
+
+
+@pytest.mark.parametrize("filename", test_files)
+def test_process_file_include_filename_pagenum(filename: str):
+    ingest_doc = GitIngestDoc(
+        path=filename,
+        config=SimpleGitConfig(
+          download_dir=EXAMPLE_DOCS_DIRECTORY,
+          metadata_include="filename,page_number",
+        ),
+    )
+    isd_elems = ingest_doc.process_file()
+
+    for elem in isd_elems:
+        for k in elem["metadata"]:
+            assert k in ["filename", "page_number"]
+
+
+@pytest.mark.parametrize("filename", test_files)
+def test_process_file_exclude_filename(filename: str):
+    ingest_doc = GitIngestDoc(
+        path=filename,
+        config=SimpleGitConfig(
+          download_dir=EXAMPLE_DOCS_DIRECTORY,
+          metadata_exclude="filename",
+        ),
+    )
+    isd_elems = ingest_doc.process_file()
+
+    for elem in isd_elems:
+        for k in elem["metadata"]:
+            assert k != "filename"
+
+
+@pytest.mark.parametrize("filename", test_files)
+def test_process_file_exclude_filename_pagenum(filename: str):
+    ingest_doc = GitIngestDoc(
+        path=filename,
+        config=SimpleGitConfig(
+          download_dir=EXAMPLE_DOCS_DIRECTORY,
+          metadata_exclude="filename,page_number",
+        ),
+    )
+    isd_elems = ingest_doc.process_file()
+
+    for elem in isd_elems:
+        for k in elem["metadata"]:
+            assert k not in ["filename", "page_number"]

unstructured/__version__.py

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

unstructured/ingest/connector/biomed.py

@@ -48,6 +48,9 @@ class SimpleBiomedConfig(BaseConnectorConfig):
     output_dir: str
     re_download: bool = False
     preserve_downloads: bool = False
+    metadata_include: str = ""
+    metadata_exclude: str = ""
+    fields_include: str = ""
 
     def _validate_date_args(self, date):
         date_formats = ["%Y-%m-%d", "%Y-%m-%d+%H:%M:%S"]

unstructured/ingest/connector/fsspec.py

@@ -29,6 +29,9 @@ class SimpleFsspecConfig(BaseConnectorConfig):
     output_dir: str
     preserve_downloads: bool = False
     re_download: bool = False
+    metadata_include: str = ""
+    metadata_exclude: str = ""
+    fields_include: str = ""
 
     # fsspec specific options
     access_kwargs: dict = field(default_factory=dict)

unstructured/ingest/connector/git.py

@@ -26,6 +26,9 @@ class SimpleGitConfig(BaseConnectorConfig):
     output_dir: str
     preserve_downloads: bool = False
     re_download: bool = False
+    metadata_include: str = ""
+    metadata_exclude: str = ""
+    fields_include: str = ""
 
     repo_path: str = field(init=False, repr=False)
 

unstructured/ingest/connector/google_drive.py

@@ -77,6 +77,9 @@ class SimpleGoogleDriveConfig(BaseConnectorConfig):
     output_dir: str
     re_download: bool = False
     preserve_downloads: bool = False
+    metadata_include: str = ""
+    metadata_exclude: str = ""
+    fields_include: str = ""
 
     recursive: bool = False
 

unstructured/ingest/connector/reddit.py

@@ -31,6 +31,9 @@ class SimpleRedditConfig(BaseConnectorConfig):
     output_dir: str
     preserve_downloads: bool = False
     re_download: bool = False
+    metadata_include: str = ""
+    metadata_exclude: str = ""
+    fields_include: str = ""
 
     def __post_init__(self):
         if self.num_posts <= 0:

unstructured/ingest/connector/wikipedia.py

@@ -26,6 +26,9 @@ class SimpleWikipediaConfig(BaseConnectorConfig):
     output_dir: str
     preserve_downloads: bool = False
     re_download: bool = False
+    metadata_include: str = ""
+    metadata_exclude: str = ""
+    fields_include: str = ""
 
 
 @dataclass

unstructured/ingest/interfaces.py

@@ -47,6 +47,9 @@ class BaseConnectorConfig(ABC):
     # where to write structured data outputs
     output_dir: str
     re_download: bool = False
+    metadata_include: str = ""
+    metadata_exclude: str = ""
+    fields_include: str = ""
 
 
 class BaseIngestDoc(ABC):
@@ -58,6 +61,8 @@ class BaseIngestDoc(ABC):
     Crucially, it is not responsible for the actual processing of the raw document.
     """
 
+    config: BaseConnectorConfig
+
     @property
     @abstractmethod
     def filename(self):
@@ -94,8 +99,21 @@ def process_file(self):
         self.isd_elems_no_filename = []
         for elem in isd_elems:
             # type: ignore
-            elem["metadata"].pop("filename", None)  # type: ignore[attr-defined]
-            elem.pop("coordinates")  # type: ignore[attr-defined]
+            if self.config.metadata_exclude:
+                ex_list = self.config.metadata_exclude.split(",")
+                for ex in ex_list:
+                    elem["metadata"].pop(ex, None)  # type: ignore[attr-defined]
+            elif self.config.metadata_include:
+                in_list = self.config.metadata_include.split(",")
+                for k in elem["metadata"]:
+                    if k not in in_list:
+                        elem["metadata"].pop(k, None)  # type: ignore[attr-defined]
+
+            in_list = self.config.fields_include.split(",")
+            for k in elem:
+                if k not in in_list:
+                    elem.pop(k, None)  # type: ignore[attr-defined]
+
             self.isd_elems_no_filename.append(elem)
 
         return self.isd_elems_no_filename