mirror of
https://github.com/hwchase17/langchain
synced 2024-10-29 17:07:25 +00:00
3c489be773
### Summary
Adds a post-processing method for Unstructured loaders that allows users
to optionally modify or clean extracted elements.
### Testing
```python
from langchain.document_loaders import UnstructuredFileLoader
from unstructured.cleaners.core import clean_extra_whitespace
loader = UnstructuredFileLoader(
"./example_data/layout-parser-paper.pdf",
mode="elements",
post_processors=[clean_extra_whitespace],
)
docs = loader.load()
docs[:5]
```
### Reviewrs
- @rlancemartin
- @eyurtsev
- @hwchase17
99 lines
2.6 KiB
Python
99 lines
2.6 KiB
Python
import os
|
|
from contextlib import ExitStack
|
|
from pathlib import Path
|
|
|
|
from unstructured.cleaners.core import clean_extra_whitespace
|
|
|
|
from langchain.document_loaders import (
|
|
UnstructuredAPIFileIOLoader,
|
|
UnstructuredAPIFileLoader,
|
|
UnstructuredFileLoader,
|
|
)
|
|
|
|
EXAMPLE_DOCS_DIRECTORY = str(Path(__file__).parent.parent / "examples/")
|
|
|
|
|
|
def test_unstructured_loader_with_post_processor() -> None:
|
|
file_path = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper.pdf")
|
|
loader = UnstructuredFileLoader(
|
|
file_path=file_path,
|
|
pos_processors=[clean_extra_whitespace],
|
|
strategy="fast",
|
|
mode="elements",
|
|
)
|
|
docs = loader.load()
|
|
|
|
assert len(docs) > 1
|
|
|
|
|
|
def test_unstructured_api_file_loader() -> None:
|
|
"""Test unstructured loader."""
|
|
file_path = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper.pdf")
|
|
loader = UnstructuredAPIFileLoader(
|
|
file_path=file_path,
|
|
api_key="FAKE_API_KEY",
|
|
strategy="fast",
|
|
mode="elements",
|
|
)
|
|
docs = loader.load()
|
|
|
|
assert len(docs) > 1
|
|
|
|
|
|
def test_unstructured_api_file_loader_multiple_files() -> None:
|
|
"""Test unstructured loader."""
|
|
file_paths = [
|
|
os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper.pdf"),
|
|
os.path.join(EXAMPLE_DOCS_DIRECTORY, "whatsapp_chat.txt"),
|
|
]
|
|
|
|
loader = UnstructuredAPIFileLoader(
|
|
file_path=file_paths,
|
|
api_key="FAKE_API_KEY",
|
|
strategy="fast",
|
|
mode="elements",
|
|
)
|
|
docs = loader.load()
|
|
|
|
assert len(docs) > 1
|
|
|
|
|
|
def test_unstructured_api_file_io_loader() -> None:
|
|
"""Test unstructured loader."""
|
|
file_path = os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper.pdf")
|
|
|
|
with open(file_path, "rb") as f:
|
|
loader = UnstructuredAPIFileIOLoader(
|
|
file=f,
|
|
api_key="FAKE_API_KEY",
|
|
strategy="fast",
|
|
mode="elements",
|
|
file_filename=file_path,
|
|
)
|
|
docs = loader.load()
|
|
|
|
assert len(docs) > 1
|
|
|
|
|
|
def test_unstructured_api_file_loader_io_multiple_files() -> None:
|
|
"""Test unstructured loader."""
|
|
file_paths = [
|
|
os.path.join(EXAMPLE_DOCS_DIRECTORY, "layout-parser-paper.pdf"),
|
|
os.path.join(EXAMPLE_DOCS_DIRECTORY, "whatsapp_chat.txt"),
|
|
]
|
|
|
|
with ExitStack() as stack:
|
|
files = [stack.enter_context(open(file_path, "rb")) for file_path in file_paths]
|
|
|
|
loader = UnstructuredAPIFileIOLoader(
|
|
file=files, # type: ignore
|
|
api_key="FAKE_API_KEY",
|
|
strategy="fast",
|
|
mode="elements",
|
|
file_filenames=file_paths,
|
|
)
|
|
|
|
docs = loader.load()
|
|
|
|
assert len(docs) > 1
|