修改为东南天坐标系
This commit is contained in:
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,168 @@
|
||||
"""Test suite for the `unstructured.chunking.basic` module.
|
||||
|
||||
That module implements the baseline chunking strategy. The baseline strategy has all behaviors
|
||||
shared by all chunking strategies and no extra rules like perserve section or page boundaries.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import FixtureRequest, Mock, function_mock
|
||||
from unstructured.chunking.basic import chunk_elements
|
||||
from unstructured.documents.elements import CompositeElement, Text, Title
|
||||
from unstructured.partition.docx import partition_docx
|
||||
|
||||
|
||||
def test_it_chunks_a_document_when_basic_chunking_strategy_is_specified_on_partition_function():
|
||||
"""Basic chunking can be combined with partitioning, exercising the decorator."""
|
||||
filename = "example-docs/handbook-1p.docx"
|
||||
|
||||
chunks = partition_docx(filename, chunking_strategy="basic")
|
||||
|
||||
assert chunks == [
|
||||
CompositeElement(
|
||||
"US Trustee Handbook\n\nCHAPTER 1\n\nINTRODUCTION\n\nCHAPTER 1 – INTRODUCTION"
|
||||
"\n\nA. PURPOSE"
|
||||
),
|
||||
CompositeElement(
|
||||
"The United States Trustee appoints and supervises standing trustees and monitors and"
|
||||
" supervises cases under chapter 13 of title 11 of the United States Code. 28 U.S.C."
|
||||
" § 586(b). The Handbook, issued as part of our duties under 28 U.S.C. § 586,"
|
||||
" establishes or clarifies the position of the United States Trustee Program (Program)"
|
||||
" on the duties owed by a standing trustee to the debtors, creditors, other parties in"
|
||||
" interest, and the United States Trustee. The Handbook does not present a full and"
|
||||
),
|
||||
CompositeElement(
|
||||
"complete statement of the law; it should not be used as a substitute for legal"
|
||||
" research and analysis. The standing trustee must be familiar with relevant"
|
||||
" provisions of the Bankruptcy Code, Federal Rules of Bankruptcy Procedure (Rules),"
|
||||
" any local bankruptcy rules, and case law. 11 U.S.C. § 321, 28 U.S.C. § 586,"
|
||||
" 28 C.F.R. § 58.6(a)(3). Standing trustees are encouraged to follow Practice Tips"
|
||||
" identified in this Handbook but these are not considered mandatory."
|
||||
),
|
||||
CompositeElement(
|
||||
"Nothing in this Handbook should be construed to excuse the standing trustee from"
|
||||
" complying with all duties imposed by the Bankruptcy Code and Rules, local rules, and"
|
||||
" orders of the court. The standing trustee should notify the United States Trustee"
|
||||
" whenever the provision of the Handbook conflicts with the local rules or orders of"
|
||||
" the court. The standing trustee is accountable for all duties set forth in this"
|
||||
" Handbook, but need not personally perform any duty unless otherwise indicated. All"
|
||||
),
|
||||
CompositeElement(
|
||||
"statutory references in this Handbook refer to the Bankruptcy Code, 11 U.S.C. § 101"
|
||||
" et seq., unless otherwise indicated."
|
||||
),
|
||||
CompositeElement(
|
||||
"This Handbook does not create additional rights against the standing trustee or"
|
||||
" United States Trustee in favor of other parties.\n\nB. ROLE OF THE UNITED STATES"
|
||||
" TRUSTEE"
|
||||
),
|
||||
CompositeElement(
|
||||
"The Bankruptcy Reform Act of 1978 removed the bankruptcy judge from the"
|
||||
" responsibilities for daytoday administration of cases. Debtors, creditors, and"
|
||||
" third parties with adverse interests to the trustee were concerned that the court,"
|
||||
" which previously appointed and supervised the trustee, would not impartially"
|
||||
" adjudicate their rights as adversaries of that trustee. To address these concerns,"
|
||||
" judicial and administrative functions within the bankruptcy system were bifurcated."
|
||||
),
|
||||
CompositeElement(
|
||||
"Many administrative functions formerly performed by the court were placed within the"
|
||||
" Department of Justice through the creation of the Program. Among the administrative"
|
||||
" functions assigned to the United States Trustee were the appointment and supervision"
|
||||
" of chapter 13 trustees./ This Handbook is issued under the authority of the"
|
||||
" Program’s enabling statutes.\n\nC. STATUTORY DUTIES OF A STANDING TRUSTEE"
|
||||
),
|
||||
CompositeElement(
|
||||
"The standing trustee has a fiduciary responsibility to the bankruptcy estate. The"
|
||||
" standing trustee is more than a mere disbursing agent. The standing trustee must"
|
||||
" be personally involved in the trustee operation. If the standing trustee is or"
|
||||
" becomes unable to perform the duties and responsibilities of a standing trustee,"
|
||||
" the standing trustee must immediately advise the United States Trustee."
|
||||
" 28 U.S.C. § 586(b), 28 C.F.R. § 58.4(b) referencing 28 C.F.R. § 58.3(b)."
|
||||
),
|
||||
CompositeElement(
|
||||
"Although this Handbook is not intended to be a complete statutory reference, the"
|
||||
" standing trustee’s primary statutory duties are set forth in 11 U.S.C. § 1302, which"
|
||||
" incorporates by reference some of the duties of chapter 7 trustees found in"
|
||||
" 11 U.S.C. § 704. These duties include, but are not limited to, the"
|
||||
" following:\n\nCopyright"
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_it_chunks_elements_when_the_user_already_has_them():
|
||||
elements = [
|
||||
Title("Introduction"),
|
||||
Text(
|
||||
# --------------------------------------------------------- 64 -v
|
||||
"Lorem ipsum dolor sit amet consectetur adipiscing elit. In rhoncus ipsum sed lectus"
|
||||
" porta volutpat.",
|
||||
),
|
||||
]
|
||||
|
||||
chunks = chunk_elements(elements, max_characters=64)
|
||||
|
||||
assert chunks == [
|
||||
CompositeElement("Introduction"),
|
||||
# -- splits on even word boundary, not mid-"rhoncus" --
|
||||
CompositeElement("Lorem ipsum dolor sit amet consectetur adipiscing elit. In"),
|
||||
CompositeElement("rhoncus ipsum sed lectus porta volutpat."),
|
||||
]
|
||||
|
||||
|
||||
def test_it_includes_original_elements_as_metadata_when_requested():
|
||||
element = Title("Introduction")
|
||||
element_2 = Text("Lorem ipsum dolor sit amet consectetur adipiscing elit.")
|
||||
element_3 = Text("In rhoncus ipsum sed lectus porta volutpat.")
|
||||
|
||||
chunks = chunk_elements(
|
||||
[element, element_2, element_3], max_characters=70, include_orig_elements=True
|
||||
)
|
||||
|
||||
assert len(chunks) == 2
|
||||
chunk = chunks[0]
|
||||
assert chunk == CompositeElement(
|
||||
"Introduction\n\nLorem ipsum dolor sit amet consectetur adipiscing elit."
|
||||
)
|
||||
assert chunk.metadata.orig_elements == [element, element_2]
|
||||
# --
|
||||
chunk = chunks[1]
|
||||
assert chunk == CompositeElement("In rhoncus ipsum sed lectus porta volutpat.")
|
||||
assert chunk.metadata.orig_elements == [element_3]
|
||||
|
||||
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
# UNIT TESTS
|
||||
# ------------------------------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Describe_chunk_elements:
|
||||
"""Unit-test suite for `unstructured.chunking.basic.chunk_elements()` function."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "expected_value"),
|
||||
[
|
||||
({"include_orig_elements": True}, True),
|
||||
({"include_orig_elements": False}, False),
|
||||
({"include_orig_elements": None}, True),
|
||||
({}, True),
|
||||
],
|
||||
)
|
||||
def it_supports_the_include_orig_elements_option(
|
||||
self, kwargs: dict[str, Any], expected_value: bool, _chunk_elements_: Mock
|
||||
):
|
||||
# -- this line would raise if "include_orig_elements" was not an available parameter on
|
||||
# -- `chunk_elements()`.
|
||||
chunk_elements([], **kwargs)
|
||||
|
||||
_, opts = _chunk_elements_.call_args.args
|
||||
assert opts.include_orig_elements is expected_value
|
||||
|
||||
# -- fixtures --------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture()
|
||||
def _chunk_elements_(self, request: FixtureRequest):
|
||||
return function_mock(request, "unstructured.chunking.basic._chunk_elements")
|
||||
@@ -0,0 +1,92 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Unit-test suite for the `unstructured.chunking.dispatch` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Iterable, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from unstructured.chunking import add_chunking_strategy, register_chunking_strategy
|
||||
from unstructured.chunking.dispatch import _ChunkerSpec, chunk
|
||||
from unstructured.documents.elements import CompositeElement, Element, Text
|
||||
|
||||
|
||||
class Describe_add_chunking_strategy:
|
||||
"""Unit-test suite for `unstructured.chunking.add_chunking_strategy()` decorator."""
|
||||
|
||||
def it_dispatches_the_partitioned_elements_to_the_indicated_chunker(self):
|
||||
decorated_partitioner = add_chunking_strategy(partition_this)
|
||||
|
||||
chunks = decorated_partitioner(chunking_strategy="basic")
|
||||
|
||||
assert chunks == [CompositeElement("Lorem ipsum.\n\nSit amet.")]
|
||||
|
||||
def but_it_skips_dispatch_when_no_chunking_strategy_is_specified(self):
|
||||
decorated_partitioner = add_chunking_strategy(partition_this)
|
||||
|
||||
elements = decorated_partitioner()
|
||||
|
||||
assert elements == [Text("Lorem ipsum."), Text("Sit amet.")]
|
||||
|
||||
|
||||
class Describe_chunk:
|
||||
"""Unit-test suite for `unstructured.chunking.dispatch.chunk()` function."""
|
||||
|
||||
def it_dispatches_to_the_chunker_registered_for_the_chunking_strategy(self):
|
||||
register_chunking_strategy("by_something_else", chunk_by_something_else)
|
||||
kwargs = {
|
||||
"max_characters": 750,
|
||||
# -- unused kwargs shouldn't cause a problem; in general `kwargs` will contain all
|
||||
# -- keyword arguments used in the partitioning call.
|
||||
"foo": "bar",
|
||||
}
|
||||
|
||||
chunks = chunk([Text("Lorem"), Text("Ipsum")], "by_something_else", **kwargs)
|
||||
|
||||
assert chunks == [
|
||||
CompositeElement("chunked 2 elements with `(max_characters=750, whizbang=None)`")
|
||||
]
|
||||
|
||||
def it_raises_when_the_requested_chunking_strategy_is_not_registered(self):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="unrecognized chunking strategy 'foobar'",
|
||||
):
|
||||
chunk(elements=[], chunking_strategy="foobar")
|
||||
|
||||
|
||||
class Describe_ChunkerSpec:
|
||||
"""Unit-test suite for `unstructured.chunking.dispatch._ChunkerSpec` objects."""
|
||||
|
||||
def it_provides_access_to_the_chunking_function(self):
|
||||
spec = _ChunkerSpec(chunk_by_something_else)
|
||||
assert spec.chunker is chunk_by_something_else
|
||||
|
||||
def it_knows_which_keyword_args_the_chunking_function_can_accept(self):
|
||||
spec = _ChunkerSpec(chunk_by_something_else)
|
||||
assert spec.kw_arg_names == ("max_characters", "whizbang")
|
||||
|
||||
|
||||
# -- MODULE-LEVEL FIXTURES -----------------------------------------------------------------------
|
||||
|
||||
|
||||
def chunk_by_something_else(
|
||||
elements: Iterable[Element],
|
||||
max_characters: Optional[int] = None,
|
||||
whizbang: Optional[float] = None,
|
||||
) -> list[Element]:
|
||||
"""A "fake" minimal chunker suitable for use in tests."""
|
||||
els = list(elements)
|
||||
return [
|
||||
CompositeElement(
|
||||
f"chunked {len(els)} elements with"
|
||||
f" `(max_characters={max_characters}, whizbang={whizbang})`"
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def partition_this(**kwargs: Any) -> list[Element]:
|
||||
"""A fake partitioner."""
|
||||
return [Text("Lorem ipsum."), Text("Sit amet.")]
|
||||
@@ -0,0 +1,90 @@
|
||||
from functools import partial
|
||||
|
||||
import pytest
|
||||
|
||||
from unstructured.chunking.basic import chunk_elements
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import ElementMetadata, NarrativeText, Text, Title
|
||||
|
||||
|
||||
@pytest.fixture(params=[chunk_elements, partial(chunk_by_title, combine_text_under_n_chars=0)])
|
||||
def chunking_fn(request):
|
||||
return request.param
|
||||
|
||||
|
||||
def test_combining_html_metadata_when_multiple_elements_in_composite_element(chunking_fn):
|
||||
metadata_1 = '<h1 class="Title" id="1">Header </h1>'
|
||||
metadata_2 = '<time class="CalendarDate" id="2">Date: October 30, 2023 </time>'
|
||||
metadata_3 = (
|
||||
'<form class="Form" id="3"> '
|
||||
'<label class="FormField" for="company-name" id="4">Form field name </label>'
|
||||
'<input class="FormFieldValue" id="5" value="Example value" />'
|
||||
"</form>"
|
||||
)
|
||||
combined_metadata = " ".join([metadata_1, metadata_2, metadata_3])
|
||||
|
||||
elements = [
|
||||
Title(text="Header", metadata=ElementMetadata(text_as_html=metadata_1)),
|
||||
Text(text="Date: October 30, 2023", metadata=ElementMetadata(text_as_html=metadata_2)),
|
||||
Text(
|
||||
text="Form field name Example value", metadata=ElementMetadata(text_as_html=metadata_3)
|
||||
),
|
||||
]
|
||||
chunks = chunking_fn(elements)
|
||||
assert len(chunks) == 1
|
||||
assert chunks[0].metadata.text_as_html == combined_metadata
|
||||
|
||||
|
||||
def test_combining_html_metadata_with_nested_relationship_between_elements(chunking_fn):
|
||||
"""
|
||||
Ground truth
|
||||
<Document>
|
||||
<Page>
|
||||
<Section>
|
||||
<p>First</p>
|
||||
<p>Second</p>
|
||||
</Section>
|
||||
</Page>
|
||||
</Document>
|
||||
Elements: Document, Page, Section, Paragraph, Paragraph
|
||||
Chunk 1: Document, Page, Section, Paragraph
|
||||
|
||||
Chunk 2:
|
||||
Paragraph
|
||||
"""
|
||||
|
||||
metadata_1 = '<div class="Section" id="1" />'
|
||||
metadata_2 = '<p class="Paragraph" id="2">First </p>'
|
||||
metadata_3 = '<p class="Paragraph" id="3">Second </p>'
|
||||
|
||||
elements = [
|
||||
Text(text="", metadata=ElementMetadata(text_as_html=metadata_1)),
|
||||
NarrativeText(
|
||||
text="First", metadata=ElementMetadata(text_as_html=metadata_2, parent_id="1")
|
||||
),
|
||||
NarrativeText(
|
||||
text="Second", metadata=ElementMetadata(text_as_html=metadata_3, parent_id="1")
|
||||
),
|
||||
]
|
||||
chunks = chunking_fn(elements, max_characters=6)
|
||||
assert len(chunks) == 2
|
||||
assert chunks[0].text == "First"
|
||||
assert chunks[1].text == "Second"
|
||||
|
||||
assert chunks[0].metadata.text_as_html == metadata_1 + " " + metadata_2
|
||||
assert chunks[1].metadata.text_as_html == metadata_3
|
||||
|
||||
|
||||
def test_html_metadata_exist_in_both_element_when_text_is_split(chunking_fn):
|
||||
"""Mimic behaviour of elements with non-html metadata"""
|
||||
metadata_1 = '<h1 class="Title" id="1">Header </h1>'
|
||||
elements = [
|
||||
Title(text="Header", metadata=ElementMetadata(text_as_html=metadata_1)),
|
||||
]
|
||||
chunks = chunking_fn(elements, max_characters=3)
|
||||
assert len(chunks) == 2
|
||||
|
||||
assert chunks[0].text == "Hea"
|
||||
assert chunks[1].text == "der"
|
||||
assert chunks[0].metadata.text_as_html == '<h1 class="Title" id="1">Header </h1>'
|
||||
assert chunks[1].metadata.text_as_html == '<h1 class="Title" id="1">Header </h1>'
|
||||
@@ -0,0 +1,541 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Test suite for the `unstructured.chunking.title` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import FixtureRequest, Mock, function_mock, input_path
|
||||
from unstructured.chunking.base import CHUNK_MULTI_PAGE_DEFAULT
|
||||
from unstructured.chunking.title import _ByTitleChunkingOptions, chunk_by_title
|
||||
from unstructured.documents.coordinates import CoordinateSystem
|
||||
from unstructured.documents.elements import (
|
||||
CheckBox,
|
||||
CompositeElement,
|
||||
CoordinatesMetadata,
|
||||
Element,
|
||||
ElementMetadata,
|
||||
ListItem,
|
||||
Table,
|
||||
TableChunk,
|
||||
Text,
|
||||
Title,
|
||||
)
|
||||
from unstructured.partition.html import partition_html
|
||||
from unstructured.staging.base import elements_from_json
|
||||
|
||||
# ================================================================================================
|
||||
# INTEGRATION-TESTS
|
||||
# ================================================================================================
|
||||
# These test `chunk_by_title()` as an integrated whole, calling `chunk_by_title()` and inspecting
|
||||
# the outputs.
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
def test_it_chunks_text_followed_by_table_together_when_both_fit():
|
||||
elements = elements_from_json(input_path("chunking/title_table_200.json"))
|
||||
|
||||
chunks = chunk_by_title(elements, combine_text_under_n_chars=0)
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert isinstance(chunks[0], CompositeElement)
|
||||
|
||||
|
||||
def test_it_chunks_table_followed_by_text_together_when_both_fit():
|
||||
elements = elements_from_json(input_path("chunking/table_text_200.json"))
|
||||
|
||||
# -- disable chunk combining so we test pre-chunking behavior, not chunk-combining --
|
||||
chunks = chunk_by_title(elements, combine_text_under_n_chars=0)
|
||||
|
||||
assert len(chunks) == 1
|
||||
assert isinstance(chunks[0], CompositeElement)
|
||||
|
||||
|
||||
def test_it_splits_oversized_table():
|
||||
elements = elements_from_json(input_path("chunking/table_2000.json"))
|
||||
|
||||
chunks = chunk_by_title(elements)
|
||||
|
||||
assert len(chunks) == 5
|
||||
assert all(isinstance(chunk, TableChunk) for chunk in chunks)
|
||||
|
||||
|
||||
def test_it_starts_new_chunk_for_table_after_full_text_chunk():
|
||||
elements = elements_from_json(input_path("chunking/long_text_table_200.json"))
|
||||
|
||||
chunks = chunk_by_title(elements, max_characters=250)
|
||||
|
||||
assert len(chunks) == 2
|
||||
assert [type(chunk) for chunk in chunks] == [CompositeElement, Table]
|
||||
|
||||
|
||||
def test_it_starts_new_chunk_for_text_after_full_table_chunk():
|
||||
elements = elements_from_json(input_path("chunking/full_table_long_text_250.json"))
|
||||
|
||||
chunks = chunk_by_title(elements, max_characters=250)
|
||||
|
||||
assert len(chunks) == 2
|
||||
assert [type(chunk) for chunk in chunks] == [Table, CompositeElement]
|
||||
|
||||
|
||||
def test_it_splits_a_large_text_element_into_multiple_chunks():
|
||||
elements: list[Element] = [
|
||||
Title("Introduction"),
|
||||
Text(
|
||||
"Lorem ipsum dolor sit amet consectetur adipiscing elit. In rhoncus ipsum sed lectus"
|
||||
" porta volutpat.",
|
||||
),
|
||||
]
|
||||
|
||||
chunks = chunk_by_title(elements, max_characters=50)
|
||||
|
||||
assert chunks == [
|
||||
CompositeElement("Introduction"),
|
||||
CompositeElement("Lorem ipsum dolor sit amet consectetur adipiscing"),
|
||||
CompositeElement("elit. In rhoncus ipsum sed lectus porta volutpat."),
|
||||
]
|
||||
|
||||
|
||||
def test_it_splits_elements_by_title_and_table():
|
||||
elements: list[Element] = [
|
||||
Title("A Great Day"),
|
||||
Text("Today is a great day."),
|
||||
Text("It is sunny outside."),
|
||||
Table("Heading\nCell text"),
|
||||
Title("An Okay Day"),
|
||||
Text("Today is an okay day."),
|
||||
Text("It is rainy outside."),
|
||||
Title("A Bad Day"),
|
||||
Text("Today is a bad day."),
|
||||
Text("It is storming outside."),
|
||||
CheckBox(),
|
||||
]
|
||||
|
||||
chunks = chunk_by_title(elements, combine_text_under_n_chars=0, include_orig_elements=True)
|
||||
|
||||
assert len(chunks) == 3
|
||||
# --
|
||||
chunk = chunks[0]
|
||||
assert isinstance(chunk, CompositeElement)
|
||||
assert chunk.metadata.orig_elements == [
|
||||
Title("A Great Day"),
|
||||
Text("Today is a great day."),
|
||||
Text("It is sunny outside."),
|
||||
Table("Heading\nCell text"),
|
||||
]
|
||||
# --
|
||||
chunk = chunks[1]
|
||||
assert isinstance(chunk, CompositeElement)
|
||||
assert chunk.metadata.orig_elements == [
|
||||
Title("An Okay Day"),
|
||||
Text("Today is an okay day."),
|
||||
Text("It is rainy outside."),
|
||||
]
|
||||
# --
|
||||
chunk = chunks[2]
|
||||
assert isinstance(chunk, CompositeElement)
|
||||
assert chunk.metadata.orig_elements == [
|
||||
Title("A Bad Day"),
|
||||
Text("Today is a bad day."),
|
||||
Text("It is storming outside."),
|
||||
CheckBox(),
|
||||
]
|
||||
|
||||
|
||||
def test_chunk_by_title():
|
||||
elements: list[Element] = [
|
||||
Title("A Great Day", metadata=ElementMetadata(emphasized_text_contents=["Day"])),
|
||||
Text("Today is a great day.", metadata=ElementMetadata(emphasized_text_contents=["day"])),
|
||||
Text("It is sunny outside."),
|
||||
Table("Heading\nCell text"),
|
||||
Title("An Okay Day"),
|
||||
Text("Today is an okay day."),
|
||||
Text("It is rainy outside."),
|
||||
Title("A Bad Day"),
|
||||
Text("Today is a bad day."),
|
||||
Text("It is storming outside."),
|
||||
CheckBox(),
|
||||
]
|
||||
|
||||
chunks = chunk_by_title(elements, combine_text_under_n_chars=0, include_orig_elements=False)
|
||||
|
||||
assert chunks == [
|
||||
CompositeElement(
|
||||
"A Great Day\n\nToday is a great day.\n\nIt is sunny outside.\n\nHeading Cell text"
|
||||
),
|
||||
CompositeElement("An Okay Day\n\nToday is an okay day.\n\nIt is rainy outside."),
|
||||
CompositeElement(
|
||||
"A Bad Day\n\nToday is a bad day.\n\nIt is storming outside.",
|
||||
),
|
||||
]
|
||||
assert chunks[0].metadata == ElementMetadata(emphasized_text_contents=["Day", "day"])
|
||||
|
||||
|
||||
def test_chunk_by_title_separates_by_page_number():
|
||||
elements: list[Element] = [
|
||||
Title("A Great Day", metadata=ElementMetadata(page_number=1)),
|
||||
Text("Today is a great day.", metadata=ElementMetadata(page_number=2)),
|
||||
Text("It is sunny outside.", metadata=ElementMetadata(page_number=2)),
|
||||
Table("Heading\nCell text"),
|
||||
Title("An Okay Day"),
|
||||
Text("Today is an okay day."),
|
||||
Text("It is rainy outside."),
|
||||
Title("A Bad Day"),
|
||||
Text("Today is a bad day."),
|
||||
Text("It is storming outside."),
|
||||
CheckBox(),
|
||||
]
|
||||
chunks = chunk_by_title(elements, multipage_sections=False, combine_text_under_n_chars=0)
|
||||
|
||||
assert chunks == [
|
||||
CompositeElement(
|
||||
"A Great Day",
|
||||
),
|
||||
CompositeElement("Today is a great day.\n\nIt is sunny outside.\n\nHeading Cell text"),
|
||||
CompositeElement("An Okay Day\n\nToday is an okay day.\n\nIt is rainy outside."),
|
||||
CompositeElement(
|
||||
"A Bad Day\n\nToday is a bad day.\n\nIt is storming outside.",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_chuck_by_title_respects_multipage():
|
||||
elements: list[Element] = [
|
||||
Title("A Great Day", metadata=ElementMetadata(page_number=1)),
|
||||
Text("Today is a great day.", metadata=ElementMetadata(page_number=2)),
|
||||
Text("It is sunny outside.", metadata=ElementMetadata(page_number=2)),
|
||||
Table("Heading\nCell text"),
|
||||
Title("An Okay Day"),
|
||||
Text("Today is an okay day."),
|
||||
Text("It is rainy outside."),
|
||||
Title("A Bad Day"),
|
||||
Text("Today is a bad day."),
|
||||
Text("It is storming outside."),
|
||||
CheckBox(),
|
||||
]
|
||||
chunks = chunk_by_title(elements, multipage_sections=True, combine_text_under_n_chars=0)
|
||||
assert chunks == [
|
||||
CompositeElement(
|
||||
"A Great Day\n\nToday is a great day.\n\nIt is sunny outside.\n\nHeading Cell text"
|
||||
),
|
||||
CompositeElement("An Okay Day\n\nToday is an okay day.\n\nIt is rainy outside."),
|
||||
CompositeElement(
|
||||
"A Bad Day\n\nToday is a bad day.\n\nIt is storming outside.",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_chunk_by_title_groups_across_pages():
|
||||
elements: list[Element] = [
|
||||
Title("A Great Day", metadata=ElementMetadata(page_number=1)),
|
||||
Text("Today is a great day.", metadata=ElementMetadata(page_number=2)),
|
||||
Text("It is sunny outside.", metadata=ElementMetadata(page_number=2)),
|
||||
Table("Heading\nCell text"),
|
||||
Title("An Okay Day"),
|
||||
Text("Today is an okay day."),
|
||||
Text("It is rainy outside."),
|
||||
Title("A Bad Day"),
|
||||
Text("Today is a bad day."),
|
||||
Text("It is storming outside."),
|
||||
CheckBox(),
|
||||
]
|
||||
chunks = chunk_by_title(elements, multipage_sections=True, combine_text_under_n_chars=0)
|
||||
|
||||
assert chunks == [
|
||||
CompositeElement(
|
||||
"A Great Day\n\nToday is a great day.\n\nIt is sunny outside.\n\nHeading Cell text"
|
||||
),
|
||||
CompositeElement("An Okay Day\n\nToday is an okay day.\n\nIt is rainy outside."),
|
||||
CompositeElement(
|
||||
"A Bad Day\n\nToday is a bad day.\n\nIt is storming outside.",
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_add_chunking_strategy_on_partition_html():
|
||||
filename = "example-docs/example-10k-1p.html"
|
||||
chunk_elements = partition_html(filename, chunking_strategy="by_title")
|
||||
elements = partition_html(filename)
|
||||
chunks = chunk_by_title(elements)
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_add_chunking_strategy_respects_max_characters():
|
||||
filename = "example-docs/example-10k-1p.html"
|
||||
chunk_elements = partition_html(
|
||||
filename,
|
||||
chunking_strategy="by_title",
|
||||
combine_text_under_n_chars=0,
|
||||
new_after_n_chars=50,
|
||||
max_characters=100,
|
||||
)
|
||||
elements = partition_html(filename)
|
||||
chunks = chunk_by_title(
|
||||
elements,
|
||||
combine_text_under_n_chars=0,
|
||||
new_after_n_chars=50,
|
||||
max_characters=100,
|
||||
)
|
||||
|
||||
for chunk in chunks:
|
||||
assert isinstance(chunk, Text)
|
||||
assert len(chunk.text) <= 100
|
||||
for chunk_element in chunk_elements:
|
||||
assert isinstance(chunk_element, Text)
|
||||
assert len(chunk_element.text) <= 100
|
||||
assert chunk_elements != elements
|
||||
assert chunk_elements == chunks
|
||||
|
||||
|
||||
def test_chunk_by_title_drops_detection_class_prob():
|
||||
elements: list[Element] = [
|
||||
Title(
|
||||
"A Great Day",
|
||||
metadata=ElementMetadata(
|
||||
detection_class_prob=0.5,
|
||||
),
|
||||
),
|
||||
Text(
|
||||
"Today is a great day.",
|
||||
metadata=ElementMetadata(
|
||||
detection_class_prob=0.62,
|
||||
),
|
||||
),
|
||||
Text(
|
||||
"It is sunny outside.",
|
||||
metadata=ElementMetadata(
|
||||
detection_class_prob=0.73,
|
||||
),
|
||||
),
|
||||
Title(
|
||||
"An Okay Day",
|
||||
metadata=ElementMetadata(
|
||||
detection_class_prob=0.84,
|
||||
),
|
||||
),
|
||||
Text(
|
||||
"Today is an okay day.",
|
||||
metadata=ElementMetadata(
|
||||
detection_class_prob=0.95,
|
||||
),
|
||||
),
|
||||
]
|
||||
chunks = chunk_by_title(elements, combine_text_under_n_chars=0)
|
||||
assert str(chunks[0]) == str(
|
||||
CompositeElement("A Great Day\n\nToday is a great day.\n\nIt is sunny outside."),
|
||||
)
|
||||
assert str(chunks[1]) == str(CompositeElement("An Okay Day\n\nToday is an okay day."))
|
||||
|
||||
|
||||
def test_chunk_by_title_drops_extra_metadata():
|
||||
elements: list[Element] = [
|
||||
Title(
|
||||
"A Great Day",
|
||||
metadata=ElementMetadata(
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=(
|
||||
(0.1, 0.1),
|
||||
(0.2, 0.1),
|
||||
(0.1, 0.2),
|
||||
(0.2, 0.2),
|
||||
),
|
||||
system=CoordinateSystem(width=0.1, height=0.1),
|
||||
),
|
||||
),
|
||||
),
|
||||
Text(
|
||||
"Today is a great day.",
|
||||
metadata=ElementMetadata(
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=(
|
||||
(0.2, 0.2),
|
||||
(0.3, 0.2),
|
||||
(0.2, 0.3),
|
||||
(0.3, 0.3),
|
||||
),
|
||||
system=CoordinateSystem(width=0.2, height=0.2),
|
||||
),
|
||||
),
|
||||
),
|
||||
Text(
|
||||
"It is sunny outside.",
|
||||
metadata=ElementMetadata(
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=(
|
||||
(0.3, 0.3),
|
||||
(0.4, 0.3),
|
||||
(0.3, 0.4),
|
||||
(0.4, 0.4),
|
||||
),
|
||||
system=CoordinateSystem(width=0.3, height=0.3),
|
||||
),
|
||||
),
|
||||
),
|
||||
Title(
|
||||
"An Okay Day",
|
||||
metadata=ElementMetadata(
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=(
|
||||
(0.3, 0.3),
|
||||
(0.4, 0.3),
|
||||
(0.3, 0.4),
|
||||
(0.4, 0.4),
|
||||
),
|
||||
system=CoordinateSystem(width=0.3, height=0.3),
|
||||
),
|
||||
),
|
||||
),
|
||||
Text(
|
||||
"Today is an okay day.",
|
||||
metadata=ElementMetadata(
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=(
|
||||
(0.4, 0.4),
|
||||
(0.5, 0.4),
|
||||
(0.4, 0.5),
|
||||
(0.5, 0.5),
|
||||
),
|
||||
system=CoordinateSystem(width=0.4, height=0.4),
|
||||
),
|
||||
),
|
||||
),
|
||||
]
|
||||
|
||||
chunks = chunk_by_title(elements, combine_text_under_n_chars=0)
|
||||
|
||||
assert str(chunks[0]) == str(
|
||||
CompositeElement("A Great Day\n\nToday is a great day.\n\nIt is sunny outside."),
|
||||
)
|
||||
|
||||
assert str(chunks[1]) == str(CompositeElement("An Okay Day\n\nToday is an okay day."))
|
||||
|
||||
|
||||
def test_it_considers_separator_length_when_pre_chunking():
|
||||
"""PreChunker includes length of separators when computing remaining space."""
|
||||
elements: list[Element] = [
|
||||
Title("Chunking Priorities"), # 19 chars
|
||||
ListItem("Divide text into manageable chunks"), # 34 chars
|
||||
ListItem("Preserve semantic boundaries"), # 28 chars
|
||||
ListItem("Minimize mid-text chunk-splitting"), # 33 chars
|
||||
] # 114 chars total but 120 chars with separators
|
||||
|
||||
chunks = chunk_by_title(elements, max_characters=115)
|
||||
|
||||
assert chunks == [
|
||||
CompositeElement(
|
||||
"Chunking Priorities"
|
||||
"\n\nDivide text into manageable chunks"
|
||||
"\n\nPreserve semantic boundaries",
|
||||
),
|
||||
CompositeElement("Minimize mid-text chunk-splitting"),
|
||||
]
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# UNIT-TESTS
|
||||
# ================================================================================================
|
||||
# These test individual components in isolation so can exercise all edge cases while still
|
||||
# performing well.
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class Describe_chunk_by_title:
|
||||
"""Unit-test suite for `unstructured.chunking.title.chunk_by_title()` function."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("kwargs", "expected_value"),
|
||||
[
|
||||
({"include_orig_elements": True}, True),
|
||||
({"include_orig_elements": False}, False),
|
||||
({"include_orig_elements": None}, True),
|
||||
({}, True),
|
||||
],
|
||||
)
|
||||
def it_supports_the_include_orig_elements_option(
|
||||
self, kwargs: dict[str, Any], expected_value: bool, _chunk_by_title_: Mock
|
||||
):
|
||||
# -- this line would raise if "include_orig_elements" was not an available parameter on
|
||||
# -- `chunk_by_title()`.
|
||||
chunk_by_title([], **kwargs)
|
||||
|
||||
_, opts = _chunk_by_title_.call_args.args
|
||||
assert opts.include_orig_elements is expected_value
|
||||
|
||||
# -- fixtures --------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture()
|
||||
def _chunk_by_title_(self, request: FixtureRequest):
|
||||
return function_mock(request, "unstructured.chunking.title._chunk_by_title")
|
||||
|
||||
|
||||
class Describe_ByTitleChunkingOptions:
|
||||
"""Unit-test suite for `unstructured.chunking.title._ByTitleChunkingOptions` objects."""
|
||||
|
||||
@pytest.mark.parametrize("n_chars", [-1, -42])
|
||||
def it_rejects_combine_text_under_n_chars_for_n_less_than_zero(self, n_chars: int):
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=f"'combine_text_under_n_chars' argument must be >= 0, got {n_chars}",
|
||||
):
|
||||
_ByTitleChunkingOptions.new(combine_text_under_n_chars=n_chars)
|
||||
|
||||
def it_accepts_0_for_combine_text_under_n_chars_to_disable_chunk_combining(self):
|
||||
"""Specifying `combine_text_under_n_chars=0` is how a caller disables chunk-combining."""
|
||||
opts = _ByTitleChunkingOptions(combine_text_under_n_chars=0)
|
||||
assert opts.combine_text_under_n_chars == 0
|
||||
|
||||
def it_does_not_complain_when_specifying_combine_text_under_n_chars_by_itself(self):
|
||||
"""Caller can specify `combine_text_under_n_chars` arg without specifying other options."""
|
||||
try:
|
||||
opts = _ByTitleChunkingOptions(combine_text_under_n_chars=50)
|
||||
except ValueError:
|
||||
pytest.fail("did not accept `combine_text_under_n_chars` as option by itself")
|
||||
|
||||
assert opts.combine_text_under_n_chars == 50
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("combine_text_under_n_chars", "max_characters", "expected_hard_max"),
|
||||
[(600, None, 500), (600, 450, 450)],
|
||||
)
|
||||
def it_rejects_combine_text_under_n_chars_greater_than_maxchars(
|
||||
self, combine_text_under_n_chars: int, max_characters: Optional[int], expected_hard_max: int
|
||||
):
|
||||
"""`combine_text_under_n_chars` > `max_characters` can produce behavior confusing to users.
|
||||
|
||||
The behavior is no different from `combine_text_under_n_chars == max_characters`, but if
|
||||
`max_characters` is left to default (500) and `combine_text_under_n_chars` is set to a
|
||||
larger number like 1500 then it can look like chunk-combining isn't working.
|
||||
"""
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
"'combine_text_under_n_chars' argument must not exceed `max_characters` value,"
|
||||
f" got {combine_text_under_n_chars} > {expected_hard_max}"
|
||||
),
|
||||
):
|
||||
_ByTitleChunkingOptions.new(
|
||||
max_characters=max_characters, combine_text_under_n_chars=combine_text_under_n_chars
|
||||
)
|
||||
|
||||
def it_does_not_complain_when_specifying_new_after_n_chars_by_itself(self):
|
||||
"""Caller can specify `new_after_n_chars` arg without specifying any other options."""
|
||||
try:
|
||||
opts = _ByTitleChunkingOptions.new(new_after_n_chars=200)
|
||||
except ValueError:
|
||||
pytest.fail("did not accept `new_after_n_chars` as option by itself")
|
||||
|
||||
assert opts.soft_max == 200
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("multipage_sections", "expected_value"),
|
||||
[(True, True), (False, False), (None, CHUNK_MULTI_PAGE_DEFAULT)],
|
||||
)
|
||||
def it_knows_whether_to_break_chunks_on_page_boundaries(
|
||||
self, multipage_sections: bool, expected_value: bool
|
||||
):
|
||||
opts = _ByTitleChunkingOptions(multipage_sections=multipage_sections)
|
||||
assert opts.multipage_sections is expected_value
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,302 @@
|
||||
import re
|
||||
|
||||
import pytest
|
||||
|
||||
from unstructured.cleaners import core
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
(
|
||||
"\x88This text contains non-ascii characters!\x88",
|
||||
"This text contains non-ascii characters!",
|
||||
),
|
||||
("\x93A lovely quote!\x94", "A lovely quote!"),
|
||||
("● An excellent point! ●●●", " An excellent point! "),
|
||||
("Item\xa01A", "Item1A"),
|
||||
("Our dog's bowl.", "Our dog's bowl."),
|
||||
("5 w=E2=80=99s", "5 w=E2=80=99s"),
|
||||
],
|
||||
)
|
||||
def test_clean_non_ascii_chars(text, expected):
|
||||
assert core.clean_non_ascii_chars(text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("● An excellent point!", "An excellent point!"),
|
||||
("● An excellent point! ●●●", "An excellent point! ●●●"),
|
||||
("An excellent point!", "An excellent point!"),
|
||||
("Morse code! ●●●", "Morse code! ●●●"),
|
||||
],
|
||||
)
|
||||
def test_clean_bullets(text, expected):
|
||||
assert core.clean_bullets(text=text) == expected
|
||||
assert core.clean(text=text, bullets=True) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("1. Introduction:", "Introduction:"),
|
||||
("a. Introduction:", "Introduction:"),
|
||||
("20.3 Morse code ●●●", "Morse code ●●●"),
|
||||
("5.3.1 Convolutional Networks ", "Convolutional Networks"),
|
||||
("D.b.C Recurrent Neural Networks", "Recurrent Neural Networks"),
|
||||
("2.b.1 Recurrent Neural Networks", "Recurrent Neural Networks"),
|
||||
("eins. Neural Networks", "eins. Neural Networks"),
|
||||
("bb.c Feed Forward Neural Networks", "Feed Forward Neural Networks"),
|
||||
("aaa.ccc Metrics", "aaa.ccc Metrics"),
|
||||
(" version = 3.8", " version = 3.8"),
|
||||
("1 2. 3 4", "1 2. 3 4"),
|
||||
("1) 2. 3 4", "1) 2. 3 4"),
|
||||
("2,3. Morse code 3. ●●●", "2,3. Morse code 3. ●●●"),
|
||||
("1..2.3 four", "1..2.3 four"),
|
||||
("Fig. 2: The relationship", "Fig. 2: The relationship"),
|
||||
("23 is everywhere", "23 is everywhere"),
|
||||
],
|
||||
)
|
||||
def test_clean_ordered_bullets(text, expected):
|
||||
assert core.clean_ordered_bullets(text=text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("The æther is a classic element.", "The aether is a classic element."),
|
||||
("In old texts, Æsop's fables are", "In old texts, AEsop's fables are"),
|
||||
("The buffer zone is there.", "The buffer zone is there."),
|
||||
("The file was found in the system.", "The file was found in the system."),
|
||||
("She had a flower in her hair.", "She had a flower in her hair."),
|
||||
("The coffin was placed in the grave.", "The coffin was placed in the grave."),
|
||||
("The buffle zone was clearly marked.", "The buffle zone was clearly marked."),
|
||||
("The craſtsman worked with dedication.", "The craftsman worked with dedication."),
|
||||
("The symbol ʪ is very rare.", "The symbol ls is very rare."),
|
||||
("The word 'cœur' means 'heart' in French.", "The word 'coeur' means 'heart' in French."),
|
||||
("The word 'Œuvre' refers to the works", "The word 'OEuvre' refers to the works"),
|
||||
("The ȹ symbol is used in some contexts.", "The qp symbol is used in some contexts."),
|
||||
("The postman delivers mail daily.", "The postman delivers mail daily."),
|
||||
(
|
||||
"The symbol ʦ can be found in certain alphabets.",
|
||||
"The symbol ts can be found in certain alphabets.",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_clean_ligatures(text, expected):
|
||||
assert core.clean_ligatures(text=text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("\x93A lovely quote!\x94", "“A lovely quote!”"),
|
||||
("\x91A lovely quote!\x92", "‘A lovely quote!’"),
|
||||
("Our dog's bowl.", "Our dog's bowl."),
|
||||
],
|
||||
)
|
||||
def test_replace_unicode_quotes(text, expected):
|
||||
assert core.replace_unicode_quotes(text=text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[("5 w=E2=80=99s", "5 w’s")],
|
||||
)
|
||||
def test_replace_mime_encodings(text, expected):
|
||||
assert core.replace_mime_encodings(text=text) == expected
|
||||
|
||||
|
||||
def test_replace_mime_encodings_works_with_different_encodings():
|
||||
text = "5 w=E2=80-99s=E2=80-92"
|
||||
assert core.replace_mime_encodings(text=text, encoding="latin-1") == "5 wâ\x80-99sâ\x80-92"
|
||||
|
||||
|
||||
def test_replace_mime_encodings_works_with_right_to_left_encodings():
|
||||
text = "=EE=E0=E9=E4"
|
||||
assert core.replace_mime_encodings(text=text, encoding="iso-8859-8") == "מאיה"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("“A lovely quote!”", "A lovely quote"),
|
||||
("‘A lovely quote!’", "A lovely quote"),
|
||||
("'()[]{};:'\",.?/\\-_", ""),
|
||||
],
|
||||
)
|
||||
def test_remove_punctuation(text, expected):
|
||||
assert core.remove_punctuation(text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("RISK\n\nFACTORS", "RISK FACTORS"),
|
||||
("Item\xa01A", "Item 1A"),
|
||||
(" Risk factors ", "Risk factors"),
|
||||
("Risk factors ", "Risk factors"),
|
||||
],
|
||||
)
|
||||
def test_clean_extra_whitespace(text, expected):
|
||||
assert core.clean_extra_whitespace(text) == expected
|
||||
assert core.clean(text=text, extra_whitespace=True) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("Risk-factors", "Risk factors"),
|
||||
("Risk – factors", "Risk factors"),
|
||||
("Risk\u2013factors", "Risk factors"),
|
||||
("Risk factors-\u2013", "Risk factors"),
|
||||
],
|
||||
)
|
||||
def test_clean_dashes(text, expected):
|
||||
assert core.clean_dashes(text) == expected
|
||||
assert core.clean(text=text, dashes=True) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("Item 1A:", "Item 1A"),
|
||||
("Item 1A;", "Item 1A"),
|
||||
("Item 1A.", "Item 1A"),
|
||||
("Item 1A,", "Item 1A"),
|
||||
("Item, 1A: ", "Item, 1A"),
|
||||
],
|
||||
)
|
||||
def test_clean_trailing_punctuation(text, expected):
|
||||
assert core.clean_trailing_punctuation(text) == expected
|
||||
assert core.clean(text=text, trailing_punctuation=True) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "pattern", "ignore_case", "strip", "expected"),
|
||||
[
|
||||
("SUMMARY: A great SUMMARY", r"(SUMMARY|DESC):", False, True, "A great SUMMARY"),
|
||||
("DESC: A great SUMMARY", r"(SUMMARY|DESC):", False, True, "A great SUMMARY"),
|
||||
("SUMMARY: A great SUMMARY", r"(SUMMARY|DESC):", False, False, " A great SUMMARY"),
|
||||
("summary: A great SUMMARY", r"(SUMMARY|DESC):", True, True, "A great SUMMARY"),
|
||||
],
|
||||
)
|
||||
def test_clean_prefix(text, pattern, ignore_case, strip, expected):
|
||||
assert core.clean_prefix(text, pattern, ignore_case, strip) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "pattern", "ignore_case", "strip", "expected"),
|
||||
[
|
||||
("The END! END", r"(END|STOP)", False, True, "The END!"),
|
||||
("The END! STOP", r"(END|STOP)", False, True, "The END!"),
|
||||
("The END! END", r"(END|STOP)", False, False, "The END! "),
|
||||
("The END! end", r"(END|STOP)", True, True, "The END!"),
|
||||
],
|
||||
)
|
||||
def test_clean_postfix(text, pattern, ignore_case, strip, expected):
|
||||
assert core.clean_postfix(text, pattern, ignore_case, strip) == expected
|
||||
|
||||
|
||||
def test_group_broken_paragraphs():
|
||||
text = """The big red fox
|
||||
is walking down the lane.
|
||||
|
||||
At the end of the lane
|
||||
the fox met a friendly bear."""
|
||||
|
||||
assert (
|
||||
core.group_broken_paragraphs(text)
|
||||
== """The big red fox is walking down the lane.
|
||||
|
||||
At the end of the lane the fox met a friendly bear."""
|
||||
)
|
||||
|
||||
|
||||
def test_group_broken_paragraphs_non_default_settings():
|
||||
text = """The big red fox
|
||||
|
||||
is walking down the lane.
|
||||
|
||||
|
||||
At the end of the lane
|
||||
|
||||
the fox met a friendly bear."""
|
||||
|
||||
para_split_re = re.compile(r"(\s*\n\s*){3}")
|
||||
|
||||
clean_text = core.group_broken_paragraphs(text, paragraph_split=para_split_re)
|
||||
assert (
|
||||
clean_text
|
||||
== """The big red fox is walking down the lane.
|
||||
|
||||
At the end of the lane the fox met a friendly bear."""
|
||||
)
|
||||
|
||||
|
||||
def test_group_broken_paragraphs_with_bullets():
|
||||
text = """○The big red fox
|
||||
is walking down the lane.
|
||||
|
||||
○At the end of the lane
|
||||
the fox met a friendly bear."""
|
||||
assert core.group_bullet_paragraph(text) == [
|
||||
"○The big red fox is walking down the lane. ",
|
||||
"○At the end of the lane the fox met a friendly bear.",
|
||||
]
|
||||
|
||||
|
||||
def test_group_bullet_paragraph_with_e_bullets():
|
||||
text = """e The big red fox
|
||||
is walking down the lane.
|
||||
|
||||
e At the end of the lane
|
||||
the fox met a friendly bear."""
|
||||
assert core.group_bullet_paragraph(text) == [
|
||||
"· The big red fox is walking down the lane. ",
|
||||
"· At the end of the lane the fox met a friendly bear.",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
# NOTE(yuming): Tests combined cleaners
|
||||
(
|
||||
"text",
|
||||
"extra_whitespace",
|
||||
"dashes",
|
||||
"bullets",
|
||||
"lowercase",
|
||||
"trailing_punctuation",
|
||||
"expected",
|
||||
),
|
||||
[
|
||||
(" Risk-factors ", True, True, False, False, False, "Risk factors"),
|
||||
("● Point! ●●● ", True, False, True, False, False, "Point! ●●●"),
|
||||
("Risk- factors ", True, False, False, True, False, "risk- factors"),
|
||||
("Risk factors: ", True, False, False, False, True, "Risk factors"),
|
||||
("● Risk-factors●●● ", False, True, True, False, False, "Risk factors●●●"),
|
||||
("Risk-factors ", False, True, False, True, False, "risk factors"),
|
||||
("Risk-factors: ", False, True, False, False, True, "Risk factors"),
|
||||
("● Point! ●●● ", False, False, True, True, False, "point! ●●●"),
|
||||
("● Point! ●●●: ", False, False, True, False, True, "Point! ●●●"),
|
||||
("Risk factors: ", False, False, False, True, True, "risk factors"),
|
||||
],
|
||||
)
|
||||
def test_clean(text, extra_whitespace, dashes, bullets, lowercase, trailing_punctuation, expected):
|
||||
assert (
|
||||
core.clean(
|
||||
text=text,
|
||||
extra_whitespace=extra_whitespace,
|
||||
dashes=dashes,
|
||||
bullets=bullets,
|
||||
trailing_punctuation=trailing_punctuation,
|
||||
lowercase=lowercase,
|
||||
)
|
||||
== expected
|
||||
)
|
||||
|
||||
|
||||
def test_bytes_string_to_string():
|
||||
text = "\xe6\xaf\x8f\xe6\x97\xa5\xe6\x96\xb0\xe9\x97\xbb"
|
||||
assert core.bytes_string_to_string(text, "utf-8") == "每日新闻"
|
||||
@@ -0,0 +1,156 @@
|
||||
import datetime
|
||||
|
||||
import pytest
|
||||
|
||||
from unstructured.cleaners import extract
|
||||
|
||||
EMAIL_META_DATA_INPUT = """from ABC.DEF.local ([ba23::58b5:2236:45g2:88h2]) by
|
||||
\n ABC.DEF.local ([68.183.71.12]) with mapi id\
|
||||
n 32.88.5467.123; Fri, 26 Mar 2021 11:04:09 +1200"""
|
||||
|
||||
|
||||
def test_get_indexed_match_raises_with_bad_index():
|
||||
with pytest.raises(ValueError):
|
||||
extract._get_indexed_match("BLAH BLAH BLAH", "BLAH", -1)
|
||||
|
||||
|
||||
def test_get_indexed_match_raises_with_index_too_high():
|
||||
with pytest.raises(ValueError):
|
||||
extract._get_indexed_match("BLAH BLAH BLAH", "BLAH", 4)
|
||||
|
||||
|
||||
def test_extract_text_before():
|
||||
text = "Teacher: BLAH BLAH BLAH; Student: BLAH BLAH BLAH!"
|
||||
assert extract.extract_text_before(text, "BLAH", 1) == "Teacher: BLAH"
|
||||
|
||||
|
||||
def test_extract_text_after():
|
||||
text = "Teacher: BLAH BLAH BLAH; Student: BLAH BLAH BLAH!"
|
||||
assert extract.extract_text_after(text, "BLAH;", 0) == "Student: BLAH BLAH BLAH!"
|
||||
|
||||
|
||||
def test_extract_email_address():
|
||||
text = "Im Rabn <Im.Rabn@npf.gov.nr>"
|
||||
assert extract.extract_email_address(text) == ["im.rabn@npf.gov.nr"]
|
||||
|
||||
|
||||
def test_extract_ip_address():
|
||||
assert extract.extract_ip_address(EMAIL_META_DATA_INPUT) == [
|
||||
"ba23::58b5:2236:45g2:88h2",
|
||||
"68.183.71.12",
|
||||
]
|
||||
|
||||
|
||||
def test_extract_ip_address_name():
|
||||
assert extract.extract_ip_address_name(EMAIL_META_DATA_INPUT) == [
|
||||
"ABC.DEF.local",
|
||||
"ABC.DEF.local",
|
||||
]
|
||||
|
||||
|
||||
def test_extract_mapi_id():
|
||||
assert extract.extract_mapi_id(EMAIL_META_DATA_INPUT) == ["32.88.5467.123"]
|
||||
|
||||
|
||||
def test_extract_datetimetz():
|
||||
assert extract.extract_datetimetz(EMAIL_META_DATA_INPUT) == datetime.datetime(
|
||||
2021,
|
||||
3,
|
||||
26,
|
||||
11,
|
||||
4,
|
||||
9,
|
||||
tzinfo=datetime.timezone(datetime.timedelta(seconds=43200)),
|
||||
)
|
||||
|
||||
|
||||
def test_extract_datetimetz_works_with_no_date():
|
||||
assert extract.extract_datetimetz("NO DATE HERE") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("215-867-5309", "215-867-5309"),
|
||||
("Phone Number: +1 215.867.5309", "+1 215.867.5309"),
|
||||
("Phone Number: Just Kidding", ""),
|
||||
],
|
||||
)
|
||||
def test_extract_us_phone_number(text, expected):
|
||||
phone_number = extract.extract_us_phone_number(text)
|
||||
assert phone_number == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
("1. Introduction:", ("1", None, None)),
|
||||
("a. Introduction:", ("a", None, None)),
|
||||
("20.3 Morse code ●●●", ("20", "3", None)),
|
||||
("5.3.1 Convolutional Networks ", ("5", "3", "1")),
|
||||
("D.b.C Recurrent Neural Networks", ("D", "b", "C")),
|
||||
("2.b.1 Recurrent Neural Networks", ("2", "b", "1")),
|
||||
("eins. Neural Networks", (None, None, None)),
|
||||
("bb.c Feed Forward Neural Networks", ("bb", "c", None)),
|
||||
("aaa.ccc Metrics", (None, None, None)),
|
||||
(" version = 3.8", (None, None, None)),
|
||||
("1 2. 3 4", (None, None, None)),
|
||||
("1) 2. 3 4", (None, None, None)),
|
||||
("2,3. Morse code 3. ●●●", (None, None, None)),
|
||||
("1..2.3 four", (None, None, None)),
|
||||
("Fig. 2: The relationship", (None, None, None)),
|
||||
("23 is everywhere", (None, None, None)),
|
||||
],
|
||||
)
|
||||
def test_extract_ordered_bullets(text, expected):
|
||||
assert extract.extract_ordered_bullets(text=text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
(
|
||||
"https://my-image.jpg",
|
||||
(["https://my-image.jpg"]),
|
||||
),
|
||||
(
|
||||
"https://my-image.png with some text",
|
||||
(["https://my-image.png"]),
|
||||
),
|
||||
(
|
||||
"https://my-image/with/some/path.png",
|
||||
(["https://my-image/with/some/path.png"]),
|
||||
),
|
||||
(
|
||||
"some text https://my-image.jpg with another http://my-image.bmp",
|
||||
(["https://my-image.jpg", "http://my-image.bmp"]),
|
||||
),
|
||||
(
|
||||
"http://not-an-image.com",
|
||||
([]),
|
||||
),
|
||||
(
|
||||
"some text",
|
||||
([]),
|
||||
),
|
||||
(
|
||||
"some text https://my-image.JPG with another http://my-image.BMP",
|
||||
(["https://my-image.JPG", "http://my-image.BMP"]),
|
||||
),
|
||||
(
|
||||
"http://my-path-with-CAPS/my-image.JPG",
|
||||
(["http://my-path-with-CAPS/my-image.JPG"]),
|
||||
),
|
||||
(
|
||||
"http://my-path/my%20image.JPG",
|
||||
(["http://my-path/my%20image.JPG"]),
|
||||
),
|
||||
# url with reference #
|
||||
(
|
||||
"https://my-image.jpg#ref",
|
||||
(["https://my-image.jpg"]),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_extract_image_urls_from_html(text, expected):
|
||||
assert extract.extract_image_urls_from_html(text=text) == expected
|
||||
@@ -0,0 +1,64 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from unstructured.cleaners import translate
|
||||
|
||||
IS_CI = os.getenv("CI") == "true"
|
||||
|
||||
|
||||
def test_get_opus_mt_model_name():
|
||||
model_name = translate._get_opus_mt_model_name("ru", "en")
|
||||
assert model_name == "Helsinki-NLP/opus-mt-ru-en"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("code", ["way-too-long", "a", "", None])
|
||||
def test_validate_language_code(code):
|
||||
with pytest.raises(ValueError):
|
||||
translate._validate_language_code(code)
|
||||
|
||||
|
||||
def test_translate_returns_same_text_if_dest_is_same():
|
||||
text = "This is already in English!"
|
||||
assert translate.translate_text(text, "en", "en") == text
|
||||
|
||||
|
||||
def test_translate_returns_same_text_text_is_empty():
|
||||
text = " "
|
||||
assert translate.translate_text(text) == text
|
||||
|
||||
|
||||
@pytest.mark.skipif(IS_CI, reason="Skipping this test in CI pipeline")
|
||||
def test_translate_with_language_specified():
|
||||
text = "Ich bin ein Berliner!"
|
||||
assert translate.translate_text(text, "de") == "I'm a Berliner!"
|
||||
|
||||
|
||||
@pytest.mark.skipif(IS_CI, reason="Skipping this test in CI pipeline")
|
||||
def test_translate_with_no_language_specified():
|
||||
text = "Ich bin ein Berliner!"
|
||||
assert translate.translate_text(text) == "I'm a Berliner!"
|
||||
|
||||
|
||||
@pytest.mark.skipif(IS_CI, reason="Skipping this test in CI pipeline")
|
||||
def test_translate_raises_with_bad_language():
|
||||
text = "Ich bin ein Berliner!"
|
||||
with pytest.raises(ValueError):
|
||||
translate.translate_text(text, "zz")
|
||||
|
||||
|
||||
@pytest.mark.skipif(IS_CI, reason="Skipping this test in CI pipeline")
|
||||
def test_tranlate_works_with_russian():
|
||||
text = "Я тоже можно переводать русский язык!"
|
||||
assert translate.translate_text(text) == "I can also translate Russian!"
|
||||
|
||||
|
||||
@pytest.mark.skipif(IS_CI, reason="Skipping this test in CI pipeline")
|
||||
def test_translate_works_with_chinese():
|
||||
text = "網站有中、英文版本"
|
||||
translate.translate_text(text) == "Website available in Chinese and English"
|
||||
|
||||
|
||||
def translate_works_with_arabic():
|
||||
text = "مرحباً بكم في متجرنا"
|
||||
translate.translate_text(text) == "Welcome to our store."
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,203 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Unit-test suite for the `unstructured.common.html_table` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from lxml.html import fragment_fromstring
|
||||
|
||||
from unstructured.common.html_table import (
|
||||
HtmlCell,
|
||||
HtmlRow,
|
||||
HtmlTable,
|
||||
htmlify_matrix_of_cell_texts,
|
||||
)
|
||||
|
||||
|
||||
class Describe_htmlify_matrix_of_cell_texts:
|
||||
"""Unit-test suite for `unstructured.common.html_table.htmlify_matrix_of_cell_texts()`."""
|
||||
|
||||
def test_htmlify_matrix_handles_empty_cells(self):
|
||||
assert htmlify_matrix_of_cell_texts([["cell1", "", "cell3"], ["", "cell5", ""]]) == (
|
||||
"<table>"
|
||||
"<tr><td>cell1</td><td/><td>cell3</td></tr>"
|
||||
"<tr><td/><td>cell5</td><td/></tr>"
|
||||
"</table>"
|
||||
)
|
||||
|
||||
def test_htmlify_matrix_handles_special_characters(self):
|
||||
assert htmlify_matrix_of_cell_texts([['<>&"', "newline\n"]]) == (
|
||||
"<table><tr><td><>&"</td><td>newline<br/></td></tr></table>"
|
||||
)
|
||||
|
||||
def test_htmlify_matrix_handles_multiple_rows_and_cells(self):
|
||||
assert htmlify_matrix_of_cell_texts([["cell1", "cell2"], ["cell3", "cell4"]]) == (
|
||||
"<table>"
|
||||
"<tr><td>cell1</td><td>cell2</td></tr>"
|
||||
"<tr><td>cell3</td><td>cell4</td></tr>"
|
||||
"</table>"
|
||||
)
|
||||
|
||||
def test_htmlify_matrix_handles_empty_matrix(self):
|
||||
assert htmlify_matrix_of_cell_texts([]) == ""
|
||||
|
||||
|
||||
class DescribeHtmlTable:
|
||||
"""Unit-test suite for `unstructured.common.html_table.HtmlTable`."""
|
||||
|
||||
def it_can_construct_from_html_text(self):
|
||||
html_table = HtmlTable.from_html_text("<table><tr><td>foobar</td></tr></table>")
|
||||
|
||||
assert isinstance(html_table, HtmlTable)
|
||||
assert html_table._table.tag == "table"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"html_text",
|
||||
[
|
||||
"<table><tr><td>foobar</td></tr></table>",
|
||||
"<body><table><tr><td>foobar</td></tr></table></body>",
|
||||
"<html><body><table><tr><td>foobar</td></tr></table></body></html>",
|
||||
],
|
||||
)
|
||||
def it_can_find_a_table_wrapped_in_an_html_or_body_element(self, html_text: str):
|
||||
html_table = HtmlTable.from_html_text(html_text)
|
||||
|
||||
assert isinstance(html_table, HtmlTable)
|
||||
assert html_table._table.tag == "table"
|
||||
|
||||
def but_it_raises_when_no_table_element_is_present_in_the_html(self):
|
||||
with pytest.raises(ValueError, match="`html_text` contains no `<table>` element"):
|
||||
HtmlTable.from_html_text("<html><body><tr><td>foobar</td></tr></body></html>")
|
||||
|
||||
def it_removes_any_attributes_present_on_the_table_element(self):
|
||||
html_table = HtmlTable.from_html_text(
|
||||
'<table border="1", class="foobar"><tr><td>foobar</td></tr></table>',
|
||||
)
|
||||
assert html_table.html == "<table><tr><td>foobar</td></tr></table>"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"html_text",
|
||||
[
|
||||
"<table><thead><tr><td>foobar</td></tr></thead></table>",
|
||||
"<table><thead><tr><td>foobar</td></tr></thead><tbody></tbody></table>",
|
||||
"<table><tbody><tr><td>foobar</td></tr></tbody><tfoot></tfoot></table>",
|
||||
],
|
||||
)
|
||||
def it_removes_any_thead_tbody_or_tfoot_elements_present_within_the_table_element(
|
||||
self, html_text: str
|
||||
):
|
||||
html_table = HtmlTable.from_html_text(html_text)
|
||||
assert html_table.html == "<table><tr><td>foobar</td></tr></table>"
|
||||
|
||||
def it_changes_any_th_elements_to_td_elements_for_cell_element_uniformity(self):
|
||||
html_table = HtmlTable.from_html_text(
|
||||
"<table>"
|
||||
" <tr><th>a</th><th/><th>b</th></tr>"
|
||||
" <tr><td/><td>c</td><td/></tr>"
|
||||
"</table>"
|
||||
)
|
||||
assert html_table.html == (
|
||||
"<table><tr><td>a</td><td/><td>b</td></tr><tr><td/><td>c</td><td/></tr></table>"
|
||||
)
|
||||
|
||||
def it_removes_any_extra_whitespace_between_elements_and_normalizes_whitespace_in_text(self):
|
||||
html_table = HtmlTable.from_html_text(
|
||||
"\n <table>\n <tr>\n <td>\tabc def\nghi </td>\n </tr>\n</table>\n ",
|
||||
)
|
||||
assert html_table.html == "<table><tr><td>abc def ghi</td></tr></table>"
|
||||
|
||||
def it_can_serialize_the_table_element_to_str_html_text(self):
|
||||
table = fragment_fromstring("<table><tr><td>foobar</td></tr></table>")
|
||||
html_table = HtmlTable(table)
|
||||
|
||||
assert html_table.html == "<table><tr><td>foobar</td></tr></table>"
|
||||
|
||||
def it_can_iterate_the_rows_in_the_table(self):
|
||||
html_table = HtmlTable.from_html_text(
|
||||
"<table>"
|
||||
" <tr><td>abc</td><td>def</td><td>ghi</td></tr>"
|
||||
" <tr><td>jkl</td><td>mno</td><td>pqr</td></tr>"
|
||||
" <tr><td>stu</td><td>vwx</td><td>yz</td></tr>"
|
||||
"</table>"
|
||||
)
|
||||
|
||||
row_iter = html_table.iter_rows()
|
||||
|
||||
row = next(row_iter)
|
||||
assert isinstance(row, HtmlRow)
|
||||
assert row.html == "<tr><td>abc</td><td>def</td><td>ghi</td></tr>"
|
||||
# --
|
||||
row = next(row_iter)
|
||||
assert isinstance(row, HtmlRow)
|
||||
assert row.html == "<tr><td>jkl</td><td>mno</td><td>pqr</td></tr>"
|
||||
# --
|
||||
row = next(row_iter)
|
||||
assert isinstance(row, HtmlRow)
|
||||
assert row.html == "<tr><td>stu</td><td>vwx</td><td>yz</td></tr>"
|
||||
# --
|
||||
with pytest.raises(StopIteration):
|
||||
next(row_iter)
|
||||
|
||||
def it_provides_access_to_the_clear_concatenated_text_of_the_table(self):
|
||||
html_table = HtmlTable.from_html_text(
|
||||
"<table>"
|
||||
" <tr><th> a\n b c </th><th/><th>def</th></tr>"
|
||||
" <tr><td>gh \ti</td><td/><td>\n jk l </td></tr>"
|
||||
" <tr><td/><td> m n op\n</td><td/></tr>"
|
||||
"</table>"
|
||||
)
|
||||
assert html_table.text == "a b c def gh i jk l m n op"
|
||||
|
||||
|
||||
class DescribeHtmlRow:
|
||||
"""Unit-test suite for `unstructured.common.html_table.HtmlRow`."""
|
||||
|
||||
def it_can_serialize_the_row_to_html(self):
|
||||
assert HtmlRow(fragment_fromstring("<tr><td>a</td><td>b</td><td/></tr>")).html == (
|
||||
"<tr><td>a</td><td>b</td><td/></tr>"
|
||||
)
|
||||
|
||||
def it_can_iterate_the_cells_in_the_row(self):
|
||||
row = HtmlRow(fragment_fromstring("<tr><td>a</td><td>b</td><td/></tr>"))
|
||||
|
||||
cell_iter = row.iter_cells()
|
||||
|
||||
cell = next(cell_iter)
|
||||
assert isinstance(cell, HtmlCell)
|
||||
assert cell.html == "<td>a</td>"
|
||||
# --
|
||||
cell = next(cell_iter)
|
||||
assert isinstance(cell, HtmlCell)
|
||||
assert cell.html == "<td>b</td>"
|
||||
# --
|
||||
cell = next(cell_iter)
|
||||
assert isinstance(cell, HtmlCell)
|
||||
assert cell.html == "<td/>"
|
||||
# --
|
||||
with pytest.raises(StopIteration):
|
||||
next(cell_iter)
|
||||
|
||||
def it_can_iterate_the_texts_of_the_cells_in_the_row(self):
|
||||
row = HtmlRow(fragment_fromstring("<tr><td>a</td><td>b</td><td/></tr>"))
|
||||
|
||||
text_iter = row.iter_cell_texts()
|
||||
|
||||
assert next(text_iter) == "a"
|
||||
assert next(text_iter) == "b"
|
||||
with pytest.raises(StopIteration):
|
||||
next(text_iter)
|
||||
|
||||
|
||||
class DescribeHtmlCell:
|
||||
"""Unit-test suite for `unstructured.common.html_table.HtmlCell`."""
|
||||
|
||||
def it_can_serialize_the_cell_to_html(self):
|
||||
assert HtmlCell(fragment_fromstring("<td>a b c</td>")).html == "<td>a b c</td>"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("cell_html", "expected_value"),
|
||||
[("<td> Lorem ipsum </td>", "Lorem ipsum"), ("<td/>", "")],
|
||||
)
|
||||
def it_knows_the_text_in_the_cell(self, cell_html: str, expected_value: str):
|
||||
assert HtmlCell(fragment_fromstring(cell_html)).text == expected_value
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,73 @@
|
||||
import pytest
|
||||
|
||||
from unstructured.documents.coordinates import (
|
||||
CoordinateSystem,
|
||||
Orientation,
|
||||
RelativeCoordinateSystem,
|
||||
convert_coordinate,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("old_t", "old_t_max", "new_t_max", "t_orientation", "expected"),
|
||||
[(0, 7, 5, 1, 0), (7, 7, 5, 1, 5), (0, 7, 5, -1, 5), (7, 7, 5, -1, 0)],
|
||||
)
|
||||
def test_convert_coordinate(old_t, old_t_max, new_t_max, t_orientation, expected):
|
||||
assert convert_coordinate(old_t, old_t_max, new_t_max, t_orientation) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("width", "height", "orientation", "x", "y", "expected_x", "expected_y"),
|
||||
[
|
||||
(100, 300, Orientation.CARTESIAN, 0.8, 0.4, 80, 120),
|
||||
(100, 300, Orientation.SCREEN, 0.8, 0.6, 80, 120),
|
||||
],
|
||||
)
|
||||
def test_convert_from_relative(width, height, orientation, x, y, expected_x, expected_y):
|
||||
coord1 = CoordinateSystem(width, height)
|
||||
coord1.orientation = orientation
|
||||
assert coord1.convert_from_relative(x, y) == (expected_x, expected_y)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("width", "height", "orientation", "x", "y", "expected_x", "expected_y"),
|
||||
[
|
||||
(100, 300, Orientation.CARTESIAN, 80, 120, 0.8, 0.4),
|
||||
(100, 300, Orientation.SCREEN, 80, 120, 0.8, 0.6),
|
||||
],
|
||||
)
|
||||
def test_convert_to_relative(width, height, orientation, x, y, expected_x, expected_y):
|
||||
coord1 = CoordinateSystem(width, height)
|
||||
coord1.orientation = orientation
|
||||
assert coord1.convert_to_relative(x, y) == (expected_x, expected_y)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("orientation1", "orientation2", "x", "y", "expected_x", "expected_y"),
|
||||
[
|
||||
(Orientation.CARTESIAN, Orientation.CARTESIAN, 80, 120, 800, 1200),
|
||||
(Orientation.CARTESIAN, Orientation.SCREEN, 80, 120, 800, 800),
|
||||
(Orientation.SCREEN, Orientation.CARTESIAN, 80, 120, 800, 800),
|
||||
(Orientation.SCREEN, Orientation.SCREEN, 80, 120, 800, 1200),
|
||||
],
|
||||
)
|
||||
def test_convert_to_new_system(orientation1, orientation2, x, y, expected_x, expected_y):
|
||||
coord1 = CoordinateSystem(width=100, height=200)
|
||||
coord1.orientation = orientation1
|
||||
coord2 = CoordinateSystem(width=1000, height=2000)
|
||||
coord2.orientation = orientation2
|
||||
assert coord1.convert_coordinates_to_new_system(coord2, x, y) == (expected_x, expected_y)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("width", "height", "orientation", "x", "y", "expected_x", "expected_y"),
|
||||
[
|
||||
(100, 300, Orientation.CARTESIAN, 80, 120, 0.8, 0.4),
|
||||
(100, 300, Orientation.SCREEN, 80, 120, 0.8, 0.6),
|
||||
],
|
||||
)
|
||||
def test_relative_system(width, height, orientation, x, y, expected_x, expected_y):
|
||||
coord1 = CoordinateSystem(width, height)
|
||||
coord1.orientation = orientation
|
||||
coord2 = RelativeCoordinateSystem()
|
||||
assert coord1.convert_coordinates_to_new_system(coord2, x, y) == (expected_x, expected_y)
|
||||
@@ -0,0 +1,756 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Test-suite for `unstructured.documents.elements` module."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import io
|
||||
import json
|
||||
import pathlib
|
||||
from functools import partial
|
||||
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import assign_hash_ids, example_doc_path
|
||||
from unstructured.cleaners.core import clean_bullets, clean_prefix
|
||||
from unstructured.documents.coordinates import (
|
||||
CoordinateSystem,
|
||||
Orientation,
|
||||
RelativeCoordinateSystem,
|
||||
)
|
||||
from unstructured.documents.elements import (
|
||||
CheckBox,
|
||||
ConsolidationStrategy,
|
||||
CoordinatesMetadata,
|
||||
DataSourceMetadata,
|
||||
Element,
|
||||
ElementMetadata,
|
||||
Points,
|
||||
Text,
|
||||
Title,
|
||||
assign_and_map_hash_ids,
|
||||
)
|
||||
from unstructured.partition.json import partition_json
|
||||
|
||||
|
||||
@pytest.mark.parametrize("element", [Element(), Text(text=""), CheckBox()])
|
||||
def test_Element_autoassigns_a_UUID_then_becomes_an_idempotent_and_deterministic_hash(
|
||||
element: Element,
|
||||
):
|
||||
# -- element self-assigns itself a UUID --
|
||||
assert isinstance(element.id, str)
|
||||
assert len(element.id) == 36
|
||||
assert element.id.count("-") == 4
|
||||
|
||||
expected_hash = "5336294a19f32ff03ef80066fbc3e0f7"
|
||||
# -- calling `.id_to_hash()` changes the element's id-type to hash --
|
||||
assert element.id_to_hash(0) == expected_hash
|
||||
assert element.id == expected_hash
|
||||
|
||||
# -- `.id_to_hash()` is idempotent --
|
||||
assert element.id_to_hash(0) == expected_hash
|
||||
assert element.id == expected_hash
|
||||
|
||||
|
||||
def test_Text_is_JSON_serializable():
|
||||
# -- This shold run without an error --
|
||||
json.dumps(Text(text="hello there!", element_id=None).to_dict())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"element",
|
||||
[
|
||||
Element(),
|
||||
Text(text=""), # -- element_id should be implicitly None --
|
||||
Text(text="", element_id=None), # -- setting explicitly to None --
|
||||
CheckBox(),
|
||||
],
|
||||
)
|
||||
def test_Element_self_assigns_itself_a_UUID_id(element: Element):
|
||||
assert isinstance(element.id, str)
|
||||
assert len(element.id) == 36
|
||||
assert element.id.count("-") == 4
|
||||
|
||||
|
||||
def test_text_element_apply_cleaners():
|
||||
text_element = Text(text="[1] A Textbook on Crocodile Habitats")
|
||||
|
||||
text_element.apply(partial(clean_prefix, pattern=r"\[\d{1,2}\]"))
|
||||
assert str(text_element) == "A Textbook on Crocodile Habitats"
|
||||
|
||||
|
||||
def test_text_element_apply_multiple_cleaners():
|
||||
cleaners = [partial(clean_prefix, pattern=r"\[\d{1,2}\]"), partial(clean_bullets)]
|
||||
text_element = Text(text="[1] \u2022 A Textbook on Crocodile Habitats")
|
||||
text_element.apply(*cleaners)
|
||||
assert str(text_element) == "A Textbook on Crocodile Habitats"
|
||||
|
||||
|
||||
def test_non_text_elements_are_serializable_to_text():
|
||||
element = CheckBox()
|
||||
assert hasattr(element, "text")
|
||||
assert element.text is not None
|
||||
assert element.text == ""
|
||||
assert str(element) == ""
|
||||
|
||||
|
||||
def test_apply_raises_if_func_does_not_produce_string():
|
||||
def bad_cleaner(s: str):
|
||||
return 1
|
||||
|
||||
text_element = Text(text="[1] A Textbook on Crocodile Habitats")
|
||||
|
||||
with pytest.raises(ValueError, match="Cleaner produced a non-string output."):
|
||||
text_element.apply(bad_cleaner) # pyright: ignore[reportArgumentType]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("coordinates", "orientation1", "orientation2", "expected_coords"),
|
||||
[
|
||||
(
|
||||
((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
Orientation.CARTESIAN,
|
||||
Orientation.CARTESIAN,
|
||||
((10, 20), (10, 40), (30, 40), (30, 20)),
|
||||
),
|
||||
(
|
||||
((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
Orientation.CARTESIAN,
|
||||
Orientation.SCREEN,
|
||||
((10, 1980), (10, 1960), (30, 1960), (30, 1980)),
|
||||
),
|
||||
(
|
||||
((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
Orientation.SCREEN,
|
||||
Orientation.CARTESIAN,
|
||||
((10, 1980), (10, 1960), (30, 1960), (30, 1980)),
|
||||
),
|
||||
(
|
||||
((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
Orientation.SCREEN,
|
||||
Orientation.SCREEN,
|
||||
((10, 20), (10, 40), (30, 40), (30, 20)),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_convert_coordinates_to_new_system(
|
||||
coordinates: Points,
|
||||
orientation1: Orientation,
|
||||
orientation2: Orientation,
|
||||
expected_coords: Points,
|
||||
):
|
||||
coord1 = CoordinateSystem(100, 200)
|
||||
coord1.orientation = orientation1
|
||||
coord2 = CoordinateSystem(1000, 2000)
|
||||
coord2.orientation = orientation2
|
||||
element = Element(coordinates=coordinates, coordinate_system=coord1)
|
||||
|
||||
new_coords = element.convert_coordinates_to_new_system(coord2)
|
||||
|
||||
assert new_coords is not None
|
||||
for new_coord, expected in zip(new_coords, expected_coords):
|
||||
assert new_coord == pytest.approx(expected) # pyright: ignore[reportUnknownMemberType]
|
||||
element.convert_coordinates_to_new_system(coord2, in_place=True)
|
||||
assert element.metadata.coordinates is not None
|
||||
assert element.metadata.coordinates.points is not None
|
||||
for new_coord, expected in zip(element.metadata.coordinates.points, expected_coords):
|
||||
assert new_coord == pytest.approx(expected) # pyright: ignore[reportUnknownMemberType]
|
||||
assert element.metadata.coordinates.system == coord2
|
||||
|
||||
|
||||
def test_convert_coordinate_to_new_system_none():
|
||||
element = Element(coordinates=None, coordinate_system=None)
|
||||
coord = CoordinateSystem(100, 200)
|
||||
coord.orientation = Orientation.SCREEN
|
||||
assert element.convert_coordinates_to_new_system(coord) is None
|
||||
|
||||
|
||||
def test_element_constructor_coordinates_all_present():
|
||||
coordinates = ((1, 2), (1, 4), (3, 4), (3, 2))
|
||||
coordinate_system = RelativeCoordinateSystem()
|
||||
element = Element(coordinates=coordinates, coordinate_system=coordinate_system)
|
||||
expected_coordinates_metadata = CoordinatesMetadata(
|
||||
points=coordinates,
|
||||
system=coordinate_system,
|
||||
)
|
||||
assert element.metadata.coordinates == expected_coordinates_metadata
|
||||
|
||||
|
||||
def test_element_constructor_coordinates_points_absent():
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
Element(coordinate_system=RelativeCoordinateSystem())
|
||||
assert (
|
||||
str(exc_info.value)
|
||||
== "Coordinates points should not exist without coordinates system and vice versa."
|
||||
)
|
||||
|
||||
|
||||
def test_element_constructor_coordinates_system_absent():
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
Element(coordinates=((1, 2), (1, 4), (3, 4), (3, 2)))
|
||||
assert (
|
||||
str(exc_info.value)
|
||||
== "Coordinates points should not exist without coordinates system and vice versa."
|
||||
)
|
||||
|
||||
|
||||
def test_coordinate_metadata_serdes():
|
||||
coordinates = ((1, 2), (1, 4), (3, 4), (3, 2))
|
||||
coordinate_system = RelativeCoordinateSystem()
|
||||
coordinates_metadata = CoordinatesMetadata(points=coordinates, system=coordinate_system)
|
||||
expected_schema = {
|
||||
"layout_height": 1,
|
||||
"layout_width": 1,
|
||||
"points": ((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
"system": "RelativeCoordinateSystem",
|
||||
}
|
||||
coordinates_metadata_dict = coordinates_metadata.to_dict()
|
||||
assert coordinates_metadata_dict == expected_schema
|
||||
assert CoordinatesMetadata.from_dict(coordinates_metadata_dict) == coordinates_metadata
|
||||
|
||||
|
||||
def test_element_to_dict():
|
||||
coordinates = ((1, 2), (1, 4), (3, 4), (3, 2))
|
||||
coordinate_system = RelativeCoordinateSystem()
|
||||
element = Element(
|
||||
element_id="awt32t1",
|
||||
coordinates=coordinates,
|
||||
coordinate_system=coordinate_system,
|
||||
)
|
||||
|
||||
assert element.to_dict() == {
|
||||
"metadata": {
|
||||
"coordinates": {
|
||||
"layout_height": 1,
|
||||
"layout_width": 1,
|
||||
"points": ((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
"system": "RelativeCoordinateSystem",
|
||||
},
|
||||
},
|
||||
"type": None,
|
||||
"text": "",
|
||||
"element_id": "awt32t1",
|
||||
}
|
||||
|
||||
|
||||
class DescribeElementMetadata:
|
||||
"""Unit-test suite for `unstructured.documents.elements.ElementMetadata`."""
|
||||
|
||||
# -- It can be constructed with known keyword arguments. In particular, including a non-known
|
||||
# -- keyword argument produces a type-error at development time and raises an exception at
|
||||
# -- runtime. This catches typos before they reach production.
|
||||
|
||||
def it_detects_unknown_constructor_args_at_both_development_time_and_runtime(self):
|
||||
with pytest.raises(TypeError, match="got an unexpected keyword argument 'file_name'"):
|
||||
ElementMetadata(file_name="memo.docx") # pyright: ignore[reportCallIssue]
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_path",
|
||||
[
|
||||
pathlib.Path("documents/docx") / "memos" / "memo-2023-11-10.docx",
|
||||
"documents/docx/memos/memo-2023-11-10.docx",
|
||||
],
|
||||
)
|
||||
def it_accommodates_either_a_pathlib_Path_or_str_for_its_filename_arg(
|
||||
self, file_path: pathlib.Path | str
|
||||
):
|
||||
meta = ElementMetadata(filename=file_path)
|
||||
|
||||
assert meta.file_directory == "documents/docx/memos"
|
||||
assert meta.filename == "memo-2023-11-10.docx"
|
||||
|
||||
def it_leaves_both_filename_and_file_directory_None_when_neither_is_specified(self):
|
||||
meta = ElementMetadata()
|
||||
|
||||
assert meta.file_directory is None
|
||||
assert meta.filename is None
|
||||
|
||||
@pytest.mark.parametrize("file_path", [pathlib.Path("memo.docx"), "memo.docx"])
|
||||
def and_it_leaves_file_directory_None_when_not_specified_and_filename_is_not_a_path(
|
||||
self, file_path: pathlib.Path | str
|
||||
):
|
||||
meta = ElementMetadata(filename=file_path)
|
||||
|
||||
assert meta.file_directory is None
|
||||
assert meta.filename == "memo.docx"
|
||||
|
||||
def and_it_splits_off_directory_path_from_its_filename_arg_when_it_is_a_file_path(self):
|
||||
meta = ElementMetadata(filename="documents/docx/memo-2023-11-11.docx")
|
||||
|
||||
assert meta.file_directory == "documents/docx"
|
||||
assert meta.filename == "memo-2023-11-11.docx"
|
||||
|
||||
def but_it_prefers_a_specified_file_directory_when_filename_also_contains_a_path(self):
|
||||
meta = ElementMetadata(filename="tmp/staging/memo.docx", file_directory="documents/docx")
|
||||
|
||||
assert meta.file_directory == "documents/docx"
|
||||
assert meta.filename == "memo.docx"
|
||||
|
||||
# -- It knows the types of its known members so type-checking support is available. --
|
||||
|
||||
def it_knows_the_types_of_its_known_members_so_type_checking_support_is_available(self):
|
||||
ElementMetadata(
|
||||
category_depth="2", # pyright: ignore[reportArgumentType]
|
||||
file_directory=True, # pyright: ignore[reportArgumentType]
|
||||
text_as_html=42, # pyright: ignore[reportArgumentType]
|
||||
)
|
||||
# -- it does not check types at runtime however (choosing to avoid validation overhead) --
|
||||
|
||||
# -- It only stores a field's value when it is not None. --
|
||||
|
||||
def it_returns_the_value_of_an_attribute_it_has(self):
|
||||
meta = ElementMetadata(url="https://google.com")
|
||||
assert "url" in meta.__dict__
|
||||
assert meta.url == "https://google.com"
|
||||
|
||||
def and_it_returns_None_for_a_known_attribute_it_does_not_have(self):
|
||||
meta = ElementMetadata()
|
||||
assert "url" not in meta.__dict__
|
||||
assert meta.url is None
|
||||
|
||||
def but_it_raises_AttributeError_for_an_unknown_attribute_it_does_not_have(self):
|
||||
meta = ElementMetadata()
|
||||
assert "coefficient" not in meta.__dict__
|
||||
with pytest.raises(AttributeError, match="object has no attribute 'coefficient'"):
|
||||
meta.coefficient
|
||||
|
||||
def it_stores_a_non_None_field_value_when_assigned(self):
|
||||
meta = ElementMetadata()
|
||||
assert "file_directory" not in meta.__dict__
|
||||
meta.file_directory = "tmp/"
|
||||
assert "file_directory" in meta.__dict__
|
||||
assert meta.file_directory == "tmp/"
|
||||
|
||||
def it_removes_a_field_when_None_is_assigned_to_it(self):
|
||||
meta = ElementMetadata(file_directory="tmp/")
|
||||
assert "file_directory" in meta.__dict__
|
||||
assert meta.file_directory == "tmp/"
|
||||
|
||||
meta.file_directory = None
|
||||
assert "file_directory" not in meta.__dict__
|
||||
assert meta.file_directory is None
|
||||
|
||||
# -- It can serialize itself to a dict -------------------------------------------------------
|
||||
|
||||
def it_can_serialize_itself_to_a_dict(self):
|
||||
meta = ElementMetadata(
|
||||
category_depth=1,
|
||||
file_directory="tmp/",
|
||||
page_number=2,
|
||||
text_as_html="<table></table>",
|
||||
url="https://google.com",
|
||||
)
|
||||
assert meta.to_dict() == {
|
||||
"category_depth": 1,
|
||||
"file_directory": "tmp/",
|
||||
"page_number": 2,
|
||||
"text_as_html": "<table></table>",
|
||||
"url": "https://google.com",
|
||||
}
|
||||
|
||||
def and_it_serializes_a_coordinates_sub_object_to_a_dict_when_it_is_present(self):
|
||||
meta = ElementMetadata(
|
||||
category_depth=1,
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=((2, 2), (1, 4), (3, 4), (3, 2)),
|
||||
system=RelativeCoordinateSystem(),
|
||||
),
|
||||
page_number=2,
|
||||
)
|
||||
assert meta.to_dict() == {
|
||||
"category_depth": 1,
|
||||
"coordinates": {
|
||||
"layout_height": 1,
|
||||
"layout_width": 1,
|
||||
"points": ((2, 2), (1, 4), (3, 4), (3, 2)),
|
||||
"system": "RelativeCoordinateSystem",
|
||||
},
|
||||
"page_number": 2,
|
||||
}
|
||||
|
||||
def and_it_serializes_a_data_source_sub_object_to_a_dict_when_it_is_present(self):
|
||||
meta = ElementMetadata(
|
||||
category_depth=1,
|
||||
data_source=DataSourceMetadata(
|
||||
url="https://www.nih.gov/about-nih/who-we-are/nih-director",
|
||||
date_created="2023-11-09",
|
||||
),
|
||||
page_number=2,
|
||||
)
|
||||
assert meta.to_dict() == {
|
||||
"category_depth": 1,
|
||||
"data_source": {
|
||||
"url": "https://www.nih.gov/about-nih/who-we-are/nih-director",
|
||||
"date_created": "2023-11-09",
|
||||
},
|
||||
"page_number": 2,
|
||||
}
|
||||
|
||||
def and_it_serializes_an_orig_elements_sub_object_to_base64_when_it_is_present(self):
|
||||
elements = assign_hash_ids([Title("Lorem"), Text("Lorem Ipsum")])
|
||||
meta = ElementMetadata(
|
||||
category_depth=1,
|
||||
orig_elements=elements,
|
||||
page_number=2,
|
||||
)
|
||||
|
||||
assert meta.to_dict() == {
|
||||
"category_depth": 1,
|
||||
"orig_elements": (
|
||||
"eJyFzcsKwjAQheFXKVm7MGkzbXwDocu6EpFcTqTQG3UEtfTdbZa"
|
||||
"6cTnDd/jPi0CHHgNf2yAOmXCljjqXoErKoIw3hqJRXlPuyphrEr"
|
||||
"tM9GAbLNvNL+t2M56ctvU4o0+AXxPSo2m5g9jIb6VwBE0VBSujp"
|
||||
"1LJ6EiRLpwiSBf3fyvZcbo/vlqnwVvGbZzbN0KT7Hr5AG/eQyM="
|
||||
),
|
||||
"page_number": 2,
|
||||
}
|
||||
|
||||
def but_unlike_in_ElementMetadata_unknown_fields_in_sub_objects_are_ignored(self):
|
||||
"""Metadata sub-objects ignore fields they do not explicitly define.
|
||||
|
||||
This is _not_ the case for ElementMetadata itself where an non-known field is welcomed as a
|
||||
user-defined ad-hoc metadata field.
|
||||
"""
|
||||
element_metadata = {
|
||||
"new_field": "hello",
|
||||
"data_source": {
|
||||
"new_field": "world",
|
||||
},
|
||||
"coordinates": {
|
||||
"new_field": "foo",
|
||||
},
|
||||
}
|
||||
|
||||
metadata = ElementMetadata.from_dict(element_metadata)
|
||||
metadata_dict = metadata.to_dict()
|
||||
|
||||
assert "new_field" in metadata_dict
|
||||
assert "new_field" not in metadata_dict["coordinates"]
|
||||
assert "new_field" not in metadata_dict["data_source"]
|
||||
|
||||
# -- It can deserialize itself from a dict ---------------------------------------------------
|
||||
|
||||
def it_can_deserialize_itself_from_a_dict(self):
|
||||
meta_dict = {
|
||||
"category_depth": 1,
|
||||
"coefficient": 0.58,
|
||||
"coordinates": {
|
||||
"layout_height": 4,
|
||||
"layout_width": 2,
|
||||
"points": ((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
"system": "RelativeCoordinateSystem",
|
||||
},
|
||||
"data_source": {
|
||||
"url": "https://www.nih.gov/about-nih/who-we-are/nih-director",
|
||||
"date_created": "2023-11-09",
|
||||
},
|
||||
"languages": ["eng"],
|
||||
}
|
||||
|
||||
meta = ElementMetadata.from_dict(meta_dict)
|
||||
|
||||
# -- known fields present in dict are present in meta --
|
||||
assert meta.category_depth == 1
|
||||
|
||||
# -- known sub-object fields present in dict are present in meta --
|
||||
assert meta.coordinates == CoordinatesMetadata(
|
||||
points=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
system=RelativeCoordinateSystem(),
|
||||
)
|
||||
assert meta.data_source == DataSourceMetadata(
|
||||
url="https://www.nih.gov/about-nih/who-we-are/nih-director",
|
||||
date_created="2023-11-09",
|
||||
)
|
||||
|
||||
# -- known fields absent from dict report None but are not present in meta --
|
||||
assert meta.file_directory is None
|
||||
assert "file_directory" not in meta.__dict__
|
||||
|
||||
# -- non-known fields present in dict are present in meta (we have no way to tell whether
|
||||
# -- they are "ad-hoc" or not because we lack indication of user-intent)
|
||||
assert meta.coefficient == 0.58
|
||||
|
||||
# -- ad-hoc fields absent from dict raise on attempted access --
|
||||
with pytest.raises(AttributeError, match="ntMetadata' object has no attribute 'quotient'"):
|
||||
meta.quotient
|
||||
|
||||
# -- but that can be worked around by end-user --
|
||||
assert (meta.quotient if hasattr(meta, "quotient") else None) is None
|
||||
|
||||
# -- mutating a mutable (collection) field does not affect the original value --
|
||||
assert isinstance(meta.languages, list)
|
||||
assert meta.languages == ["eng"]
|
||||
meta.languages.append("spa")
|
||||
assert meta.languages == ["eng", "spa"]
|
||||
assert meta_dict["languages"] == ["eng"]
|
||||
|
||||
# -- It allows downstream users to add an arbitrary new member by assignment. ----------------
|
||||
|
||||
def it_allows_an_end_user_to_add_an_arbitrary_field(self):
|
||||
meta = ElementMetadata()
|
||||
meta.foobar = 7
|
||||
assert "foobar" in meta.__dict__
|
||||
assert meta.foobar == 7
|
||||
|
||||
def and_fields_so_added_appear_in_the_metadata_JSON(self):
|
||||
meta = ElementMetadata()
|
||||
meta.foobar = 7
|
||||
assert meta.to_dict() == {"foobar": 7}
|
||||
|
||||
def and_it_removes_an_end_user_field_when_it_is_assigned_None(self):
|
||||
meta = ElementMetadata()
|
||||
meta.foobar = 7
|
||||
assert "foobar" in meta.__dict__
|
||||
meta.foobar = None
|
||||
assert "foobar" not in meta.__dict__
|
||||
with pytest.raises(
|
||||
AttributeError, match="'ElementMetadata' object has no attribute 'foobar'"
|
||||
):
|
||||
meta.foobar
|
||||
|
||||
# -- It can update itself from another instance ----------------------------------------------
|
||||
|
||||
def it_can_update_itself_from_another_instance(self):
|
||||
meta = ElementMetadata(category_depth=1, page_number=1)
|
||||
meta.coefficient = 0.58
|
||||
meta.stem_length = 18
|
||||
other = ElementMetadata(file_directory="tmp/", page_number=2)
|
||||
other.quotient = 1.4
|
||||
other.stem_length = 20
|
||||
|
||||
meta.update(other)
|
||||
|
||||
# -- known-fields present on self but not other are unchanged --
|
||||
assert meta.category_depth == 1
|
||||
# -- known-fields present on other but not self are added --
|
||||
assert meta.file_directory == "tmp/"
|
||||
# -- known-fields present on both self and other are updated --
|
||||
assert meta.page_number == 2
|
||||
# -- ad-hoc-fields present on self but not other are unchanged --
|
||||
assert meta.coefficient == 0.58
|
||||
# -- ad-hoc-fields present on other but not self are added --
|
||||
assert meta.quotient == 1.4
|
||||
# -- ad-hoc-fields present on both self and other are updated --
|
||||
assert meta.stem_length == 20
|
||||
# -- other is left unchanged --
|
||||
assert other.category_depth is None
|
||||
assert other.file_directory == "tmp/"
|
||||
assert other.page_number == 2
|
||||
assert other.text_as_html is None
|
||||
assert other.url is None
|
||||
assert other.quotient == 1.4
|
||||
assert other.stem_length == 20
|
||||
with pytest.raises(AttributeError, match="etadata' object has no attribute 'coefficient'"):
|
||||
other.coefficient
|
||||
|
||||
def but_it_raises_on_attempt_to_update_from_a_non_ElementMetadata_object(self):
|
||||
meta = ElementMetadata()
|
||||
with pytest.raises(ValueError, match=r"ate\(\)' must be an instance of 'ElementMetadata'"):
|
||||
meta.update({"coefficient": "0.56"}) # pyright: ignore[reportArgumentType]
|
||||
|
||||
# -- It knows when it is equal to another instance -------------------------------------------
|
||||
|
||||
def it_is_equal_to_another_instance_with_the_same_known_field_values(self):
|
||||
meta = ElementMetadata(
|
||||
category_depth=1,
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
system=RelativeCoordinateSystem(),
|
||||
),
|
||||
data_source=DataSourceMetadata(
|
||||
url="https://www.nih.gov/about-nih/who-we-are/nih-director",
|
||||
date_created="2023-11-08",
|
||||
),
|
||||
file_directory="tmp/",
|
||||
languages=["eng"],
|
||||
page_number=2,
|
||||
text_as_html="<table></table>",
|
||||
url="https://google.com",
|
||||
)
|
||||
assert meta == ElementMetadata(
|
||||
category_depth=1,
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
system=RelativeCoordinateSystem(),
|
||||
),
|
||||
data_source=DataSourceMetadata(
|
||||
url="https://www.nih.gov/about-nih/who-we-are/nih-director",
|
||||
date_created="2023-11-08",
|
||||
),
|
||||
file_directory="tmp/",
|
||||
languages=["eng"],
|
||||
page_number=2,
|
||||
text_as_html="<table></table>",
|
||||
url="https://google.com",
|
||||
)
|
||||
|
||||
def but_it_is_never_equal_to_a_non_ElementMetadata_object(self):
|
||||
class NotElementMetadata:
|
||||
pass
|
||||
|
||||
meta = ElementMetadata()
|
||||
other = NotElementMetadata()
|
||||
|
||||
# -- all the "fields" are the same --
|
||||
assert meta.__dict__ == other.__dict__
|
||||
# -- but it is rejected solely because its type is different --
|
||||
assert meta != other
|
||||
|
||||
def it_is_equal_to_another_instance_with_the_same_ad_hoc_field_values(self):
|
||||
meta = ElementMetadata(category_depth=1)
|
||||
meta.coefficient = 0.58
|
||||
other = ElementMetadata(category_depth=1)
|
||||
other.coefficient = 0.58
|
||||
|
||||
assert meta == other
|
||||
|
||||
def but_it_is_not_equal_to_an_instance_with_ad_hoc_fields_that_differ(self):
|
||||
meta = ElementMetadata(category_depth=1)
|
||||
meta.coefficient = 0.58
|
||||
other = ElementMetadata(category_depth=1)
|
||||
other.coefficient = 0.72
|
||||
|
||||
assert meta != other
|
||||
|
||||
def it_is_not_equal_when_a_list_field_contains_different_items(self):
|
||||
meta = ElementMetadata(languages=["eng"])
|
||||
assert meta != ElementMetadata(languages=["eng", "spa"])
|
||||
|
||||
def and_it_is_not_equal_when_the_coordinates_sub_object_field_differs(self):
|
||||
meta = ElementMetadata(
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=((1, 2), (1, 4), (3, 4), (3, 2)),
|
||||
system=RelativeCoordinateSystem(),
|
||||
)
|
||||
)
|
||||
assert meta != ElementMetadata(
|
||||
coordinates=CoordinatesMetadata(
|
||||
points=((2, 2), (2, 4), (3, 4), (4, 2)),
|
||||
system=RelativeCoordinateSystem(),
|
||||
)
|
||||
)
|
||||
|
||||
def and_it_is_not_equal_when_the_data_source_sub_object_field_differs(self):
|
||||
meta = ElementMetadata(
|
||||
data_source=DataSourceMetadata(
|
||||
url="https://www.nih.gov/about-nih/who-we-are/nih-director",
|
||||
date_created="2023-11-08",
|
||||
)
|
||||
)
|
||||
assert meta != ElementMetadata(
|
||||
data_source=DataSourceMetadata(
|
||||
url="https://www.nih.gov/about-nih/who-we-are/nih-director",
|
||||
date_created="2023-11-09",
|
||||
)
|
||||
)
|
||||
|
||||
# -- There is a consolidation-strategy for all known fields ----------------------------------
|
||||
|
||||
def it_can_find_the_consolidation_strategy_for_each_of_its_known_fields(self):
|
||||
metadata = ElementMetadata()
|
||||
metadata_field_names = sorted(metadata._known_field_names)
|
||||
consolidation_strategies = ConsolidationStrategy.field_consolidation_strategies()
|
||||
|
||||
for field_name in metadata_field_names:
|
||||
assert field_name in consolidation_strategies, (
|
||||
f"ElementMetadata field `.{field_name}` does not have a consolidation strategy."
|
||||
f" Add one in `ConsolidationStrategy.field_consolidation_strategies()."
|
||||
)
|
||||
|
||||
|
||||
def test_hash_ids_are_unique_for_duplicate_elements():
|
||||
# GIVEN
|
||||
parent = Text(text="Parent", metadata=ElementMetadata(page_number=1))
|
||||
elements: list[Element] = [
|
||||
parent,
|
||||
Text(text="Element", metadata=ElementMetadata(page_number=1, parent_id=parent.id)),
|
||||
Text(text="Element", metadata=ElementMetadata(page_number=1, parent_id=parent.id)),
|
||||
]
|
||||
|
||||
# WHEN
|
||||
updated_elements = assign_and_map_hash_ids(copy.deepcopy(elements))
|
||||
ids = [element.id for element in updated_elements]
|
||||
|
||||
# THEN
|
||||
assert len(ids) == len(set(ids)), "Recalculated IDs must be unique."
|
||||
assert elements[1].metadata.parent_id == elements[2].metadata.parent_id
|
||||
|
||||
for idx, updated_element in enumerate(updated_elements):
|
||||
assert updated_element.id != elements[idx].id, "IDs haven't changed after recalculation"
|
||||
if updated_element.metadata.parent_id is not None:
|
||||
assert updated_element.metadata.parent_id in ids, "Parent ID not in the list of IDs"
|
||||
assert (
|
||||
updated_element.metadata.parent_id != elements[idx].metadata.parent_id
|
||||
), "Parent ID hasn't changed after recalculation"
|
||||
|
||||
|
||||
def test_hash_ids_can_handle_duplicated_element_instances():
|
||||
# GIVEN
|
||||
parent = Text(text="Parent", metadata=ElementMetadata(page_number=1))
|
||||
element = Text(text="Element", metadata=ElementMetadata(page_number=1, parent_id=parent.id))
|
||||
elements: list[Element] = [parent, element, element]
|
||||
|
||||
# WHEN
|
||||
updated_elements = assign_and_map_hash_ids(copy.deepcopy(elements))
|
||||
ids = [element.id for element in updated_elements]
|
||||
|
||||
# THEN
|
||||
assert len(ids) == len(set(ids)) + 1, "One element is duplicated so uniques should be one less."
|
||||
assert elements[1].metadata.parent_id == elements[2].metadata.parent_id
|
||||
|
||||
|
||||
def test_hash_ids_are_deterministic():
|
||||
parent = Text(text="Parent", metadata=ElementMetadata(page_number=1))
|
||||
elements: list[Element] = [
|
||||
parent,
|
||||
Text(text="Element", metadata=ElementMetadata(page_number=1, parent_id=parent.id)),
|
||||
Text(text="Element", metadata=ElementMetadata(page_number=1, parent_id=parent.id)),
|
||||
]
|
||||
|
||||
updated_elements = assign_and_map_hash_ids(elements)
|
||||
ids = [element.id for element in updated_elements]
|
||||
parent_ids = [element.metadata.parent_id for element in updated_elements]
|
||||
|
||||
assert ids == [
|
||||
"ea9eb7e80383c190f8cafce1ad666624",
|
||||
"4112a8d24886276e18e759d06956021b",
|
||||
"eba84bbe7f03e8b91a1527323040ee3d",
|
||||
]
|
||||
assert parent_ids == [
|
||||
None,
|
||||
"ea9eb7e80383c190f8cafce1ad666624",
|
||||
"ea9eb7e80383c190f8cafce1ad666624",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "sequence_number", "filename", "page_number", "expected_hash"),
|
||||
[
|
||||
# -- pdf files support page numbers --
|
||||
("foo", 1, "foo.pdf", 1, "4bb264eb23ceb44cd8fcc5af44f8dc71"),
|
||||
("foo", 2, "foo.pdf", 1, "75fc1de48cf724ec00aa8d1c5a0d3758"),
|
||||
# -- txt files don't have a page number --
|
||||
("some text", 0, "some.txt", None, "1a2627b5760c06b1440102f11a1edb0f"),
|
||||
("some text", 1, "some.txt", None, "e3fd10d867c4a1c0264dde40e3d7e45a"),
|
||||
],
|
||||
)
|
||||
def test_id_to_hash_calculates(
|
||||
text: str, sequence_number: int, filename: str, page_number: int | None, expected_hash: str
|
||||
):
|
||||
element = Text(
|
||||
text=text,
|
||||
metadata=ElementMetadata(filename=filename, page_number=page_number),
|
||||
)
|
||||
assert element.id_to_hash(sequence_number) == expected_hash, "Returned ID does not match"
|
||||
assert element.id == expected_hash, "ID should be set"
|
||||
|
||||
|
||||
def test_formskeysvalues_reads_saves():
|
||||
filename = example_doc_path("test_evaluate_files/unstructured_output/form.json")
|
||||
as_read = partition_json(filename=filename)
|
||||
tmp_file = io.StringIO()
|
||||
json.dump([element.to_dict() for element in as_read], tmp_file)
|
||||
tmp_file.seek(0)
|
||||
as_read_2 = partition_json(file=tmp_file) # type: ignore[arg-type]
|
||||
assert as_read == as_read_2
|
||||
@@ -0,0 +1,46 @@
|
||||
from collections import defaultdict
|
||||
from typing import Type
|
||||
|
||||
from unstructured.documents import elements, ontology
|
||||
from unstructured.documents.mappings import (
|
||||
ALL_ONTOLOGY_ELEMENT_TYPES,
|
||||
HTML_TAG_TO_DEFAULT_ELEMENT_TYPE_MAP,
|
||||
ONTOLOGY_CLASS_TO_UNSTRUCTURED_ELEMENT_TYPE,
|
||||
get_all_subclasses,
|
||||
)
|
||||
from unstructured.documents.ontology import OntologyElement
|
||||
|
||||
|
||||
def test_if_all_html_tags_have_default_ontology_type():
|
||||
html_tag_to_possible_ontology_classes: dict[str, list[Type[ontology.OntologyElement]]] = (
|
||||
defaultdict(list)
|
||||
)
|
||||
|
||||
for ontology_class in ALL_ONTOLOGY_ELEMENT_TYPES:
|
||||
for tag in ontology_class().allowed_tags:
|
||||
html_tag_to_possible_ontology_classes[tag].append(ontology_class)
|
||||
|
||||
for html_tag, possible_ontology_classes in html_tag_to_possible_ontology_classes.items():
|
||||
assert html_tag in HTML_TAG_TO_DEFAULT_ELEMENT_TYPE_MAP
|
||||
assert HTML_TAG_TO_DEFAULT_ELEMENT_TYPE_MAP[html_tag] in possible_ontology_classes + [
|
||||
ontology.UncategorizedText
|
||||
] # In some cases it is better to use unknown type than assign incorrect type
|
||||
|
||||
|
||||
def test_all_expected_ontology_types_are_subclasses_of_OntologyElement():
|
||||
for element_type in HTML_TAG_TO_DEFAULT_ELEMENT_TYPE_MAP.values():
|
||||
assert issubclass(element_type, OntologyElement)
|
||||
|
||||
|
||||
def test_ontology_to_unstructured_mapping_has_valid_types():
|
||||
for (
|
||||
ontology_element,
|
||||
unstructured_element,
|
||||
) in ONTOLOGY_CLASS_TO_UNSTRUCTURED_ELEMENT_TYPE.items():
|
||||
assert issubclass(unstructured_element, elements.Element)
|
||||
assert issubclass(ontology_element, ontology.OntologyElement)
|
||||
|
||||
|
||||
def test_all_ontology_elements_are_defined_in_mapping_to_unstructured():
|
||||
for ontology_element in get_all_subclasses(ontology.OntologyElement):
|
||||
assert ontology_element in ONTOLOGY_CLASS_TO_UNSTRUCTURED_ELEMENT_TYPE
|
||||
@@ -0,0 +1,379 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from bs4 import BeautifulSoup
|
||||
|
||||
from unstructured.chunking.basic import chunk_elements
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.ontology import (
|
||||
Column,
|
||||
Document,
|
||||
Hyperlink,
|
||||
Image,
|
||||
Page,
|
||||
Paragraph,
|
||||
Section,
|
||||
Table,
|
||||
remove_ids_and_class_from_table,
|
||||
)
|
||||
from unstructured.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
|
||||
from unstructured.partition.html import partition_html
|
||||
from unstructured.partition.html.transformations import (
|
||||
ontology_to_unstructured_elements,
|
||||
parse_html_to_ontology,
|
||||
)
|
||||
from unstructured.partition.json import partition_json
|
||||
from unstructured.staging.base import elements_from_json
|
||||
|
||||
|
||||
def test_remove_ids_and_class_from_table():
|
||||
html_text = """
|
||||
<table>
|
||||
<tr class="TableRow">
|
||||
<td><img class="Signature" alt="cell 1"/></td>
|
||||
<td>cell 2</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><IMG class="Signature" alt="cell 3"/></td>
|
||||
<td>cell 4</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><input class="Checkbox" type="checkbox"/></td>
|
||||
<td>Option 1</td>
|
||||
</tr>
|
||||
</table>
|
||||
"""
|
||||
soup = BeautifulSoup(html_text, "html.parser")
|
||||
assert (
|
||||
str(remove_ids_and_class_from_table(soup))
|
||||
== """
|
||||
<table>
|
||||
<tr>
|
||||
<td><img alt="cell 1" class="Signature"/></td>
|
||||
<td>cell 2</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><img alt="cell 3" class="Signature"/></td>
|
||||
<td>cell 4</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><input class="Checkbox" type="checkbox"/></td>
|
||||
<td>Option 1</td>
|
||||
</tr>
|
||||
</table>
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def test_page_number_is_passed_correctly():
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(
|
||||
children=[Paragraph(text="Paragraph1")],
|
||||
additional_attributes={"data-page-number": "1"},
|
||||
),
|
||||
Page(
|
||||
children=[Paragraph(text="Paragraph2")],
|
||||
additional_attributes={"data-page-number": "2"},
|
||||
),
|
||||
]
|
||||
)
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
page1, p1, page2, p2 = unstructured_elements
|
||||
assert p1.metadata.page_number == 1
|
||||
assert p2.metadata.page_number == 2
|
||||
|
||||
|
||||
def test_invalid_page_number_is_not_passed():
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(
|
||||
children=[Paragraph(text="Paragraph1")],
|
||||
additional_attributes={"data-page-number": "invalid"},
|
||||
)
|
||||
]
|
||||
)
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
page1, p1 = unstructured_elements
|
||||
assert not p1.metadata.page_number
|
||||
|
||||
|
||||
def test_depth_is_passed_correctly():
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(children=[Paragraph(text="Paragraph1")]),
|
||||
Page(
|
||||
children=[
|
||||
Column(children=[Paragraph(text="Paragraph2")]),
|
||||
Column(children=[Paragraph(text="Paragraph3")]),
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
page1, p1, page2, c1, p2, c2, p3 = unstructured_elements
|
||||
|
||||
assert page1.metadata.category_depth == 0
|
||||
assert page2.metadata.category_depth == 0
|
||||
|
||||
assert p1.metadata.category_depth == 1
|
||||
|
||||
assert c2.metadata.category_depth == 1
|
||||
assert c1.metadata.category_depth == 1
|
||||
|
||||
assert p2.metadata.category_depth == 2
|
||||
assert p3.metadata.category_depth == 2
|
||||
|
||||
|
||||
def test_chunking_is_applied_on_elements():
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(children=[Paragraph(text="Paragraph1")]),
|
||||
Page(
|
||||
children=[
|
||||
Column(children=[Paragraph(text="Paragraph2")]),
|
||||
Column(children=[Paragraph(text="Paragraph3")]),
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
|
||||
chunked_basic = chunk_elements(unstructured_elements)
|
||||
assert str(chunked_basic[0]) == "Paragraph1\n\nParagraph2\n\nParagraph3"
|
||||
chunked_by_title = chunk_by_title(unstructured_elements)
|
||||
assert str(chunked_by_title[0]) == "Paragraph1\n\nParagraph2\n\nParagraph3"
|
||||
|
||||
|
||||
def test_embeddings_are_applied_on_elements(mocker):
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(children=[Paragraph(text="Paragraph1")]),
|
||||
Page(
|
||||
children=[
|
||||
Column(children=[Paragraph(text="Paragraph2")]),
|
||||
Column(children=[Paragraph(text="Paragraph3")]),
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
# Mocked client with the desired behavior for embed_documents
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.embed_documents.return_value = [1, 2, 3, 4, 5, 6, 7]
|
||||
|
||||
# Mock get_client to return our mock_client
|
||||
mocker.patch.object(OpenAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(api_key="api_key"))
|
||||
elements = encoder.embed_documents(
|
||||
elements=unstructured_elements,
|
||||
)
|
||||
|
||||
assert len(elements) == 7
|
||||
|
||||
page1, p1, page2, c1, p2, c2, p3 = elements
|
||||
|
||||
assert p1.embeddings == 2
|
||||
assert p2.embeddings == 5
|
||||
assert p3.embeddings == 7
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("html_file_path", "json_file_path"),
|
||||
[
|
||||
("html_files/example.html", "unstructured_json_output/example.json"),
|
||||
],
|
||||
)
|
||||
def test_ingest(html_file_path, json_file_path):
|
||||
html_file_path = Path(__file__).parent / html_file_path
|
||||
json_file_path = Path(__file__).parent / json_file_path
|
||||
|
||||
html_code = html_file_path.read_text()
|
||||
expected_json_elements = elements_from_json(str(json_file_path))
|
||||
|
||||
ontology = parse_html_to_ontology(html_code)
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
assert unstructured_elements == expected_json_elements
|
||||
|
||||
|
||||
@pytest.mark.parametrize("json_file_path", ["unstructured_json_output/example.json"])
|
||||
def test_parsed_ontology_can_be_serialized_from_json(json_file_path):
|
||||
json_file_path = Path(__file__).parent / json_file_path
|
||||
|
||||
expected_json_elements = elements_from_json(str(json_file_path))
|
||||
|
||||
json_elements_text = json_file_path.read_text()
|
||||
elements = partition_json(text=json_elements_text)
|
||||
|
||||
assert len(elements) == len(expected_json_elements)
|
||||
for i in range(len(elements)):
|
||||
assert elements[i] == expected_json_elements[i]
|
||||
# The partitioning output comes from PDF file, so only stem is compared
|
||||
# as the suffix is different .pdf != .json
|
||||
assert Path(elements[i].metadata.filename).stem == json_file_path.stem
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("html_file_path", "json_file_path"),
|
||||
[
|
||||
("html_files/example.html", "unstructured_json_output/example.json"),
|
||||
("html_files/example_full_doc.html", "unstructured_json_output/example_full_doc.json"),
|
||||
(
|
||||
"html_files/example_with_alternative_text.html",
|
||||
"unstructured_json_output/example_with_alternative_text.json",
|
||||
),
|
||||
("html_files/three_tables.html", "unstructured_json_output/three_tables.json"),
|
||||
(
|
||||
"html_files/example_with_inline_fields.html",
|
||||
"unstructured_json_output/example_with_inline_fields.json",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_parsed_ontology_can_be_serialized_from_html(html_file_path, json_file_path):
|
||||
html_file_path = Path(__file__).parent / html_file_path
|
||||
json_file_path = Path(__file__).parent / json_file_path
|
||||
expected_json_elements = elements_from_json(str(json_file_path))
|
||||
html_code = html_file_path.read_text()
|
||||
|
||||
predicted_elements = partition_html(
|
||||
text=html_code, html_parser_version="v2", unique_element_ids=True
|
||||
)
|
||||
|
||||
assert len(expected_json_elements) == len(predicted_elements)
|
||||
|
||||
for i in range(len(expected_json_elements)):
|
||||
assert expected_json_elements[i] == predicted_elements[i]
|
||||
assert (
|
||||
expected_json_elements[i].metadata.text_as_html
|
||||
== predicted_elements[i].metadata.text_as_html
|
||||
)
|
||||
|
||||
|
||||
def test_inline_elements_are_squeezed():
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(
|
||||
children=[
|
||||
Hyperlink(text="Hyperlink1"),
|
||||
Hyperlink(text="Hyperlink2"),
|
||||
Hyperlink(text="Hyperlink3"),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
assert len(unstructured_elements) == 2
|
||||
|
||||
page, text1 = unstructured_elements
|
||||
assert text1.text == "Hyperlink1 Hyperlink2 Hyperlink3"
|
||||
|
||||
|
||||
def test_text_elements_are_squeezed():
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(
|
||||
children=[
|
||||
Paragraph(text="Paragraph1"),
|
||||
Paragraph(text="Paragraph2"),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
assert len(unstructured_elements) == 2
|
||||
|
||||
page, text1 = unstructured_elements
|
||||
assert text1.text == "Paragraph1 Paragraph2"
|
||||
|
||||
|
||||
def test_inline_elements_are_squeezed_when_image():
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(
|
||||
children=[
|
||||
Paragraph(text="Paragraph1"),
|
||||
Hyperlink(text="Hyperlink1"),
|
||||
Image(text="Image1"),
|
||||
Hyperlink(text="Hyperlink2"),
|
||||
Hyperlink(text="Hyperlink3"),
|
||||
Paragraph(text="Paragraph2"),
|
||||
Paragraph(text="Paragraph3"),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
assert len(unstructured_elements) == 4
|
||||
|
||||
page, text1, image, text2 = unstructured_elements
|
||||
assert text1.text == "Paragraph1 Hyperlink1"
|
||||
assert text2.text == "Hyperlink2 Hyperlink3 Paragraph2 Paragraph3"
|
||||
|
||||
assert '<a class="Hyperlink"' in text1.metadata.text_as_html
|
||||
assert '<p class="Paragraph"' in text1.metadata.text_as_html
|
||||
|
||||
assert '<a class="Hyperlink"' in text2.metadata.text_as_html
|
||||
assert '<p class="Paragraph"' in text2.metadata.text_as_html
|
||||
|
||||
|
||||
def test_inline_elements_are_squeezed_when_table():
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(
|
||||
children=[
|
||||
Hyperlink(text="Hyperlink1"),
|
||||
Paragraph(text="Paragraph1"),
|
||||
Paragraph(text="Paragraph2"),
|
||||
Table(text="Table1"),
|
||||
Paragraph(text="Paragraph2"),
|
||||
Hyperlink(text="Hyperlink2"),
|
||||
Hyperlink(text="Hyperlink3"),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
assert len(unstructured_elements) == 4
|
||||
|
||||
page, text1, table1, text3 = unstructured_elements
|
||||
assert text1.text == "Hyperlink1 Paragraph1 Paragraph2"
|
||||
assert table1.text == "Table1"
|
||||
assert text3.text == "Paragraph2 Hyperlink2 Hyperlink3"
|
||||
|
||||
|
||||
def test_inline_elements_are_on_many_depths():
|
||||
ontology = Document(
|
||||
children=[
|
||||
Page(
|
||||
children=[
|
||||
Hyperlink(text="Hyperlink1"),
|
||||
Paragraph(text="Paragraph1"),
|
||||
Section(
|
||||
children=[
|
||||
Section(
|
||||
children=[
|
||||
Hyperlink(text="Hyperlink2"),
|
||||
Hyperlink(text="Hyperlink3"),
|
||||
]
|
||||
),
|
||||
Paragraph(text="Paragraph2"),
|
||||
Hyperlink(text="Hyperlink4"),
|
||||
]
|
||||
),
|
||||
],
|
||||
)
|
||||
]
|
||||
)
|
||||
unstructured_elements = ontology_to_unstructured_elements(ontology)
|
||||
|
||||
assert len(unstructured_elements) == 6
|
||||
|
||||
page, text1, section1, section2, text2, text3 = unstructured_elements
|
||||
|
||||
assert text1.text == "Hyperlink1 Paragraph1"
|
||||
assert text2.text == "Hyperlink2 Hyperlink3"
|
||||
assert text3.text == "Paragraph2 Hyperlink4"
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,41 @@
|
||||
from unstructured.documents.elements import Text
|
||||
from unstructured.embed.mixedbreadai import (
|
||||
MixedbreadAIEmbeddingConfig,
|
||||
MixedbreadAIEmbeddingEncoder,
|
||||
)
|
||||
|
||||
|
||||
def test_embed_documents_does_not_break_element_to_dict(mocker):
|
||||
mock_client = mocker.MagicMock()
|
||||
|
||||
def mock_embeddings(
|
||||
model,
|
||||
normalized,
|
||||
encoding_format,
|
||||
truncation_strategy,
|
||||
request_options,
|
||||
input,
|
||||
):
|
||||
mock_response = mocker.MagicMock()
|
||||
mock_response.data = [mocker.MagicMock(embedding=[i, i + 1]) for i in range(len(input))]
|
||||
return mock_response
|
||||
|
||||
mock_client.embeddings.side_effect = mock_embeddings
|
||||
|
||||
# Mock get_client to return our mock_client
|
||||
mocker.patch.object(MixedbreadAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = MixedbreadAIEmbeddingEncoder(
|
||||
config=MixedbreadAIEmbeddingConfig(
|
||||
api_key="api_key", model_name="mixedbread-ai/mxbai-embed-large-v1"
|
||||
)
|
||||
)
|
||||
|
||||
elements = encoder.embed_documents(
|
||||
elements=[Text("This is sentence 1"), Text("This is sentence 2")],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
assert elements[0].to_dict()["text"] == "This is sentence 1"
|
||||
assert elements[1].to_dict()["text"] == "This is sentence 2"
|
||||
assert elements[0].embeddings is not None
|
||||
assert elements[1].embeddings is not None
|
||||
@@ -0,0 +1,19 @@
|
||||
from unstructured.documents.elements import Text
|
||||
from unstructured.embed.octoai import OctoAiEmbeddingConfig, OctoAIEmbeddingEncoder
|
||||
|
||||
|
||||
def test_embed_documents_does_not_break_element_to_dict(mocker):
|
||||
# Mocked client with the desired behavior for embed_documents
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.embed_documents.return_value = [1, 2]
|
||||
|
||||
# Mock get_client to return our mock_client
|
||||
mocker.patch.object(OctoAiEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = OctoAIEmbeddingEncoder(config=OctoAiEmbeddingConfig(api_key="api_key"))
|
||||
elements = encoder.embed_documents(
|
||||
elements=[Text("This is sentence 1"), Text("This is sentence 2")],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
assert elements[0].to_dict()["text"] == "This is sentence 1"
|
||||
assert elements[1].to_dict()["text"] == "This is sentence 2"
|
||||
@@ -0,0 +1,19 @@
|
||||
from unstructured.documents.elements import Text
|
||||
from unstructured.embed.openai import OpenAIEmbeddingConfig, OpenAIEmbeddingEncoder
|
||||
|
||||
|
||||
def test_embed_documents_does_not_break_element_to_dict(mocker):
|
||||
# Mocked client with the desired behavior for embed_documents
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.embed_documents.return_value = [1, 2]
|
||||
|
||||
# Mock get_client to return our mock_client
|
||||
mocker.patch.object(OpenAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = OpenAIEmbeddingEncoder(config=OpenAIEmbeddingConfig(api_key="api_key"))
|
||||
elements = encoder.embed_documents(
|
||||
elements=[Text("This is sentence 1"), Text("This is sentence 2")],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
assert elements[0].to_dict()["text"] == "This is sentence 1"
|
||||
assert elements[1].to_dict()["text"] == "This is sentence 2"
|
||||
@@ -0,0 +1,19 @@
|
||||
from unstructured.documents.elements import Text
|
||||
from unstructured.embed.vertexai import VertexAIEmbeddingConfig, VertexAIEmbeddingEncoder
|
||||
|
||||
|
||||
def test_embed_documents_does_not_break_element_to_dict(mocker):
|
||||
# Mocked client with the desired behavior for embed_documents
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.embed_documents.return_value = [1, 2]
|
||||
|
||||
# Mock create_client to return our mock_client
|
||||
mocker.patch.object(VertexAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = VertexAIEmbeddingEncoder(config=VertexAIEmbeddingConfig(api_key="api_key"))
|
||||
elements = encoder.embed_documents(
|
||||
elements=[Text("This is sentence 1"), Text("This is sentence 2")],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
assert elements[0].to_dict()["text"] == "This is sentence 1"
|
||||
assert elements[1].to_dict()["text"] == "This is sentence 2"
|
||||
@@ -0,0 +1,242 @@
|
||||
from unittest.mock import Mock
|
||||
|
||||
from unstructured.documents.elements import Text
|
||||
from unstructured.embed.voyageai import VoyageAIEmbeddingConfig, VoyageAIEmbeddingEncoder
|
||||
|
||||
|
||||
def test_embed_documents_does_not_break_element_to_dict(mocker):
|
||||
# Mocked client with the desired behavior for embed_documents
|
||||
embed_response = Mock()
|
||||
embed_response.embeddings = [[1], [2]]
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.embed.return_value = embed_response
|
||||
mock_client.tokenize.return_value = [[1], [1]] # Mock token counts
|
||||
|
||||
# Mock get_client to return our mock_client
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3-large")
|
||||
)
|
||||
elements = encoder.embed_documents(
|
||||
elements=[Text("This is sentence 1"), Text("This is sentence 2")],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
assert elements[0].to_dict()["text"] == "This is sentence 1"
|
||||
assert elements[1].to_dict()["text"] == "This is sentence 2"
|
||||
|
||||
|
||||
def test_embed_documents_voyage_3_5(mocker):
|
||||
"""Test embedding with voyage-3.5 model."""
|
||||
embed_response = Mock()
|
||||
embed_response.embeddings = [[1.0] * 1024, [2.0] * 1024]
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.embed.return_value = embed_response
|
||||
mock_client.tokenize.return_value = [[1, 2, 3], [1, 2]] # Mock token counts
|
||||
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
|
||||
)
|
||||
elements = encoder.embed_documents(
|
||||
elements=[Text("Test document 1"), Text("Test document 2")],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
assert len(elements[0].embeddings) == 1024
|
||||
assert len(elements[1].embeddings) == 1024
|
||||
|
||||
|
||||
def test_embed_documents_voyage_3_5_lite(mocker):
|
||||
"""Test embedding with voyage-3.5-lite model."""
|
||||
embed_response = Mock()
|
||||
embed_response.embeddings = [[1.0] * 512, [2.0] * 512, [3.0] * 512]
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.embed.return_value = embed_response
|
||||
mock_client.tokenize.return_value = [[1], [1], [1]] # Mock token counts
|
||||
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5-lite")
|
||||
)
|
||||
elements = encoder.embed_documents(
|
||||
elements=[Text("Test 1"), Text("Test 2"), Text("Test 3")],
|
||||
)
|
||||
assert len(elements) == 3
|
||||
assert all(len(e.embeddings) == 512 for e in elements)
|
||||
|
||||
|
||||
def test_embed_documents_contextual_model(mocker):
|
||||
"""Test embedding with voyage-context-3 model."""
|
||||
# Mock contextualized_embed response
|
||||
contextualized_response = Mock()
|
||||
result_item = Mock()
|
||||
result_item.embeddings = [[1.0] * 1024, [2.0] * 1024]
|
||||
contextualized_response.results = [result_item]
|
||||
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.contextualized_embed.return_value = contextualized_response
|
||||
mock_client.tokenize.return_value = [[1, 2], [1, 2, 3]] # Mock token counts
|
||||
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-context-3")
|
||||
)
|
||||
elements = encoder.embed_documents(
|
||||
elements=[Text("Context document 1"), Text("Context document 2")],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
assert len(elements[0].embeddings) == 1024
|
||||
assert len(elements[1].embeddings) == 1024
|
||||
# Verify contextualized_embed was called
|
||||
mock_client.contextualized_embed.assert_called_once()
|
||||
|
||||
|
||||
def test_count_tokens(mocker):
|
||||
"""Test token counting functionality."""
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.tokenize.return_value = [[1, 2], [1, 2, 3, 4, 5]] # Different token counts
|
||||
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
|
||||
)
|
||||
texts = ["short text", "this is a longer text with more tokens"]
|
||||
token_counts = encoder.count_tokens(texts)
|
||||
|
||||
assert len(token_counts) == 2
|
||||
assert token_counts[0] == 2
|
||||
assert token_counts[1] == 5
|
||||
|
||||
|
||||
def test_count_tokens_empty_list(mocker):
|
||||
"""Test token counting with empty list."""
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mocker.MagicMock())
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
|
||||
)
|
||||
token_counts = encoder.count_tokens([])
|
||||
assert token_counts == []
|
||||
|
||||
|
||||
def test_get_token_limit(mocker):
|
||||
"""Test getting token limit for different models."""
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mocker.MagicMock())
|
||||
|
||||
# Test voyage-3.5 model
|
||||
config = VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
|
||||
assert config.get_token_limit() == 320_000
|
||||
|
||||
# Test voyage-3.5-lite model
|
||||
config_lite = VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5-lite")
|
||||
assert config_lite.get_token_limit() == 1_000_000
|
||||
|
||||
# Test context model
|
||||
config_context = VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-context-3")
|
||||
assert config_context.get_token_limit() == 32_000
|
||||
|
||||
# Test voyage-2 model
|
||||
config_v2 = VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-2")
|
||||
assert config_v2.get_token_limit() == 320_000
|
||||
|
||||
# Test unknown model (should use default)
|
||||
config_unknown = VoyageAIEmbeddingConfig(api_key="api_key", model_name="unknown-model")
|
||||
assert config_unknown.get_token_limit() == 120_000
|
||||
|
||||
|
||||
def test_is_context_model(mocker):
|
||||
"""Test the _is_context_model helper method."""
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mocker.MagicMock())
|
||||
|
||||
# Test with context model
|
||||
encoder_context = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-context-3")
|
||||
)
|
||||
assert encoder_context._is_context_model() is True
|
||||
|
||||
# Test with regular model
|
||||
encoder_regular = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
|
||||
)
|
||||
assert encoder_regular._is_context_model() is False
|
||||
|
||||
|
||||
def test_build_batches_with_token_limits(mocker):
|
||||
"""Test that batching respects token limits."""
|
||||
mock_client = mocker.MagicMock()
|
||||
# Simulate different token counts for each text
|
||||
mock_client.tokenize.return_value = [[1] * 10, [1] * 20, [1] * 15, [1] * 25]
|
||||
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-2")
|
||||
)
|
||||
texts = ["text1", "text2", "text3", "text4"]
|
||||
batches = list(encoder._build_batches(texts, mock_client))
|
||||
|
||||
# Should create at least one batch
|
||||
assert len(batches) >= 1
|
||||
# Total texts should be preserved
|
||||
total_texts = sum(len(batch) for batch in batches)
|
||||
assert total_texts == len(texts)
|
||||
|
||||
|
||||
def test_embed_query(mocker):
|
||||
"""Test embedding a single query."""
|
||||
embed_response = Mock()
|
||||
embed_response.embeddings = [[1.0] * 1024]
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.embed.return_value = embed_response
|
||||
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
|
||||
)
|
||||
embedding = encoder.embed_query("test query")
|
||||
|
||||
assert len(embedding) == 1024
|
||||
# Verify embed was called with input_type="query"
|
||||
mock_client.embed.assert_called_once()
|
||||
call_kwargs = mock_client.embed.call_args[1]
|
||||
assert call_kwargs["input_type"] == "query"
|
||||
|
||||
|
||||
def test_embed_documents_with_output_dimension(mocker):
|
||||
"""Test embedding with custom output dimension."""
|
||||
embed_response = Mock()
|
||||
embed_response.embeddings = [[1.0] * 512, [2.0] * 512]
|
||||
mock_client = mocker.MagicMock()
|
||||
mock_client.embed.return_value = embed_response
|
||||
mock_client.tokenize.return_value = [[1], [1]]
|
||||
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mock_client)
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(
|
||||
api_key="api_key", model_name="voyage-3.5", output_dimension=512
|
||||
)
|
||||
)
|
||||
elements = encoder.embed_documents(
|
||||
elements=[Text("Test 1"), Text("Test 2")],
|
||||
)
|
||||
assert len(elements) == 2
|
||||
# Verify output_dimension was passed
|
||||
call_kwargs = mock_client.embed.call_args[1]
|
||||
assert call_kwargs["output_dimension"] == 512
|
||||
|
||||
|
||||
def test_embed_documents_empty_list(mocker):
|
||||
"""Test embedding empty list of documents."""
|
||||
mocker.patch.object(VoyageAIEmbeddingConfig, "get_client", return_value=mocker.MagicMock())
|
||||
|
||||
encoder = VoyageAIEmbeddingEncoder(
|
||||
config=VoyageAIEmbeddingConfig(api_key="api_key", model_name="voyage-3.5")
|
||||
)
|
||||
elements = encoder.embed_documents(elements=[])
|
||||
assert elements == []
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,71 @@
|
||||
"""Test encoding detection error handling (PR #4071)."""
|
||||
|
||||
import os
|
||||
import pickle
|
||||
import sys
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from unstructured.errors import UnprocessableEntityError
|
||||
from unstructured.file_utils.encoding import detect_file_encoding
|
||||
|
||||
|
||||
def test_charset_detection_failure():
|
||||
"""Test encoding detection failure with memory safety checks."""
|
||||
large_data = b"\x80\x81\x82\x83" * 250_000 # 1MB of invalid UTF-8
|
||||
|
||||
with tempfile.NamedTemporaryFile(mode="wb", delete=False) as f:
|
||||
f.write(large_data)
|
||||
temp_file_path = f.name
|
||||
|
||||
try:
|
||||
detect_result = {"encoding": None, "confidence": None}
|
||||
with patch("unstructured.file_utils.encoding.detect", return_value=detect_result):
|
||||
with patch("unstructured.file_utils.encoding.COMMON_ENCODINGS", ["utf_8"]): # Will fail
|
||||
with pytest.raises(UnprocessableEntityError) as exc_info:
|
||||
detect_file_encoding(filename=temp_file_path)
|
||||
|
||||
exception = exc_info.value
|
||||
|
||||
assert "Unable to determine file encoding" in str(exception)
|
||||
|
||||
# Ensure no .object attribute that would store file content (prevents memory bloat)
|
||||
# See: https://docs.python.org/3/library/exceptions.html#UnicodeError.object
|
||||
assert not hasattr(exception, "object")
|
||||
|
||||
# Exception should be lightweight regardless of file size
|
||||
exception_memory = sys.getsizeof(exception)
|
||||
serialized_size = len(pickle.dumps(exception))
|
||||
|
||||
assert exception_memory < 10_000 # Small in-memory footprint
|
||||
assert serialized_size < 10_000 # Small serialization footprint
|
||||
finally:
|
||||
os.unlink(temp_file_path)
|
||||
|
||||
|
||||
def test_decode_failure():
|
||||
"""Test decode failure with memory safety checks."""
|
||||
# Invalid UTF-16: BOM followed by odd number of bytes
|
||||
invalid_utf16 = b"\xff\xfe" + b"A\x00B\x00" + b"\x00"
|
||||
|
||||
detect_result = {"encoding": "utf-16", "confidence": 0.95}
|
||||
with patch("unstructured.file_utils.encoding.detect", return_value=detect_result):
|
||||
with pytest.raises(UnprocessableEntityError) as exc_info:
|
||||
detect_file_encoding(file=invalid_utf16)
|
||||
|
||||
exception = exc_info.value
|
||||
|
||||
assert "detected 'utf-16' but decode failed" in str(exception)
|
||||
|
||||
# Ensure no .object attribute that would store file content (prevents memory bloat)
|
||||
# See: https://docs.python.org/3/library/exceptions.html#UnicodeError.object
|
||||
assert not hasattr(exception, "object")
|
||||
|
||||
# Exception should be lightweight
|
||||
exception_memory = sys.getsizeof(exception)
|
||||
serialized_size = len(pickle.dumps(exception))
|
||||
|
||||
assert exception_memory < 10_000 # Small in-memory footprint
|
||||
assert serialized_size < 10_000 # Small serialization footprint
|
||||
@@ -0,0 +1,54 @@
|
||||
import os
|
||||
import pathlib
|
||||
import tempfile
|
||||
from unittest.mock import patch
|
||||
|
||||
import pypandoc
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import FixtureRequest, example_doc_path, stdlib_fn_mock
|
||||
from unstructured.file_utils.file_conversion import (
|
||||
convert_file_to_html_text_using_pandoc,
|
||||
convert_file_to_text,
|
||||
)
|
||||
|
||||
DIRECTORY = pathlib.Path(__file__).parent.resolve()
|
||||
|
||||
|
||||
def test_convert_file_to_text():
|
||||
filename = os.path.join(DIRECTORY, "..", "..", "example-docs", "winter-sports.epub")
|
||||
html_text = convert_file_to_text(filename, source_format="epub", target_format="html")
|
||||
assert html_text.startswith("<p>")
|
||||
|
||||
|
||||
def test_convert_to_file_raises_if_pandoc_not_available():
|
||||
filename = os.path.join(DIRECTORY, "..", "..", "example-docs", "winter-sports.epub")
|
||||
with patch.object(pypandoc, "convert_file", side_effect=FileNotFoundError):
|
||||
with pytest.raises(FileNotFoundError):
|
||||
convert_file_to_text(filename, source_format="epub", target_format="html")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("source_format", "filename"),
|
||||
[
|
||||
("epub", "winter-sports.epub"),
|
||||
("org", "README.org"),
|
||||
("rst", "README.rst"),
|
||||
("rtf", "fake-doc.rtf"),
|
||||
],
|
||||
)
|
||||
def test_convert_file_to_html_text_using_pandoc(
|
||||
request: FixtureRequest, tmp_path: pathlib.Path, source_format: str, filename: str
|
||||
):
|
||||
# -- Get a real tempdir: `tmp_path`
|
||||
# -- Mock tempfile.TemporaryDirectory() using `stdlib_fn_mock`
|
||||
# -- Set the return value of mock.__enter__ to the real tempdir
|
||||
tempdir_ = stdlib_fn_mock(request, tempfile, "TemporaryDirectory")
|
||||
tempdir_.return_value.__enter__.return_value = tmp_path
|
||||
|
||||
with open(example_doc_path(filename), "rb") as f:
|
||||
html_text = convert_file_to_html_text_using_pandoc(file=f, source_format=source_format)
|
||||
|
||||
assert isinstance(html_text, str)
|
||||
assert len(list(tmp_path.iterdir())) == 1
|
||||
tempdir_.return_value.__exit__.assert_called_once()
|
||||
@@ -0,0 +1,998 @@
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
"""Test suite for `unstructured.file_utils.filetype`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import (
|
||||
FixtureRequest,
|
||||
LogCaptureFixture,
|
||||
Mock,
|
||||
example_doc_path,
|
||||
input_path,
|
||||
patch,
|
||||
property_mock,
|
||||
)
|
||||
from unstructured.file_utils.filetype import (
|
||||
_FileTypeDetectionContext,
|
||||
_OleFileDetector,
|
||||
_TextFileDifferentiator,
|
||||
_ZipFileDetector,
|
||||
detect_filetype,
|
||||
is_json_processable,
|
||||
)
|
||||
from unstructured.file_utils.model import FileType, create_file_type
|
||||
|
||||
is_in_docker = os.path.exists("/.dockerenv")
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# STRATEGY #1 - DIRECT DETECTION OF CFB/ZIP-BASED BINARY FILE TYPES (8 TYPES)
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name"),
|
||||
[
|
||||
(FileType.DOC, "simple.doc"),
|
||||
(FileType.DOCX, "simple.docx"),
|
||||
(FileType.EPUB, "winter-sports.epub"),
|
||||
(FileType.ODT, "simple.odt"),
|
||||
(FileType.PPT, "fake-power-point.ppt"),
|
||||
(FileType.PPTX, "fake-power-point.pptx"),
|
||||
(FileType.XLS, "tests-example.xls"),
|
||||
(FileType.XLSX, "stanley-cups.xlsx"),
|
||||
],
|
||||
)
|
||||
def test_it_detects_correct_file_type_for_CFB_and_ZIP_subtypes_detected_by_direct_inspection(
|
||||
file_name: str, expected_value: FileType, ctx_mime_type_: Mock
|
||||
):
|
||||
# -- disable other strategies; no content-type, guessed MIME-type or extension --
|
||||
ctx_mime_type_.return_value = None
|
||||
with open(example_doc_path(file_name), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
|
||||
file_type = detect_filetype(file=file)
|
||||
|
||||
# -- Strategy 1 should not need to refer to guessed MIME-type and detection should not
|
||||
# -- fall back to MIME-type guessing for any of these test cases.
|
||||
ctx_mime_type_.assert_not_called()
|
||||
assert file_type == expected_value
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# STRATEGY #2 - CONTENT-TYPE ASSERTED IN CALL
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name", "content_type"),
|
||||
[
|
||||
(FileType.BMP, "img/bmp_24.bmp", "image/bmp"),
|
||||
(FileType.CSV, "stanley-cups.csv", "text/csv"),
|
||||
(FileType.EML, "eml/fake-email.eml", "message/rfc822"),
|
||||
(FileType.HEIC, "img/DA-1p.heic", "image/heic"),
|
||||
(FileType.HTML, "example-10k-1p.html", "text/html"),
|
||||
(FileType.JPG, "img/example.jpg", "image/jpeg"),
|
||||
(FileType.MD, "README.md", "text/markdown"),
|
||||
(FileType.ORG, "README.org", "text/org"),
|
||||
(FileType.PDF, "pdf/layout-parser-paper-fast.pdf", "application/pdf"),
|
||||
(FileType.PNG, "img/DA-1p.png", "image/png"),
|
||||
(FileType.RST, "README.rst", "text/x-rst"),
|
||||
(FileType.RTF, "fake-doc.rtf", "text/rtf"),
|
||||
(FileType.TIFF, "img/layout-parser-paper-fast.tiff", "image/tiff"),
|
||||
(FileType.TSV, "stanley-cups.tsv", "text/tsv"),
|
||||
(FileType.TXT, "norwich-city.txt", "text/plain"),
|
||||
(FileType.WAV, "CantinaBand3.wav", "audio/wav"),
|
||||
(FileType.XML, "factbook.xml", "application/xml"),
|
||||
(FileType.ZIP, "simple.zip", "application/zip"),
|
||||
(FileType.NDJSON, "spring-weather.html.ndjson", "application/x-ndjson"),
|
||||
],
|
||||
)
|
||||
def test_it_detects_correct_file_type_from_file_path_with_correct_asserted_content_type(
|
||||
file_name: str, content_type: str, expected_value: FileType, ctx_mime_type_: Mock
|
||||
):
|
||||
# -- disable mime-guessing leaving only asserted content-type and extension --
|
||||
ctx_mime_type_.return_value = None
|
||||
|
||||
file_type = detect_filetype(example_doc_path(file_name), content_type=content_type)
|
||||
|
||||
# -- Content-type strategy should not need to refer to guessed MIME-type and detection should
|
||||
# not -- fall back to strategy 2 for any of these test cases.
|
||||
ctx_mime_type_.assert_not_called()
|
||||
assert file_type == expected_value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name", "content_type"),
|
||||
[
|
||||
(FileType.BMP, "img/bmp_24.bmp", "image/bmp"),
|
||||
(FileType.CSV, "stanley-cups.csv", "text/csv"),
|
||||
(FileType.EML, "eml/fake-email.eml", "message/rfc822"),
|
||||
(FileType.HEIC, "img/DA-1p.heic", "image/heic"),
|
||||
(FileType.HTML, "example-10k-1p.html", "text/html"),
|
||||
(FileType.JPG, "img/example.jpg", "image/jpeg"),
|
||||
(FileType.MD, "README.md", "text/markdown"),
|
||||
(FileType.ORG, "README.org", "text/org"),
|
||||
(FileType.PDF, "pdf/layout-parser-paper-fast.pdf", "application/pdf"),
|
||||
(FileType.PNG, "img/DA-1p.png", "image/png"),
|
||||
(FileType.RST, "README.rst", "text/x-rst"),
|
||||
(FileType.RTF, "fake-doc.rtf", "text/rtf"),
|
||||
(FileType.TIFF, "img/layout-parser-paper-fast.tiff", "image/tiff"),
|
||||
(FileType.TSV, "stanley-cups.tsv", "text/tsv"),
|
||||
(FileType.TXT, "norwich-city.txt", "text/plain"),
|
||||
(FileType.WAV, "CantinaBand3.wav", "audio/wav"),
|
||||
(FileType.XML, "factbook.xml", "application/xml"),
|
||||
(FileType.ZIP, "simple.zip", "application/zip"),
|
||||
],
|
||||
)
|
||||
def test_it_detects_correct_file_type_from_file_no_name_with_correct_asserted_content_type(
|
||||
file_name: str, content_type: str, expected_value: FileType, ctx_mime_type_: Mock
|
||||
):
|
||||
# -- disable mime-guessing --
|
||||
ctx_mime_type_.return_value = None
|
||||
# -- disable filename extension mapping by supplying no source of file name --
|
||||
with open(example_doc_path(file_name), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
|
||||
file_type = detect_filetype(file=file, content_type=content_type)
|
||||
|
||||
# -- Content-type strategy should not need to refer to guessed MIME-type and detection should
|
||||
# -- not fall-back to strategy 2 for any of these test cases.
|
||||
ctx_mime_type_.assert_not_called()
|
||||
assert file_type is expected_value
|
||||
|
||||
|
||||
def test_it_identifies_NDJSON_for_file_like_object_with_no_name_but_NDJSON_content_type():
|
||||
with open(example_doc_path("simple.ndjson"), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
assert detect_filetype(file=file, content_type=FileType.NDJSON.mime_type) == FileType.NDJSON
|
||||
|
||||
|
||||
def test_it_identifies_NDJSON_for_file_with_ndjson_extension_but_JSON_content_type():
|
||||
file_path = example_doc_path("simple.ndjson")
|
||||
assert detect_filetype(file_path, content_type=FileType.JSON.mime_type) == FileType.NDJSON
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# STRATEGY #3 - GUESS MIME-TYPE WITH LIBMAGIC/FILETYPE LIBRARY
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name", "mime_type"),
|
||||
[
|
||||
(FileType.BMP, "img/bmp_24.bmp", "image/bmp"),
|
||||
(FileType.CSV, "stanley-cups.csv", "text/csv"),
|
||||
(FileType.CSV, "stanley-cups.csv", "application/csv"),
|
||||
(FileType.CSV, "stanley-cups.csv", "application/x-csv"),
|
||||
(FileType.EML, "eml/fake-email.eml", "message/rfc822"),
|
||||
(FileType.HEIC, "img/DA-1p.heic", "image/heic"),
|
||||
(FileType.HTML, "example-10k-1p.html", "text/html"),
|
||||
(FileType.JPG, "img/example.jpg", "image/jpeg"),
|
||||
(FileType.JSON, "spring-weather.html.json", "application/json"),
|
||||
(FileType.MD, "README.md", "text/markdown"),
|
||||
(FileType.MD, "README.md", "text/x-markdown"),
|
||||
(FileType.ORG, "README.org", "text/org"),
|
||||
(FileType.PDF, "pdf/layout-parser-paper-fast.pdf", "application/pdf"),
|
||||
(FileType.PNG, "img/DA-1p.png", "image/png"),
|
||||
(FileType.RST, "README.rst", "text/x-rst"),
|
||||
(FileType.RTF, "fake-doc.rtf", "text/rtf"),
|
||||
(FileType.RTF, "fake-doc.rtf", "application/rtf"),
|
||||
(FileType.TIFF, "img/layout-parser-paper-fast.tiff", "image/tiff"),
|
||||
(FileType.TSV, "stanley-cups.tsv", "text/tsv"),
|
||||
(FileType.TXT, "norwich-city.txt", "text/plain"),
|
||||
(FileType.TXT, "simple.yaml", "text/yaml"),
|
||||
(FileType.WAV, "CantinaBand3.wav", "audio/wav"),
|
||||
(FileType.XML, "factbook.xml", "application/xml"),
|
||||
(FileType.XML, "factbook.xml", "text/xml"),
|
||||
],
|
||||
)
|
||||
def test_it_detects_correct_file_type_by_guessed_MIME_when_libmagic_guesses_recognized_mime_type(
|
||||
file_name: str, mime_type: str, expected_value: FileType, ctx_mime_type_: Mock
|
||||
):
|
||||
# -- libmagic guesses a MIME-type mapped to a `FileType` --
|
||||
ctx_mime_type_.return_value = mime_type
|
||||
# -- disable strategy #3 (filename extension) by not providing filename --
|
||||
with open(example_doc_path(file_name), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
|
||||
# -- disable content-type strategy by not asserting a content_type in the call --
|
||||
file_type = detect_filetype(file=file)
|
||||
|
||||
# -- ctx.mime_type may be referenced multiple times, but at least once --
|
||||
ctx_mime_type_.assert_called_with()
|
||||
assert file_type is expected_value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name"),
|
||||
[
|
||||
(FileType.BMP, "img/bmp_24.bmp"),
|
||||
(FileType.CSV, "stanley-cups.csv"),
|
||||
(FileType.EML, "eml/fake-email.eml"),
|
||||
(FileType.HEIC, "img/DA-1p.heic"),
|
||||
(FileType.HTML, "ideas-page.html"),
|
||||
(FileType.JPG, "img/example.jpg"),
|
||||
(FileType.JSON, "spring-weather.html.json"),
|
||||
(FileType.PDF, "pdf/layout-parser-paper-fast.pdf"),
|
||||
(FileType.PNG, "img/DA-1p.png"),
|
||||
(FileType.RTF, "fake-doc.rtf"),
|
||||
(FileType.TIFF, "img/layout-parser-paper-fast.tiff"),
|
||||
(FileType.TXT, "norwich-city.txt"),
|
||||
(FileType.WAV, "CantinaBand3.wav"),
|
||||
(FileType.XML, "factbook.xml"),
|
||||
(FileType.ZIP, "simple.zip"),
|
||||
],
|
||||
)
|
||||
def test_it_detects_most_file_types_using_mime_guessing_when_libmagic_guesses_mime_type_for_itself(
|
||||
file_name: str, expected_value: FileType
|
||||
):
|
||||
"""Does not work for all types, in particular:
|
||||
|
||||
TODOs:
|
||||
- TSV is identified as TXT, maybe need an `.is_tsv` predicate in `_TextFileDifferentiator`
|
||||
|
||||
NOCANDOs: w/o an extension I think these are the best we can do.
|
||||
- MD is identified as TXT
|
||||
- ORG is identified as TXT
|
||||
- RST is identified as TXT
|
||||
"""
|
||||
# -- disable content-type strategy by not asserting a content_type in the call --
|
||||
# -- disable extension-mapping strategy by passing file-like object with no `.name` attribute --
|
||||
with open(example_doc_path(file_name), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
|
||||
assert detect_filetype(file=file) is expected_value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name"),
|
||||
[
|
||||
# -- `filetype` lib recognizes all these binary file-types --
|
||||
(FileType.BMP, "img/bmp_24.bmp"),
|
||||
(FileType.HEIC, "img/DA-1p.heic"),
|
||||
(FileType.JPG, "img/example.jpg"),
|
||||
(FileType.PDF, "pdf/layout-parser-paper-fast.pdf"),
|
||||
(FileType.PNG, "img/DA-1p.png"),
|
||||
(FileType.RTF, "fake-doc.rtf"),
|
||||
(FileType.TIFF, "img/layout-parser-paper-fast.tiff"),
|
||||
(FileType.WAV, "CantinaBand3.wav"),
|
||||
(FileType.ZIP, "simple.zip"),
|
||||
# -- but it doesn't recognize textual file-types at all --
|
||||
(FileType.UNK, "stanley-cups.csv"),
|
||||
(FileType.UNK, "eml/fake-email.eml"),
|
||||
(FileType.UNK, "example-10k-1p.html"),
|
||||
(FileType.UNK, "README.md"),
|
||||
(FileType.UNK, "README.org"),
|
||||
(FileType.UNK, "README.rst"),
|
||||
(FileType.UNK, "stanley-cups.tsv"),
|
||||
(FileType.UNK, "norwich-city.txt"),
|
||||
(FileType.UNK, "factbook.xml"),
|
||||
],
|
||||
)
|
||||
def test_strategy_mime_guessing_can_detect_only_binary_file_types_when_libmagic_is_unavailable(
|
||||
file_name: str, expected_value: FileType, LIBMAGIC_AVAILABLE_False: bool
|
||||
):
|
||||
"""File-type is detected using `filetype` library when libmagic is not available.
|
||||
|
||||
`filetype.guess_mime()` does a good job on binary file types (PDF, images, legacy MS-Office),
|
||||
but doesn't even try to guess textual file-types.
|
||||
"""
|
||||
# -- disable detection by extension by passing file-like object with no `.name` attribute --
|
||||
with open(example_doc_path(file_name), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
# -- simulate libmagic is not available --
|
||||
assert LIBMAGIC_AVAILABLE_False is False
|
||||
|
||||
# -- disable strategy #1 by not asserting a content_type in the call --
|
||||
file_type = detect_filetype(file=file)
|
||||
|
||||
assert file_type is expected_value
|
||||
|
||||
|
||||
def test_detect_filetype_from_file_warns_when_libmagic_is_not_installed(
|
||||
caplog: LogCaptureFixture, LIBMAGIC_AVAILABLE_False: bool
|
||||
):
|
||||
with open(example_doc_path("fake-text.txt"), "rb") as f:
|
||||
detect_filetype(file=f)
|
||||
|
||||
assert "WARNING" in caplog.text
|
||||
assert "libmagic is unavailable but assists in filetype detection. Please cons" in caplog.text
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# STRATEGY #4 - MAP FILENAME EXTENSION TO FILETYPE
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name"),
|
||||
[
|
||||
(FileType.BMP, "img/bmp_24.bmp"),
|
||||
(FileType.CSV, "stanley-cups.csv"),
|
||||
(FileType.EML, "eml/fake-email.eml"),
|
||||
(FileType.HEIC, "img/DA-1p.heic"),
|
||||
(FileType.HTML, "example-10k-1p.html"),
|
||||
(FileType.JPG, "img/example.jpg"),
|
||||
(FileType.JSON, "spring-weather.html.json"),
|
||||
(FileType.MD, "README.md"),
|
||||
(FileType.ORG, "README.org"),
|
||||
(FileType.PDF, "pdf/layout-parser-paper-fast.pdf"),
|
||||
(FileType.PNG, "img/DA-1p.png"),
|
||||
(FileType.RST, "README.rst"),
|
||||
(FileType.RTF, "fake-doc.rtf"),
|
||||
(FileType.TIFF, "img/layout-parser-paper-fast.tiff"),
|
||||
(FileType.TSV, "stanley-cups.tsv"),
|
||||
(FileType.TXT, "norwich-city.txt"),
|
||||
(FileType.WAV, "CantinaBand3.wav"),
|
||||
(FileType.XML, "factbook.xml"),
|
||||
(FileType.NDJSON, "simple.ndjson"),
|
||||
],
|
||||
)
|
||||
def test_it_detects_correct_file_type_from_extension_when_that_maps_to_a_file_type(
|
||||
file_name: str, expected_value: FileType, ctx_mime_type_: Mock
|
||||
):
|
||||
# -- disable strategy #2 by making libmagic always guess `None` --
|
||||
ctx_mime_type_.return_value = None
|
||||
|
||||
# -- disable strategy #1 by not asserting a content_type in the call --
|
||||
# -- enable strategy #3 by passing filename as source for extension --
|
||||
file_type = detect_filetype(example_doc_path(file_name))
|
||||
|
||||
# -- ctx.mime_type may be referenced multiple times, but at least once --
|
||||
ctx_mime_type_.assert_called_with()
|
||||
assert file_type is expected_value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name", "mime_type"),
|
||||
[
|
||||
(FileType.BMP, "img/bmp_24.bmp", "application/octet-stream"),
|
||||
(FileType.HEIC, "img/DA-1p.heic", "application/octet-stream"),
|
||||
],
|
||||
)
|
||||
def test_it_falls_back_to_extension_strategy_when_prior_strategies_fail(
|
||||
file_name: str, mime_type: str | None, expected_value: FileType, ctx_mime_type_: Mock
|
||||
):
|
||||
ctx_mime_type_.return_value = mime_type
|
||||
|
||||
file_type = detect_filetype(example_doc_path(file_name))
|
||||
|
||||
ctx_mime_type_.assert_called_with()
|
||||
assert file_type is expected_value
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# SPECIAL CASES
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mime_type", [FileType.XLS.mime_type, FileType.XLSX.mime_type])
|
||||
def test_it_ignores_asserted_XLS_content_type_when_file_is_CSV(mime_type: str):
|
||||
file_path = example_doc_path("stanley-cups.csv")
|
||||
assert detect_filetype(file_path, content_type=mime_type) == FileType.CSV
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mime_type", ["application/xml", "text/xml"])
|
||||
@pytest.mark.parametrize("extension", [".html", ".htm"])
|
||||
def test_it_detects_HTML_from_guessed_mime_type_ending_with_xml_and_html_extension(
|
||||
mime_type: str, extension: str, ctx_mime_type_: Mock
|
||||
):
|
||||
ctx_mime_type_.return_value = mime_type
|
||||
with open(example_doc_path("example-10k-1p.html"), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
file.name = f"a/b/page{extension}"
|
||||
|
||||
file_type = detect_filetype(file=file)
|
||||
|
||||
ctx_mime_type_.assert_called_with()
|
||||
assert file_type is FileType.HTML
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name"),
|
||||
[(FileType.NDJSON, "simple.ndjson"), (FileType.JSON, "spring-weather.html.json")],
|
||||
)
|
||||
def test_it_detects_correct_json_type_without_extension(expected_value: FileType, file_name: str):
|
||||
with open(example_doc_path(file_name), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
|
||||
filetype = detect_filetype(file=file)
|
||||
assert filetype == expected_value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("expected_value", "file_name"),
|
||||
[(FileType.NDJSON, "simple.ndjson"), (FileType.JSON, "spring-weather.html.json")],
|
||||
)
|
||||
def test_it_detects_correct_json_type_with_extension(expected_value: FileType, file_name: str):
|
||||
filetype = detect_filetype(file_path=example_doc_path(file_name))
|
||||
assert filetype == expected_value
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mime_type", "file_name"),
|
||||
[
|
||||
("text/x-script.python", "logger.py"),
|
||||
("text/x-go", "fake.go"),
|
||||
("application/x-javascript", "fake-text.txt"),
|
||||
],
|
||||
)
|
||||
def test_it_detects_TXT_for_source_code_files(mime_type: str, file_name: str, ctx_mime_type_: Mock):
|
||||
ctx_mime_type_.return_value = mime_type
|
||||
# -- disable extension-based strategy #3 --
|
||||
with open(example_doc_path(file_name), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
|
||||
file_type = detect_filetype(file=file)
|
||||
|
||||
ctx_mime_type_.assert_called_with()
|
||||
assert file_type is FileType.TXT
|
||||
|
||||
|
||||
def test_detects_TXT_from_an_unknown_guessed_text_subtype(ctx_mime_type_: Mock):
|
||||
ctx_mime_type_.return_value = "text/new-type"
|
||||
with open(example_doc_path("fake-text.txt"), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
|
||||
filetype = detect_filetype(file=file)
|
||||
|
||||
ctx_mime_type_.assert_called_with()
|
||||
assert filetype == FileType.TXT
|
||||
|
||||
|
||||
def test_detect_filetype_raises_with_neither_path_or_file_like_object_specified():
|
||||
with pytest.raises(ValueError, match="either `file_path` or `file` argument must be provided"):
|
||||
detect_filetype()
|
||||
|
||||
|
||||
def test_it_detects_EMPTY_from_file_path_to_empty_file():
|
||||
assert detect_filetype(example_doc_path("empty.txt")) == FileType.EMPTY
|
||||
|
||||
|
||||
def test_it_detects_EMPTY_from_empty_file_like_object():
|
||||
with open(example_doc_path("empty.txt"), "rb") as f:
|
||||
assert detect_filetype(file=f) == FileType.EMPTY
|
||||
|
||||
|
||||
def test_it_detect_CSV_from_path_and_file_when_content_contains_escaped_commas():
|
||||
file_path = example_doc_path("csv-with-escaped-commas.csv")
|
||||
|
||||
assert detect_filetype(file_path) == FileType.CSV
|
||||
with open(file_path, "rb") as f:
|
||||
assert detect_filetype(file=f) == FileType.CSV
|
||||
|
||||
|
||||
def test_it_detects_correct_file_type_for_custom_types(tmp_path):
|
||||
file_type = create_file_type("FOO", canonical_mime_type="application/foo", extensions=[".foo"])
|
||||
dumb_file = tmp_path / "dumb.foo"
|
||||
dumb_file.write_bytes(b"38v8df889qw8sdfj")
|
||||
assert detect_filetype(file_path=str(dumb_file), content_type="application/foo") is file_type
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# Describe `is_json_processable()`
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
def it_affirms_JSON_is_array_of_objects_from_a_file_path():
|
||||
assert is_json_processable(example_doc_path("simple.json")) is True
|
||||
|
||||
|
||||
def and_it_affirms_JSON_is_NOT_an_array_of_objects_from_a_file_path():
|
||||
assert is_json_processable(example_doc_path("not-unstructured-payload.json")) is False
|
||||
|
||||
|
||||
def it_affirms_JSON_is_array_of_objects_from_a_file_like_object_open_for_reading_bytes():
|
||||
with open(example_doc_path("simple.json"), "rb") as f:
|
||||
assert is_json_processable(file=f) is True
|
||||
|
||||
|
||||
def and_it_affirms_JSON_is_NOT_an_array_of_objects_from_a_file_like_object_open_for_reading_bytes():
|
||||
with open(example_doc_path("not-unstructured-payload.json"), "rb") as f:
|
||||
assert is_json_processable(file=f) is False
|
||||
|
||||
|
||||
def it_affirms_JSON_is_array_of_objects_from_text():
|
||||
with open(example_doc_path("simple.json")) as f:
|
||||
text = f.read()
|
||||
assert is_json_processable(file_text=text) is True
|
||||
|
||||
|
||||
def and_it_affirms_JSON_is_NOT_an_array_of_objects_from_text():
|
||||
with open(example_doc_path("not-unstructured-payload.json")) as f:
|
||||
text = f.read()
|
||||
assert is_json_processable(file_text=text) is False
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# MODULE-LEVEL FIXTURES
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def LIBMAGIC_AVAILABLE_False():
|
||||
with patch("unstructured.file_utils.filetype.LIBMAGIC_AVAILABLE", False) as m:
|
||||
yield m
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def ctx_mime_type_(request: FixtureRequest):
|
||||
return property_mock(request, _FileTypeDetectionContext, "mime_type")
|
||||
|
||||
|
||||
# ================================================================================================
|
||||
# UNIT-TESTS
|
||||
# ================================================================================================
|
||||
|
||||
|
||||
class Describe_FileTypeDetectionContext:
|
||||
"""Unit-test suite for `unstructured.file_utils.filetype._FileTypeDetectionContext`."""
|
||||
|
||||
# -- .new() -------------------------------------------------
|
||||
|
||||
def it_provides_a_validating_alternate_constructor(self):
|
||||
ctx = _FileTypeDetectionContext.new(
|
||||
file_path=example_doc_path("simple.docx"),
|
||||
file=None,
|
||||
encoding="utf-8",
|
||||
content_type="text/plain",
|
||||
metadata_file_path="a/b/foo.bar",
|
||||
)
|
||||
assert isinstance(ctx, _FileTypeDetectionContext)
|
||||
|
||||
def and_the_validating_constructor_raises_on_an_invalid_context(self):
|
||||
with pytest.raises(ValueError, match="either `file_path` or `file` argument must be pro"):
|
||||
_FileTypeDetectionContext.new(
|
||||
file_path=None,
|
||||
file=None,
|
||||
encoding=None,
|
||||
content_type=None,
|
||||
metadata_file_path=None,
|
||||
)
|
||||
|
||||
# -- .content_type ------------------------------------------
|
||||
|
||||
def it_knows_the_content_type_asserted_by_the_caller(self):
|
||||
assert _FileTypeDetectionContext(content_type="TEXT/hTmL").content_type == "text/html"
|
||||
|
||||
# -- .encoding ----------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("encoding", "expected_value"),
|
||||
[
|
||||
("utf-8", "utf-8"),
|
||||
("UTF_8", "utf-8"),
|
||||
("UTF_16LE", "utf-16le"),
|
||||
("ISO_8859_6_I", "iso-8859-6"),
|
||||
# -- default value is utf-8 --
|
||||
(None, "utf-8"),
|
||||
],
|
||||
)
|
||||
def it_knows_the_encoding_asserted_by_the_caller_and_normalizes_it(
|
||||
self, encoding: str | None, expected_value: str
|
||||
):
|
||||
assert _FileTypeDetectionContext(encoding=encoding).encoding == expected_value
|
||||
|
||||
# -- .extension ---------------------------------------------
|
||||
|
||||
def it_derives_the_filename_extension_from_the_file_path_when_one_is_provided(self):
|
||||
ctx = _FileTypeDetectionContext(file_path=example_doc_path("simple.docx"))
|
||||
assert ctx.extension == ".docx"
|
||||
|
||||
def and_it_derives_the_extension_from_a_file_opened_from_a_path(self):
|
||||
with open(example_doc_path("picture.pptx"), "rb") as f:
|
||||
assert _FileTypeDetectionContext(file=f).extension == ".pptx"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_name",
|
||||
[
|
||||
# -- case 1: file-like object has no `.name` attribute
|
||||
None,
|
||||
# -- case 2: file-like object has `.name` attribute but it's value is the empty string
|
||||
"",
|
||||
],
|
||||
)
|
||||
def and_it_derives_the_extension_from_metadata_file_path_when_file_object_has_no_name(
|
||||
self, file_name: str | None
|
||||
):
|
||||
with open(example_doc_path("ideas-page.html"), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
if file_name is not None:
|
||||
file.name = file_name
|
||||
|
||||
ctx = _FileTypeDetectionContext(file=file, metadata_file_path="a/b/c.html")
|
||||
|
||||
assert ctx.extension == ".html"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_name",
|
||||
[
|
||||
# -- case 1: file-like object has no `.name` attribute
|
||||
None,
|
||||
# -- case 2: file-like object has `.name` attribute but it's value is the empty string
|
||||
"",
|
||||
],
|
||||
)
|
||||
def and_it_returns_the_empty_string_as_the_extension_when_there_are_no_file_name_sources(
|
||||
self, file_name: str | None
|
||||
):
|
||||
with open(example_doc_path("ideas-page.html"), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
if file_name is not None:
|
||||
file.name = file_name
|
||||
|
||||
assert _FileTypeDetectionContext(file=file).extension == ""
|
||||
|
||||
# -- .file_head ---------------------------------------------
|
||||
|
||||
def it_grabs_the_first_8k_bytes_of_the_file_for_use_by_magic(self):
|
||||
ctx = _FileTypeDetectionContext(file_path=example_doc_path("norwich-city.txt"))
|
||||
|
||||
head = ctx.file_head
|
||||
|
||||
assert isinstance(head, bytes)
|
||||
assert len(head) == 8192
|
||||
assert head.startswith(b"Iwan Roberts\nRoberts celebrating after")
|
||||
|
||||
# -- .file_path ---------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize("file_path", [None, "a/b/c.pdf"])
|
||||
def it_knows_the_file_path_provided_by_the_caller(self, file_path: str | None):
|
||||
assert _FileTypeDetectionContext(file_path=file_path).file_path == file_path
|
||||
|
||||
# -- .has_code_mime_type ------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mime_type", "expected_value"),
|
||||
[
|
||||
("text/plain", False),
|
||||
("text/x-csharp", True),
|
||||
("text/x-go", True),
|
||||
("text/x-java", True),
|
||||
("text/x-python", True),
|
||||
("application/xml", False),
|
||||
(None, False),
|
||||
],
|
||||
)
|
||||
def it_knows_whether_its_mime_type_indicates_programming_language_source_code(
|
||||
self, mime_type_prop_: Mock, mime_type: str | None, expected_value: bool
|
||||
):
|
||||
mime_type_prop_.return_value = mime_type
|
||||
assert _FileTypeDetectionContext().has_code_mime_type is expected_value
|
||||
|
||||
# -- .is_zipfile --------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_name", "expected_value"),
|
||||
[
|
||||
("README.md", False),
|
||||
("emoji.xlsx", True),
|
||||
("simple.doc", False),
|
||||
("simple.docx", True),
|
||||
("simple.odt", True),
|
||||
("simple.zip", True),
|
||||
("winter-sports.epub", True),
|
||||
],
|
||||
)
|
||||
def it_knows_whether_it_is_a_zipfile(self, file_name: str, expected_value: bool):
|
||||
assert _FileTypeDetectionContext(example_doc_path(file_name)).is_zipfile is expected_value
|
||||
|
||||
# -- .mime_type ---------------------------------------------
|
||||
|
||||
def it_provides_the_MIME_type_detected_by_libmagic_from_a_file_path(self):
|
||||
ctx = _FileTypeDetectionContext(file_path=example_doc_path("norwich-city.txt"))
|
||||
assert ctx.mime_type == "text/plain"
|
||||
|
||||
def and_it_provides_the_MIME_type_from_path_using_filetype_lib_when_magic_is_unavailable(self):
|
||||
with patch("unstructured.file_utils.filetype.LIBMAGIC_AVAILABLE", False):
|
||||
ctx = _FileTypeDetectionContext(file_path=example_doc_path("simple.doc"))
|
||||
assert ctx.mime_type == "application/msword"
|
||||
|
||||
def but_it_warns_to_install_libmagic_when_the_filetype_lib_cannot_detect_the_MIME_type(
|
||||
self, caplog: LogCaptureFixture
|
||||
):
|
||||
with patch("unstructured.file_utils.filetype.LIBMAGIC_AVAILABLE", False):
|
||||
ctx = _FileTypeDetectionContext(file_path=example_doc_path("norwich-city.txt"))
|
||||
assert ctx.mime_type is None
|
||||
assert "WARNING" in caplog.text
|
||||
assert "libmagic is unavailable" in caplog.text
|
||||
assert "consider installing libmagic" in caplog.text
|
||||
|
||||
def it_provides_the_MIME_type_detected_by_libmagic_from_a_file_like_object(self):
|
||||
with open(example_doc_path("norwich-city.txt"), "rb") as f:
|
||||
ctx = _FileTypeDetectionContext(file=f)
|
||||
assert ctx.mime_type == "text/plain"
|
||||
|
||||
def and_it_provides_the_MIME_type_from_file_using_filetype_lib_when_magic_is_unavailable(self):
|
||||
with patch("unstructured.file_utils.filetype.LIBMAGIC_AVAILABLE", False):
|
||||
file_path = example_doc_path("simple.doc")
|
||||
with open(file_path, "rb") as f:
|
||||
ctx = _FileTypeDetectionContext(file=f)
|
||||
assert ctx.mime_type == "application/msword"
|
||||
|
||||
# -- .open() ------------------------------------------------
|
||||
|
||||
def it_provides_transparent_access_to_the_source_file_when_it_is_a_file_like_object(self):
|
||||
with open(example_doc_path("norwich-city.txt"), "rb") as f:
|
||||
ctx = _FileTypeDetectionContext(file=f)
|
||||
with ctx.open() as file:
|
||||
assert file is f
|
||||
assert file.read(38) == b"Iwan Roberts\nRoberts celebrating after"
|
||||
|
||||
def it_provides_transparent_access_to_the_source_file_when_it_is_a_file_path(self):
|
||||
ctx = _FileTypeDetectionContext(file_path=example_doc_path("norwich-city.txt"))
|
||||
with ctx.open() as file:
|
||||
assert file.read(38) == b"Iwan Roberts\nRoberts celebrating after"
|
||||
|
||||
# -- .text_head ---------------------------------------------
|
||||
|
||||
def it_grabs_the_first_4k_chars_from_file_path_for_textual_type_differentiation(self):
|
||||
ctx = _FileTypeDetectionContext(file_path=example_doc_path("norwich-city.txt"))
|
||||
|
||||
text_head = ctx.text_head
|
||||
|
||||
assert isinstance(text_head, str)
|
||||
assert len(text_head) == 4096
|
||||
assert text_head.startswith("Iwan Roberts\nRoberts celebrating after")
|
||||
|
||||
def and_it_uses_character_detection_to_correct_a_wrong_encoding_arg_for_file_path(self):
|
||||
ctx = _FileTypeDetectionContext(
|
||||
file_path=example_doc_path("norwich-city.txt"), encoding="utf_32_be"
|
||||
)
|
||||
|
||||
text_head = ctx.text_head
|
||||
|
||||
assert isinstance(text_head, str)
|
||||
assert len(text_head) == 4096
|
||||
assert text_head.startswith("Iwan Roberts\nRoberts celebrating after")
|
||||
|
||||
def but_not_to_correct_a_wrong_encoding_arg_for_a_file_like_object_open_in_binary_mode(self):
|
||||
"""Fails silently in this case, returning empty string."""
|
||||
with open(example_doc_path("norwich-city.txt"), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
ctx = _FileTypeDetectionContext(file=file, encoding="utf_32_be")
|
||||
|
||||
text_head = ctx.text_head
|
||||
|
||||
assert text_head == ""
|
||||
|
||||
def and_it_grabs_the_first_4k_chars_from_binary_file_for_textual_type_differentiation(self):
|
||||
with open(example_doc_path("norwich-city.txt"), "rb") as f:
|
||||
ctx = _FileTypeDetectionContext(file=f)
|
||||
|
||||
text_head = ctx.text_head
|
||||
|
||||
assert isinstance(text_head, str)
|
||||
# -- some characters consume multiple bytes, so shorter than 4096 --
|
||||
assert len(text_head) == 4063
|
||||
assert text_head.startswith("Iwan Roberts\nRoberts celebrating after")
|
||||
|
||||
def and_it_grabs_the_first_4k_chars_from_text_file_for_textual_type_differentiation(self):
|
||||
"""Not a documented behavior to accept IO[str], but support is implemented."""
|
||||
with open(example_doc_path("norwich-city.txt")) as f:
|
||||
ctx = _FileTypeDetectionContext(file=f) # pyright: ignore[reportArgumentType]
|
||||
|
||||
text_head = ctx.text_head
|
||||
|
||||
assert isinstance(text_head, str)
|
||||
assert len(text_head) == 4096
|
||||
assert text_head.startswith("Iwan Roberts\nRoberts celebrating after")
|
||||
|
||||
def it_accommodates_a_utf_32_encoded_file_path(self):
|
||||
ctx = _FileTypeDetectionContext(example_doc_path("fake-text-utf-32.txt"))
|
||||
|
||||
text_head = ctx.text_head
|
||||
|
||||
assert isinstance(text_head, str)
|
||||
# -- test document is short --
|
||||
assert len(text_head) == 188
|
||||
assert text_head.startswith("This is a test document to use for unit tests.\n\n Doyle")
|
||||
|
||||
# TODO: this fails because `.text_head` ignores decoding errors on a file open for binary
|
||||
# reading. Probably better if it used chardet in that case as it does for a file-path.
|
||||
@pytest.mark.xfail(reason="WIP", raises=AssertionError, strict=True)
|
||||
def and_it_accommodates_a_utf_32_encoded_file_like_object(self):
|
||||
with open(example_doc_path("fake-text-utf-32.txt"), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
ctx = _FileTypeDetectionContext(file=file)
|
||||
|
||||
text_head = ctx.text_head
|
||||
|
||||
assert isinstance(text_head, str)
|
||||
# -- test document is short --
|
||||
assert len(text_head) == 188
|
||||
assert text_head.startswith("This is a test document to use for unit tests.\n\n Doyle")
|
||||
|
||||
# -- .validate() --------------------------------------------
|
||||
|
||||
def it_raises_when_no_file_exists_at_the_specified_file_path(self):
|
||||
with pytest.raises(FileNotFoundError, match="no such file a/b/c.foo"):
|
||||
_FileTypeDetectionContext(file_path="a/b/c.foo")._validate()
|
||||
|
||||
def it_raises_when_neither_file_path_nor_file_is_provided(self):
|
||||
with pytest.raises(ValueError, match="either `file_path` or `file` argument must be pro"):
|
||||
_FileTypeDetectionContext()._validate()
|
||||
|
||||
# -- fixtures --------------------------------------------------------------------------------
|
||||
|
||||
@pytest.fixture
|
||||
def mime_type_prop_(self, request: FixtureRequest):
|
||||
return property_mock(request, _FileTypeDetectionContext, "mime_type")
|
||||
|
||||
|
||||
class Describe_OleFileDetector:
|
||||
"""Unit-test suite for `unstructured.file_utils.filetype._OleFileDetector`."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_name", "expected_value"),
|
||||
[
|
||||
("simple.doc", FileType.DOC),
|
||||
("fake-power-point.ppt", FileType.PPT),
|
||||
("tests-example.xls", FileType.XLS),
|
||||
("fake-email.msg", FileType.MSG),
|
||||
("README.org", None),
|
||||
],
|
||||
)
|
||||
def it_distinguishes_the_file_type_of_applicable_CFB_files(
|
||||
self, file_name: str, expected_value: FileType | None
|
||||
):
|
||||
# -- no file-name available, just to make sure we're not relying on an extension --
|
||||
with open(example_doc_path(file_name), "rb") as f:
|
||||
file = io.BytesIO(f.read())
|
||||
ctx = _FileTypeDetectionContext(file=file)
|
||||
|
||||
assert _OleFileDetector.file_type(ctx) is expected_value
|
||||
|
||||
|
||||
class Describe_TextFileDifferentiator:
|
||||
"""Unit-test suite for `unstructured.file_utils.filetype._TextFileDifferentiator`."""
|
||||
|
||||
# -- .applies() ---------------------------------------------
|
||||
|
||||
def it_provides_a_qualifying_alternate_constructor_which_constructs_when_applicable(self):
|
||||
"""The constructor determines whether this differentiator is applicable.
|
||||
|
||||
It returns an instance only when differentiating a text file-type is required, which it can
|
||||
judge from the context (`ctx`).
|
||||
"""
|
||||
ctx = _FileTypeDetectionContext(example_doc_path("norwich-city.txt"))
|
||||
|
||||
differentiator = _TextFileDifferentiator.applies(ctx)
|
||||
|
||||
assert isinstance(differentiator, _TextFileDifferentiator)
|
||||
|
||||
def and_it_returns_None_when_text_differentiation_does_not_apply_to_the_detection_context(self):
|
||||
ctx = _FileTypeDetectionContext(example_doc_path("simple.docx"))
|
||||
assert _TextFileDifferentiator.applies(ctx) is None
|
||||
|
||||
# -- ._is_csv -----------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("content", "expected_value"),
|
||||
[
|
||||
# -- no commas, too few lines --
|
||||
(b"d\xe2\x80", False),
|
||||
(b'[{"key": "value"}]', False),
|
||||
# -- at least a header and one data row, at least two columns --
|
||||
(b"column1,column2,column3\nvalue1,value2,value3\n", True),
|
||||
# -- no content --
|
||||
(b"", False),
|
||||
],
|
||||
)
|
||||
def it_distinguishes_a_CSV_file_from_other_text_files(
|
||||
self, content: bytes, expected_value: bool
|
||||
):
|
||||
ctx = _FileTypeDetectionContext(file=io.BytesIO(content))
|
||||
differentiator = _TextFileDifferentiator(ctx)
|
||||
|
||||
assert differentiator._is_csv is expected_value
|
||||
|
||||
# -- ._is_eml -----------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_name", "expected_value"), [("fake-email.eml", True), ("norwich-city.txt", False)]
|
||||
)
|
||||
def it_distinguishes_an_EML_file_from_other_text_files(
|
||||
self, file_name: str, expected_value: bool
|
||||
):
|
||||
ctx = _FileTypeDetectionContext(example_doc_path(file_name))
|
||||
assert _TextFileDifferentiator(ctx)._is_eml is expected_value
|
||||
|
||||
# -- ._is_json ----------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("content", "expected_value"),
|
||||
[
|
||||
(b"d\xe2\x80", False),
|
||||
(b'[{"key": "value"}]', True),
|
||||
(b"", False),
|
||||
# -- valid JSON, but not for our purposes --
|
||||
(b'"This is not a JSON"', False),
|
||||
],
|
||||
)
|
||||
def it_distinguishes_a_JSON_file_from_other_text_files(
|
||||
self, content: bytes, expected_value: bool
|
||||
):
|
||||
ctx = _FileTypeDetectionContext(file=io.BytesIO(content))
|
||||
differentiator = _TextFileDifferentiator(ctx)
|
||||
|
||||
assert differentiator._is_json is expected_value
|
||||
|
||||
|
||||
class Describe_ZipFileDetector:
|
||||
"""Unit-test suite for `unstructured.file_utils.filetype._ZipFileDetector`."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_name", "expected_value"),
|
||||
[
|
||||
("simple.docx", FileType.DOCX),
|
||||
("winter-sports.epub", FileType.EPUB),
|
||||
("simple.odt", FileType.ODT),
|
||||
("picture.pptx", FileType.PPTX),
|
||||
("vodafone.xlsx", FileType.XLSX),
|
||||
("simple.zip", FileType.ZIP),
|
||||
("README.org", None),
|
||||
],
|
||||
)
|
||||
def it_distinguishes_the_file_type_of_applicable_zip_files(
|
||||
self, file_name: str, expected_value: FileType | None
|
||||
):
|
||||
ctx = _FileTypeDetectionContext(example_doc_path(file_name))
|
||||
assert _ZipFileDetector.file_type(ctx) is expected_value
|
||||
|
||||
|
||||
def test_mimetype_magic_detection_is_used_before_filename_when_filetype_is_detected_for_json():
|
||||
json_bytes = json.dumps([{"example": "data"}]).encode("utf-8")
|
||||
|
||||
file_buffer = io.BytesIO(json_bytes)
|
||||
predicted_type = detect_filetype(file=file_buffer, metadata_file_path="filename.pdf")
|
||||
assert predicted_type == FileType.JSON
|
||||
|
||||
file_buffer.name = "filename.pdf"
|
||||
predicted_type = detect_filetype(file=file_buffer)
|
||||
assert predicted_type == FileType.JSON
|
||||
|
||||
|
||||
def test_mimetype_magic_detection_is_used_before_filename_when_filetype_is_detected_for_ndjson():
|
||||
data = [{"example": "data1"}, {"example": "data2"}, {"example": "data3"}]
|
||||
ndjson_string = "\n".join(json.dumps(item) for item in data) + "\n"
|
||||
ndjson_bytes = ndjson_string.encode("utf-8")
|
||||
|
||||
file_buffer = io.BytesIO(ndjson_bytes)
|
||||
predicted_type = detect_filetype(file=file_buffer, metadata_file_path="filename.pdf")
|
||||
assert predicted_type == FileType.NDJSON
|
||||
|
||||
file_buffer.name = "filename.pdf"
|
||||
predicted_type = detect_filetype(file=file_buffer)
|
||||
assert predicted_type == FileType.NDJSON
|
||||
|
||||
|
||||
def test_json_content_type_is_disambiguated_for_ndjson():
|
||||
data = [{"example": "data1"}, {"example": "data2"}, {"example": "data3"}]
|
||||
ndjson_string = "\n".join(json.dumps(item) for item in data) + "\n"
|
||||
ndjson_bytes = ndjson_string.encode("utf-8")
|
||||
|
||||
file_buffer = io.BytesIO(ndjson_bytes)
|
||||
predicted_type = detect_filetype(
|
||||
file=file_buffer, metadata_file_path="filename.pdf", content_type="application/json"
|
||||
)
|
||||
assert predicted_type == FileType.NDJSON
|
||||
|
||||
file_buffer.name = "filename.pdf"
|
||||
predicted_type = detect_filetype(file=file_buffer, content_type="application/json")
|
||||
assert predicted_type == FileType.NDJSON
|
||||
|
||||
|
||||
def test_office_files_when_document_archive_has_non_standard_prefix():
|
||||
predicted_type = detect_filetype(
|
||||
file_path=input_path("file_type/test_document_from_office365.docx")
|
||||
)
|
||||
assert predicted_type == FileType.DOCX
|
||||
@@ -0,0 +1,245 @@
|
||||
"""Test suite for `unstructured.file_utils.filetype`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from unstructured.file_utils.model import FileType, create_file_type, register_partitioner
|
||||
|
||||
|
||||
class DescribeFileType:
|
||||
"""Unit-test suite for `unstructured.file_utils.model.Filetype`."""
|
||||
|
||||
# -- .__lt__() ----------------------------------------------
|
||||
|
||||
def it_is_a_collection_ordered_by_name_and_can_be_sorted(self):
|
||||
"""FileType is a total order on name, e.g. FileType.A < FileType.B."""
|
||||
assert FileType.EML < FileType.HTML < FileType.XML
|
||||
|
||||
# -- .from_extension() --------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("ext", "file_type"),
|
||||
[
|
||||
(".bmp", FileType.BMP),
|
||||
(".html", FileType.HTML),
|
||||
(".eml", FileType.EML),
|
||||
(".p7s", FileType.EML),
|
||||
(".java", FileType.TXT),
|
||||
],
|
||||
)
|
||||
def it_can_recognize_a_file_type_from_an_extension(self, ext: str, file_type: FileType | None):
|
||||
assert FileType.from_extension(ext) is file_type
|
||||
|
||||
@pytest.mark.parametrize("ext", [".foobar", ".xyz", ".mdx", "", ".", None])
|
||||
def but_not_when_that_extension_is_empty_or_None_or_not_registered(self, ext: str | None):
|
||||
assert FileType.from_extension(ext) is None
|
||||
|
||||
# -- .from_mime_type() --------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("mime_type", "file_type"),
|
||||
[
|
||||
("image/bmp", FileType.BMP),
|
||||
("text/x-csv", FileType.CSV),
|
||||
("application/msword", FileType.DOC),
|
||||
("message/rfc822", FileType.EML),
|
||||
("text/plain", FileType.TXT),
|
||||
("text/yaml", FileType.TXT),
|
||||
("application/xml", FileType.XML),
|
||||
("text/xml", FileType.XML),
|
||||
("inode/x-empty", FileType.EMPTY),
|
||||
],
|
||||
)
|
||||
def it_can_recognize_a_file_type_from_a_mime_type(
|
||||
self, mime_type: str, file_type: FileType | None
|
||||
):
|
||||
assert FileType.from_mime_type(mime_type) is file_type
|
||||
|
||||
@pytest.mark.parametrize("mime_type", ["text/css", "image/gif", "audio/mpeg", "foo/bar", None])
|
||||
def but_not_when_that_mime_type_is_not_registered_by_a_file_type_or_None(
|
||||
self, mime_type: str | None
|
||||
):
|
||||
assert FileType.from_mime_type(mime_type) is None
|
||||
|
||||
# -- .extra_name --------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_type", "expected_value"),
|
||||
[
|
||||
(FileType.BMP, "image"),
|
||||
(FileType.DOC, "doc"),
|
||||
(FileType.DOCX, "docx"),
|
||||
(FileType.EML, None),
|
||||
(FileType.EMPTY, None),
|
||||
(FileType.MSG, "msg"),
|
||||
(FileType.PDF, "pdf"),
|
||||
(FileType.XLS, "xlsx"),
|
||||
(FileType.UNK, None),
|
||||
(FileType.WAV, None),
|
||||
(FileType.ZIP, None),
|
||||
],
|
||||
)
|
||||
def and_it_knows_which_pip_extra_needs_to_be_installed_to_get_those_dependencies(
|
||||
self, file_type: FileType, expected_value: str | None
|
||||
):
|
||||
assert file_type.extra_name == expected_value
|
||||
|
||||
# -- .importable_package_dependencies -----------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_type", "expected_value"),
|
||||
[
|
||||
(FileType.BMP, ("unstructured_inference",)),
|
||||
(FileType.CSV, ("pandas",)),
|
||||
(FileType.DOC, ("docx",)),
|
||||
(FileType.EMPTY, ()),
|
||||
(FileType.HTML, ()),
|
||||
(FileType.ODT, ("docx", "pypandoc")),
|
||||
(FileType.PDF, ("pdf2image", "pdfminer", "PIL")),
|
||||
(FileType.UNK, ()),
|
||||
(FileType.WAV, ()),
|
||||
(FileType.ZIP, ()),
|
||||
],
|
||||
)
|
||||
def it_knows_which_importable_packages_its_partitioner_depends_on(
|
||||
self, file_type: FileType, expected_value: tuple[str, ...]
|
||||
):
|
||||
assert file_type.importable_package_dependencies == expected_value
|
||||
|
||||
# -- .is_partitionable --------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_type", "expected_value"),
|
||||
[
|
||||
(FileType.BMP, True),
|
||||
(FileType.CSV, True),
|
||||
(FileType.DOC, True),
|
||||
(FileType.EML, True),
|
||||
(FileType.JPG, True),
|
||||
(FileType.PDF, True),
|
||||
(FileType.PPTX, True),
|
||||
(FileType.WAV, False),
|
||||
(FileType.ZIP, False),
|
||||
(FileType.EMPTY, False),
|
||||
(FileType.UNK, False),
|
||||
],
|
||||
)
|
||||
def it_knows_whether_files_of_its_type_are_directly_partitionable(
|
||||
self, file_type: FileType, expected_value: str
|
||||
):
|
||||
assert file_type.is_partitionable is expected_value
|
||||
|
||||
# -- .mime_type ---------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_type", "mime_type"),
|
||||
[
|
||||
(FileType.BMP, "image/bmp"),
|
||||
(FileType.CSV, "text/csv"),
|
||||
(FileType.DOC, "application/msword"),
|
||||
(FileType.EML, "message/rfc822"),
|
||||
(FileType.HTML, "text/html"),
|
||||
(FileType.JPG, "image/jpeg"),
|
||||
(FileType.PDF, "application/pdf"),
|
||||
(FileType.TXT, "text/plain"),
|
||||
(FileType.XML, "application/xml"),
|
||||
(FileType.EMPTY, "inode/x-empty"),
|
||||
(FileType.UNK, "application/octet-stream"),
|
||||
],
|
||||
)
|
||||
def it_knows_its_canonical_MIME_type(self, file_type: FileType, mime_type: str):
|
||||
assert file_type.mime_type == mime_type
|
||||
|
||||
# -- .partitioner_function_name -----------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_type", "expected_value"),
|
||||
[
|
||||
(FileType.BMP, "partition_image"),
|
||||
(FileType.CSV, "partition_csv"),
|
||||
(FileType.DOC, "partition_doc"),
|
||||
(FileType.DOCX, "partition_docx"),
|
||||
(FileType.JPG, "partition_image"),
|
||||
(FileType.PNG, "partition_image"),
|
||||
(FileType.TIFF, "partition_image"),
|
||||
],
|
||||
)
|
||||
def it_knows_its_partitioner_function_name(self, file_type: FileType, expected_value: str):
|
||||
assert file_type.partitioner_function_name == expected_value
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_type", [FileType.WAV, FileType.ZIP, FileType.EMPTY, FileType.UNK]
|
||||
)
|
||||
def but_it_raises_on_partitioner_function_name_access_when_the_file_type_is_not_partitionable(
|
||||
self, file_type: FileType
|
||||
):
|
||||
with pytest.raises(ValueError, match="`.partitioner_function_name` is undefined because "):
|
||||
file_type.partitioner_function_name
|
||||
|
||||
# -- .partitioner_module_qname ------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_type", "expected_value"),
|
||||
[
|
||||
(FileType.BMP, "unstructured.partition.image"),
|
||||
(FileType.CSV, "unstructured.partition.csv"),
|
||||
(FileType.DOC, "unstructured.partition.doc"),
|
||||
(FileType.DOCX, "unstructured.partition.docx"),
|
||||
(FileType.JPG, "unstructured.partition.image"),
|
||||
(FileType.PNG, "unstructured.partition.image"),
|
||||
(FileType.TIFF, "unstructured.partition.image"),
|
||||
],
|
||||
)
|
||||
def it_knows_the_fully_qualified_name_of_its_partitioner_module(
|
||||
self, file_type: FileType, expected_value: str
|
||||
):
|
||||
assert file_type.partitioner_module_qname == expected_value
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file_type", [FileType.WAV, FileType.ZIP, FileType.EMPTY, FileType.UNK]
|
||||
)
|
||||
def but_it_raises_on_partitioner_module_qname_access_when_the_file_type_is_not_partitionable(
|
||||
self, file_type: FileType
|
||||
):
|
||||
with pytest.raises(ValueError, match="`.partitioner_module_qname` is undefined because "):
|
||||
file_type.partitioner_module_qname
|
||||
|
||||
# -- .partitioner_shortname ---------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("file_type", "expected_value"),
|
||||
[
|
||||
(FileType.BMP, "image"),
|
||||
(FileType.CSV, "csv"),
|
||||
(FileType.DOC, "doc"),
|
||||
(FileType.DOCX, "docx"),
|
||||
(FileType.JPG, "image"),
|
||||
(FileType.PNG, "image"),
|
||||
(FileType.TIFF, "image"),
|
||||
(FileType.XLS, "xlsx"),
|
||||
(FileType.XLSX, "xlsx"),
|
||||
],
|
||||
)
|
||||
def it_provides_access_to_the_partitioner_shortname(
|
||||
self, file_type: FileType, expected_value: str
|
||||
):
|
||||
assert file_type.partitioner_shortname == expected_value
|
||||
|
||||
|
||||
def test_create_file_type():
|
||||
file_type = create_file_type("FOO", canonical_mime_type="application/foo", extensions=[".foo"])
|
||||
|
||||
assert FileType.from_extension(".foo") is file_type
|
||||
assert FileType.from_mime_type("application/foo") is file_type
|
||||
|
||||
|
||||
def test_register_partitioner():
|
||||
file_type = create_file_type("FOO", canonical_mime_type="application/foo", extensions=[".foo"])
|
||||
|
||||
@register_partitioner(file_type)
|
||||
def partition_foo():
|
||||
pass
|
||||
|
||||
assert file_type.partitioner_function_name == "partition_foo"
|
||||
assert file_type.partitioner_module_qname == "test_unstructured.file_utils.test_model"
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,113 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.metrics.element_type import (
|
||||
FrequencyDict,
|
||||
calculate_element_type_percent_match,
|
||||
get_element_type_frequency,
|
||||
)
|
||||
from unstructured.partition.auto import partition
|
||||
from unstructured.staging.base import elements_to_json
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "frequency"),
|
||||
[
|
||||
(
|
||||
"fake-email.txt",
|
||||
{
|
||||
("NarrativeText", None): 1,
|
||||
("UncategorizedText", None): 1,
|
||||
("ListItem", 1): 2,
|
||||
},
|
||||
),
|
||||
(
|
||||
"sample-presentation.pptx",
|
||||
{
|
||||
("Title", 0): 4,
|
||||
("Title", 1): 1,
|
||||
("NarrativeText", 0): 3,
|
||||
("PageBreak", None): 3,
|
||||
("ListItem", 0): 6,
|
||||
("ListItem", 1): 6,
|
||||
("ListItem", 2): 3,
|
||||
("Table", None): 1,
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_element_type_frequency(filename: str, frequency: dict[tuple[str, int | None], int]):
|
||||
elements = partition(example_doc_path(filename))
|
||||
elements_freq = get_element_type_frequency(elements_to_json(elements))
|
||||
assert elements_freq == frequency
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "expected_frequency", "percent_matched"),
|
||||
[
|
||||
(
|
||||
"fake-email.txt",
|
||||
{
|
||||
("UncategorizedText", None): 1,
|
||||
("ListItem", 1): 2,
|
||||
("NarrativeText", None): 2,
|
||||
},
|
||||
(0.8, 0.8, 0.80),
|
||||
),
|
||||
(
|
||||
"sample-presentation.pptx",
|
||||
{
|
||||
("Title", 0): 3,
|
||||
("Title", 1): 1,
|
||||
("NarrativeText", None): 1,
|
||||
("NarrativeText", 0): 3,
|
||||
("ListItem", 0): 6,
|
||||
("ListItem", 1): 6,
|
||||
("ListItem", 2): 3,
|
||||
("Table", None): 1,
|
||||
},
|
||||
(0.96, 0.96, 0.96),
|
||||
),
|
||||
(
|
||||
"handbook-1p.docx",
|
||||
{
|
||||
("Header", None): 1,
|
||||
("UncategorizedText", 0): 6,
|
||||
("ListItem", 3): 3,
|
||||
("NarrativeText", 0): 7,
|
||||
("Footer", None): 1,
|
||||
},
|
||||
(0.78, 0.72, 0.81),
|
||||
),
|
||||
(
|
||||
"handbook-1p.docx",
|
||||
{
|
||||
("Header", None): 1,
|
||||
("UncategorizedText", 0): 6,
|
||||
("NarrativeText", 0): 7,
|
||||
("PageBreak", None): 1,
|
||||
("Footer", None): 1,
|
||||
},
|
||||
(0.94, 0.88, 0.98),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_calculate_element_type_percent_match(
|
||||
filename: str, expected_frequency: FrequencyDict, percent_matched: tuple[float, float, float]
|
||||
):
|
||||
elements = partition(example_doc_path(filename))
|
||||
elements_frequency = get_element_type_frequency(elements_to_json(elements))
|
||||
assert (
|
||||
round(calculate_element_type_percent_match(elements_frequency, expected_frequency), 2)
|
||||
== percent_matched[0]
|
||||
)
|
||||
assert (
|
||||
round(calculate_element_type_percent_match(elements_frequency, expected_frequency, 0.0), 2)
|
||||
== percent_matched[1]
|
||||
)
|
||||
assert (
|
||||
round(calculate_element_type_percent_match(elements_frequency, expected_frequency, 0.8), 2)
|
||||
== percent_matched[2]
|
||||
)
|
||||
@@ -0,0 +1,595 @@
|
||||
import os
|
||||
import pathlib
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from unstructured.metrics.evaluate import (
|
||||
ElementTypeMetricsCalculator,
|
||||
TableStructureMetricsCalculator,
|
||||
TextExtractionMetricsCalculator,
|
||||
filter_metrics,
|
||||
get_mean_grouping,
|
||||
)
|
||||
|
||||
is_in_docker = os.path.exists("/.dockerenv")
|
||||
|
||||
EXAMPLE_DOCS_DIRECTORY = os.path.join(
|
||||
pathlib.Path(__file__).parent.resolve(), "..", "..", "example-docs"
|
||||
)
|
||||
TESTING_FILE_DIR = os.path.join(EXAMPLE_DOCS_DIRECTORY, "test_evaluate_files")
|
||||
|
||||
UNSTRUCTURED_OUTPUT_DIRNAME = "unstructured_output"
|
||||
GOLD_CCT_DIRNAME = "gold_standard_cct"
|
||||
GOLD_ELEMENT_TYPE_DIRNAME = "gold_standard_element_type"
|
||||
GOLD_TABLE_STRUCTURE_DIRNAME = "gold_standard_table_structure"
|
||||
UNSTRUCTURED_CCT_DIRNAME = "unstructured_output_cct"
|
||||
UNSTRUCTURED_TABLE_STRUCTURE_DIRNAME = "unstructured_output_table_structure"
|
||||
|
||||
DUMMY_DF_CCT = pd.DataFrame(
|
||||
{
|
||||
"filename": [
|
||||
"Bank Good Credit Loan.pptx",
|
||||
"Performance-Audit-Discussion.pdf",
|
||||
"currency.csv",
|
||||
],
|
||||
"doctype": ["pptx", "pdf", "csv"],
|
||||
"connector": ["connector1", "connector1", "connector2"],
|
||||
"cct-accuracy": [0.812, 0.994, 0.887],
|
||||
"cct-%missing": [0.001, 0.002, 0.041],
|
||||
}
|
||||
)
|
||||
|
||||
DUMMY_DF_ELEMENT_TYPE = pd.DataFrame(
|
||||
{
|
||||
"filename": [
|
||||
"Bank Good Credit Loan.pptx",
|
||||
"Performance-Audit-Discussion.pdf",
|
||||
"currency.csv",
|
||||
],
|
||||
"doctype": ["pptx", "pdf", "csv"],
|
||||
"connector": ["connector1", "connector1", "connector2"],
|
||||
"element-type-accuracy": [0.812, 0.994, 0.887],
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_dependencies():
|
||||
with patch(
|
||||
"unstructured.metrics.evaluate.calculate_accuracy"
|
||||
) as mock_calculate_accuracy, patch(
|
||||
"unstructured.metrics.evaluate.calculate_percent_missing_text"
|
||||
) as mock_calculate_percent_missing_text, patch.object(
|
||||
TextExtractionMetricsCalculator, "_get_ccts"
|
||||
) as mock_get_ccts, patch(
|
||||
"unstructured.metrics.evaluate.get_element_type_frequency"
|
||||
) as mock_get_element_type_frequency, patch(
|
||||
"unstructured.metrics.evaluate.calculate_element_type_percent_match"
|
||||
) as mock_calculate_element_type_percent_match, patch(
|
||||
"unstructured.metrics.evaluate._read_text_file"
|
||||
) as mock_read_text_file, patch.object(
|
||||
Path, "exists"
|
||||
) as mock_path_exists, patch(
|
||||
"unstructured.metrics.evaluate.TableEvalProcessor.from_json_files"
|
||||
) as mock_table_eval_processor_from_json_files, patch.object(
|
||||
TableStructureMetricsCalculator, "supported_metric_names"
|
||||
) as mock_supported_metric_names:
|
||||
mocks = {
|
||||
"mock_calculate_accuracy": mock_calculate_accuracy,
|
||||
"mock_calculate_percent_missing_text": mock_calculate_percent_missing_text,
|
||||
"mock_get_ccts": mock_get_ccts,
|
||||
"mock_get_element_type_frequency": mock_get_element_type_frequency,
|
||||
"mock_read_text_file": mock_read_text_file,
|
||||
"mock_calculate_element_type_percent_match": mock_calculate_element_type_percent_match,
|
||||
"mock_table_eval_processor_from_json_files": mock_table_eval_processor_from_json_files,
|
||||
"mock_supported_metric_names": mock_supported_metric_names,
|
||||
"mock_path_exists": mock_path_exists,
|
||||
}
|
||||
|
||||
# setup mocks
|
||||
mocks["mock_calculate_accuracy"].return_value = 0.5
|
||||
mocks["mock_calculate_percent_missing_text"].return_value = 0.5
|
||||
mocks["mock_get_ccts"].return_value = ["output_cct", "source_cct"]
|
||||
mocks["mock_get_element_type_frequency"].side_effect = [{"ele1": 1}, {"ele2": 3}]
|
||||
mocks["mock_calculate_element_type_percent_match"].return_value = 0.5
|
||||
mocks["mock_supported_metric_names"].return_value = ["table_level_acc"]
|
||||
mocks["mock_path_exists"].return_value = True
|
||||
mocks["mock_read_text_file"].side_effect = ["output_text", "source_text"]
|
||||
|
||||
yield mocks
|
||||
|
||||
|
||||
@pytest.fixture()
|
||||
def _cleanup_after_test():
|
||||
"""Fixture for removing side-effects of running tests in this file."""
|
||||
|
||||
def remove_generated_directories():
|
||||
"""Remove directories created from running tests."""
|
||||
|
||||
# Directories to be removed:
|
||||
target_dir_names = [
|
||||
"test_evaluate_results_cct",
|
||||
"test_evaluate_results_cct_txt",
|
||||
"test_evaluate_results_element_type",
|
||||
"test_evaluate_result_table_structure",
|
||||
]
|
||||
subdirs = (d for d in os.scandir(TESTING_FILE_DIR) if d.is_dir())
|
||||
for d in subdirs:
|
||||
if d.name in target_dir_names:
|
||||
shutil.rmtree(d.path)
|
||||
|
||||
# Run test as normal
|
||||
yield
|
||||
remove_generated_directories()
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_text_extraction_evaluation():
|
||||
output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME)
|
||||
source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME)
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
TextExtractionMetricsCalculator(
|
||||
documents_dir=output_dir, ground_truths_dir=source_dir
|
||||
).calculate(export_dir=export_dir, visualize_progress=False, display_agg_df=False)
|
||||
|
||||
assert os.path.isfile(os.path.join(export_dir, "all-docs-cct.tsv"))
|
||||
df = pd.read_csv(os.path.join(export_dir, "all-docs-cct.tsv"), sep="\t")
|
||||
assert len(df) == 3
|
||||
assert len(df.columns) == 5
|
||||
assert df.iloc[0].filename == "Bank Good Credit Loan.pptx"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("calculator_class", "output_dirname", "source_dirname", "path", "expected_length", "kwargs"),
|
||||
[
|
||||
(
|
||||
TextExtractionMetricsCalculator,
|
||||
UNSTRUCTURED_CCT_DIRNAME,
|
||||
GOLD_CCT_DIRNAME,
|
||||
Path("Bank Good Credit Loan.pptx.txt"),
|
||||
5,
|
||||
{"document_type": "txt"},
|
||||
),
|
||||
(
|
||||
TableStructureMetricsCalculator,
|
||||
UNSTRUCTURED_TABLE_STRUCTURE_DIRNAME,
|
||||
GOLD_TABLE_STRUCTURE_DIRNAME,
|
||||
Path("IRS-2023-Form-1095-A.pdf.json"),
|
||||
14,
|
||||
{},
|
||||
),
|
||||
(
|
||||
ElementTypeMetricsCalculator,
|
||||
UNSTRUCTURED_OUTPUT_DIRNAME,
|
||||
GOLD_ELEMENT_TYPE_DIRNAME,
|
||||
Path("IRS-form-1987.pdf.json"),
|
||||
4,
|
||||
{},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_process_document_returns_the_correct_amount_of_values(
|
||||
calculator_class, output_dirname, source_dirname, path, expected_length, kwargs
|
||||
):
|
||||
output_dir = Path(TESTING_FILE_DIR) / output_dirname
|
||||
source_dir = Path(TESTING_FILE_DIR) / source_dirname
|
||||
|
||||
calculator = calculator_class(documents_dir=output_dir, ground_truths_dir=source_dir, **kwargs)
|
||||
output_list = calculator._process_document(path)
|
||||
assert len(output_list) == expected_length
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test", "mock_dependencies")
|
||||
@pytest.mark.parametrize(
|
||||
("calculator_class", "output_dirname", "source_dirname", "path", "kwargs"),
|
||||
[
|
||||
(
|
||||
TextExtractionMetricsCalculator,
|
||||
UNSTRUCTURED_CCT_DIRNAME,
|
||||
GOLD_CCT_DIRNAME,
|
||||
Path("2310.03502text_to_image_synthesis1-7.pdf.txt"),
|
||||
{"document_type": "txt"},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_TextExtractionMetricsCalculator_process_document_returns_the_correct_doctype(
|
||||
mock_dependencies, calculator_class, output_dirname, source_dirname, path, kwargs
|
||||
):
|
||||
output_dir = Path(TESTING_FILE_DIR) / output_dirname
|
||||
source_dir = Path(TESTING_FILE_DIR) / source_dirname
|
||||
mock_calculate_accuracy = mock_dependencies["mock_calculate_accuracy"]
|
||||
mock_calculate_percent_missing_text = mock_dependencies["mock_calculate_percent_missing_text"]
|
||||
mock_get_ccts = mock_dependencies["mock_get_ccts"]
|
||||
calculator = calculator_class(documents_dir=output_dir, ground_truths_dir=source_dir, **kwargs)
|
||||
output_list = calculator._process_document(path)
|
||||
assert output_list[1] == ".pdf"
|
||||
assert mock_calculate_accuracy.call_count == 1
|
||||
assert mock_calculate_percent_missing_text.call_count == 1
|
||||
assert mock_get_ccts.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test", "mock_dependencies")
|
||||
@pytest.mark.parametrize(
|
||||
("calculator_class", "output_dirname", "source_dirname", "path", "kwargs"),
|
||||
[
|
||||
(
|
||||
TableStructureMetricsCalculator,
|
||||
UNSTRUCTURED_TABLE_STRUCTURE_DIRNAME,
|
||||
GOLD_TABLE_STRUCTURE_DIRNAME,
|
||||
Path("tablib-627mTABLES-2310.07875-p7.pdf.json"),
|
||||
{},
|
||||
),
|
||||
# (
|
||||
# ElementTypeMetricsCalculator,
|
||||
# UNSTRUCTURED_OUTPUT_DIRNAME,
|
||||
# GOLD_ELEMENT_TYPE_DIRNAME,
|
||||
# Path("IRS-form.1987.pdf.json"),
|
||||
# {},
|
||||
# ),
|
||||
],
|
||||
)
|
||||
def test_TableStructureMetricsCalculator_process_document_returns_the_correct_doctype(
|
||||
mock_dependencies, calculator_class, output_dirname, source_dirname, path, kwargs
|
||||
):
|
||||
output_dir = Path(TESTING_FILE_DIR) / output_dirname
|
||||
source_dir = Path(TESTING_FILE_DIR) / source_dirname
|
||||
calculator = calculator_class(documents_dir=output_dir, ground_truths_dir=source_dir, **kwargs)
|
||||
calculator._ground_truths_dir = source_dir
|
||||
calculator._documents_dir = output_dir
|
||||
calculator._ground_truth_paths = [source_dir / path]
|
||||
mock_report = MagicMock()
|
||||
mock_report.total_predicted_tables = 3
|
||||
mock_report.table_evel_acc = 0.83
|
||||
mock_table_eval_processor_from_json_files = mock_dependencies[
|
||||
"mock_table_eval_processor_from_json_files"
|
||||
]
|
||||
mock_table_eval_processor_from_json_files.return_value.process_file.return_value = mock_report
|
||||
|
||||
output_list = calculator._process_document(path)
|
||||
assert output_list[1] == ".pdf"
|
||||
assert mock_table_eval_processor_from_json_files.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test", "mock_dependencies")
|
||||
@pytest.mark.parametrize(
|
||||
("calculator_class", "output_dirname", "source_dirname", "path", "kwargs"),
|
||||
[
|
||||
(
|
||||
ElementTypeMetricsCalculator,
|
||||
UNSTRUCTURED_OUTPUT_DIRNAME,
|
||||
GOLD_ELEMENT_TYPE_DIRNAME,
|
||||
Path("IRS-form.1987.pdf.json"),
|
||||
{},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_ElementTypeMetricsCalculator_process_document_returns_the_correct_doctype(
|
||||
mock_dependencies, calculator_class, output_dirname, source_dirname, path, kwargs
|
||||
):
|
||||
output_dir = Path(TESTING_FILE_DIR) / output_dirname
|
||||
source_dir = Path(TESTING_FILE_DIR) / source_dirname
|
||||
calculator = calculator_class(documents_dir=output_dir, ground_truths_dir=source_dir, **kwargs)
|
||||
mock_element_type_frequency = mock_dependencies["mock_get_element_type_frequency"]
|
||||
mock_read_text_file = mock_dependencies["mock_read_text_file"]
|
||||
mock_calculate_element_type_percent_match = mock_dependencies[
|
||||
"mock_calculate_element_type_percent_match"
|
||||
]
|
||||
output_list = calculator._process_document(path)
|
||||
assert output_list[1] == ".pdf"
|
||||
assert mock_read_text_file.call_count == 2
|
||||
assert mock_element_type_frequency.call_count == 2
|
||||
assert mock_calculate_element_type_percent_match.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_text_extraction_evaluation_type_txt():
|
||||
output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_CCT_DIRNAME)
|
||||
source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME)
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
TextExtractionMetricsCalculator(
|
||||
documents_dir=output_dir, ground_truths_dir=source_dir, document_type="txt"
|
||||
).calculate(export_dir=export_dir)
|
||||
|
||||
df = pd.read_csv(os.path.join(export_dir, "all-docs-cct.tsv"), sep="\t")
|
||||
assert len(df) == 3
|
||||
assert len(df.columns) == 5
|
||||
assert df.iloc[0].filename == "Bank Good Credit Loan.pptx"
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_element_type_evaluation():
|
||||
output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME)
|
||||
source_dir = os.path.join(TESTING_FILE_DIR, GOLD_ELEMENT_TYPE_DIRNAME)
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
ElementTypeMetricsCalculator(
|
||||
documents_dir=output_dir,
|
||||
ground_truths_dir=source_dir,
|
||||
).calculate(export_dir=export_dir, visualize_progress=False)
|
||||
|
||||
assert os.path.isfile(os.path.join(export_dir, "all-docs-element-type-frequency.tsv"))
|
||||
df = pd.read_csv(os.path.join(export_dir, "all-docs-element-type-frequency.tsv"), sep="\t")
|
||||
assert len(df) == 1
|
||||
assert len(df.columns) == 4
|
||||
assert df.iloc[0].filename == "IRS-form-1987.pdf"
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_table_structure_evaluation():
|
||||
output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_TABLE_STRUCTURE_DIRNAME)
|
||||
source_dir = os.path.join(TESTING_FILE_DIR, GOLD_TABLE_STRUCTURE_DIRNAME)
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_result_table_structure")
|
||||
|
||||
TableStructureMetricsCalculator(
|
||||
documents_dir=output_dir,
|
||||
ground_truths_dir=source_dir,
|
||||
).calculate(export_dir=export_dir, visualize_progress=False)
|
||||
|
||||
assert os.path.isfile(os.path.join(export_dir, "all-docs-table-structure-accuracy.tsv"))
|
||||
assert os.path.isfile(os.path.join(export_dir, "aggregate-table-structure-accuracy.tsv"))
|
||||
df = pd.read_csv(os.path.join(export_dir, "all-docs-table-structure-accuracy.tsv"), sep="\t")
|
||||
agg_df = pd.read_csv(
|
||||
os.path.join(export_dir, "aggregate-table-structure-accuracy.tsv"), sep="\t"
|
||||
).set_index("metric")
|
||||
assert len(df) == 2
|
||||
assert len(df.columns) == 15
|
||||
assert df.iloc[1].filename == "IRS-2023-Form-1095-A.pdf"
|
||||
assert (
|
||||
np.round(np.average(df["table_level_acc"], weights=df["total_tables"]), 3)
|
||||
== agg_df.loc["table_level_acc", "average"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_text_extraction_takes_list():
|
||||
output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME)
|
||||
output_list = ["currency.csv.json"]
|
||||
source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME)
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
TextExtractionMetricsCalculator(
|
||||
documents_dir=output_dir,
|
||||
ground_truths_dir=source_dir,
|
||||
).on_files(document_paths=output_list).calculate(export_dir=export_dir)
|
||||
|
||||
# check that only the listed files are included
|
||||
assert os.path.isfile(os.path.join(export_dir, "all-docs-cct.tsv"))
|
||||
df = pd.read_csv(os.path.join(export_dir, "all-docs-cct.tsv"), sep="\t")
|
||||
assert len(df) == len(output_list)
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_text_extraction_with_grouping():
|
||||
output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME)
|
||||
source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME)
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
TextExtractionMetricsCalculator(
|
||||
documents_dir=output_dir,
|
||||
ground_truths_dir=source_dir,
|
||||
group_by="doctype",
|
||||
).calculate(export_dir=export_dir)
|
||||
|
||||
df = pd.read_csv(os.path.join(export_dir, "all-doctype-agg-cct.tsv"), sep="\t")
|
||||
assert len(df) == 4 # metrics row and doctype rows
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_text_extraction_wrong_type():
|
||||
output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME)
|
||||
source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME)
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
with pytest.raises(ValueError):
|
||||
TextExtractionMetricsCalculator(
|
||||
documents_dir=output_dir, ground_truths_dir=source_dir, document_type="invalid type"
|
||||
).calculate(export_dir=export_dir)
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
@pytest.mark.parametrize(("grouping", "count_row"), [("doctype", 3), ("connector", 2)])
|
||||
def test_get_mean_grouping_df_input(grouping: str, count_row: int):
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
get_mean_grouping(
|
||||
group_by=grouping,
|
||||
data_input=DUMMY_DF_CCT,
|
||||
export_dir=export_dir,
|
||||
eval_name="text_extraction",
|
||||
)
|
||||
grouped_df = pd.read_csv(os.path.join(export_dir, f"all-{grouping}-agg-cct.tsv"), sep="\t")
|
||||
assert grouped_df[grouping].dropna().nunique() == count_row
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_get_mean_grouping_tsv_input():
|
||||
output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME)
|
||||
source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME)
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
TextExtractionMetricsCalculator(
|
||||
documents_dir=output_dir,
|
||||
ground_truths_dir=source_dir,
|
||||
).calculate(export_dir=export_dir)
|
||||
|
||||
filename = os.path.join(export_dir, "all-docs-cct.tsv")
|
||||
get_mean_grouping(
|
||||
group_by="doctype",
|
||||
data_input=filename,
|
||||
export_dir=export_dir,
|
||||
eval_name="text_extraction",
|
||||
)
|
||||
grouped_df = pd.read_csv(os.path.join(export_dir, "all-doctype-agg-cct.tsv"), sep="\t")
|
||||
assert grouped_df["doctype"].dropna().nunique() == 3
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_get_mean_grouping_invalid_group():
|
||||
output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME)
|
||||
source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME)
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
TextExtractionMetricsCalculator(
|
||||
documents_dir=output_dir,
|
||||
ground_truths_dir=source_dir,
|
||||
).calculate(export_dir=export_dir)
|
||||
|
||||
df = pd.read_csv(os.path.join(export_dir, "all-docs-cct.tsv"), sep="\t")
|
||||
with pytest.raises(ValueError):
|
||||
get_mean_grouping(
|
||||
group_by="invalid",
|
||||
data_input=df,
|
||||
export_dir=export_dir,
|
||||
eval_name="text_extraction",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_text_extraction_grouping_empty_df():
|
||||
empty_df = pd.DataFrame()
|
||||
with pytest.raises(SystemExit):
|
||||
get_mean_grouping("doctype", empty_df, "some_dir", eval_name="text_extraction")
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_get_mean_grouping_missing_grouping_column():
|
||||
df_with_no_grouping = pd.DataFrame({"some_column": [1, 2, 3]})
|
||||
with pytest.raises(SystemExit):
|
||||
get_mean_grouping("doctype", df_with_no_grouping, "some_dir", "text_extraction")
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_get_mean_grouping_all_null_grouping_column():
|
||||
df_with_null_grouping = pd.DataFrame({"doctype": [None, None, None]})
|
||||
with pytest.raises(SystemExit):
|
||||
get_mean_grouping("doctype", df_with_null_grouping, "some_dir", eval_name="text_extraction")
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_get_mean_grouping_invalid_eval_name():
|
||||
with pytest.raises(ValueError):
|
||||
get_mean_grouping("doctype", DUMMY_DF_ELEMENT_TYPE, "some_dir", eval_name="invalid")
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
@pytest.mark.parametrize(("group_by", "count_row"), [("doctype", 3), ("connector", 2)])
|
||||
def test_get_mean_grouping_element_type(group_by: str, count_row: int):
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_element_type")
|
||||
get_mean_grouping(
|
||||
group_by=group_by,
|
||||
data_input=DUMMY_DF_ELEMENT_TYPE,
|
||||
export_dir=export_dir,
|
||||
eval_name="element_type",
|
||||
)
|
||||
grouped_df = pd.read_csv(
|
||||
os.path.join(export_dir, f"all-{group_by}-agg-element-type.tsv"), sep="\t"
|
||||
)
|
||||
assert grouped_df[group_by].dropna().nunique() == count_row
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_filter_metrics():
|
||||
with open(os.path.join(TESTING_FILE_DIR, "filter_list.txt"), "w") as file:
|
||||
file.write("Bank Good Credit Loan.pptx\n")
|
||||
file.write("Performance-Audit-Discussion.pdf\n")
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
filter_metrics(
|
||||
data_input=DUMMY_DF_CCT,
|
||||
filter_list=os.path.join(TESTING_FILE_DIR, "filter_list.txt"),
|
||||
filter_by="filename",
|
||||
export_filename="filtered_metrics.tsv",
|
||||
export_dir=export_dir,
|
||||
return_type="file",
|
||||
)
|
||||
filtered_df = pd.read_csv(os.path.join(export_dir, "filtered_metrics.tsv"), sep="\t")
|
||||
assert len(filtered_df) == 2
|
||||
assert filtered_df["filename"].iloc[0] == "Bank Good Credit Loan.pptx"
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_get_mean_grouping_all_file():
|
||||
with open(os.path.join(TESTING_FILE_DIR, "filter_list.txt"), "w") as file:
|
||||
file.write("Bank Good Credit Loan.pptx\n")
|
||||
file.write("Performance-Audit-Discussion.pdf\n")
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
filter_metrics(
|
||||
data_input=DUMMY_DF_CCT,
|
||||
filter_list=["Bank Good Credit Loan.pptx", "Performance-Audit-Discussion.pdf"],
|
||||
filter_by="filename",
|
||||
export_filename="filtered_metrics.tsv",
|
||||
export_dir=export_dir,
|
||||
return_type="file",
|
||||
)
|
||||
filtered_df = pd.read_csv(os.path.join(export_dir, "filtered_metrics.tsv"), sep="\t")
|
||||
|
||||
get_mean_grouping(
|
||||
group_by="all",
|
||||
data_input=filtered_df,
|
||||
export_dir=export_dir,
|
||||
eval_name="text_extraction",
|
||||
export_filename="two-filename-agg-cct.tsv",
|
||||
)
|
||||
grouped_df = pd.read_csv(os.path.join(export_dir, "two-filename-agg-cct.tsv"), sep="\t")
|
||||
|
||||
assert np.isclose(float(grouped_df.iloc[1, 0]), 0.903)
|
||||
assert np.isclose(float(grouped_df.iloc[1, 1]), 0.129)
|
||||
assert np.isclose(float(grouped_df.iloc[1, 2]), 0.091)
|
||||
|
||||
|
||||
@pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container")
|
||||
@pytest.mark.usefixtures("_cleanup_after_test")
|
||||
def test_get_mean_grouping_all_file_txt():
|
||||
with open(os.path.join(TESTING_FILE_DIR, "filter_list.txt"), "w") as file:
|
||||
file.write("Bank Good Credit Loan.pptx\n")
|
||||
file.write("Performance-Audit-Discussion.pdf\n")
|
||||
export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct")
|
||||
|
||||
filter_metrics(
|
||||
data_input=DUMMY_DF_CCT,
|
||||
filter_list=os.path.join(TESTING_FILE_DIR, "filter_list.txt"),
|
||||
filter_by="filename",
|
||||
export_filename="filtered_metrics.tsv",
|
||||
export_dir=export_dir,
|
||||
return_type="file",
|
||||
)
|
||||
filtered_df = pd.read_csv(os.path.join(export_dir, "filtered_metrics.tsv"), sep="\t")
|
||||
|
||||
get_mean_grouping(
|
||||
group_by="all",
|
||||
data_input=filtered_df,
|
||||
export_dir=export_dir,
|
||||
eval_name="text_extraction",
|
||||
export_filename="two-filename-agg-cct.tsv",
|
||||
)
|
||||
grouped_df = pd.read_csv(os.path.join(export_dir, "two-filename-agg-cct.tsv"), sep="\t")
|
||||
|
||||
assert np.isclose(float(grouped_df.iloc[1, 0]), 0.903)
|
||||
assert np.isclose(float(grouped_df.iloc[1, 1]), 0.129)
|
||||
assert np.isclose(float(grouped_df.iloc[1, 2]), 0.091)
|
||||
@@ -0,0 +1,14 @@
|
||||
from unstructured.metrics.table.table_alignment import TableAlignment
|
||||
|
||||
|
||||
def test_get_element_level_alignment_when_no_match():
|
||||
example_table = [{"row_index": 0, "col_index": 0, "content": "a"}]
|
||||
metrics = TableAlignment.get_element_level_alignment(
|
||||
predicted_table_data=[example_table],
|
||||
ground_truth_table_data=[example_table],
|
||||
matched_indices=[-1],
|
||||
)
|
||||
assert metrics["col_index_acc"] == 0
|
||||
assert metrics["row_index_acc"] == 0
|
||||
assert metrics["row_content_acc"] == 0
|
||||
assert metrics["col_content_acc"] == 0
|
||||
@@ -0,0 +1,33 @@
|
||||
import pytest
|
||||
|
||||
from unstructured.metrics.table.table_eval import calculate_table_detection_metrics
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("matched_indices", "ground_truth_tables_number", "expected_metrics"),
|
||||
[
|
||||
([0, 1, 2], 3, (1, 1, 1)), # everything was predicted correctly
|
||||
([2, 1, 0], 3, (1, 1, 1)), # everything was predicted correctly
|
||||
(
|
||||
[-1, 2, -1, 1, 0, -1],
|
||||
3,
|
||||
(1, 0.5, 0.66),
|
||||
), # some false positives, all tables matched, too many predictions
|
||||
([2, 2, 1, 1], 8, (0.25, 0.5, 0.33)),
|
||||
# Some false negatives, all predictions matched with gt, not enough predictions
|
||||
# The precision here is not 1 as only one from tables matched with '1' index can be correct
|
||||
([1, -1], 2, (0.5, 0.5, 0.5)), # typical case with false positive and false negative
|
||||
([-1, -1, -1], 2, (0, 0, 0)), # nothing was matched
|
||||
([-1, -1, -1], 0, (0, 0, 0)), # there was no table in ground truth
|
||||
([], 0, (0, 0, 0)), # just zeros to account for errors
|
||||
],
|
||||
)
|
||||
def test_calculate_table_metrics(matched_indices, ground_truth_tables_number, expected_metrics):
|
||||
expected_recall, expected_precision, expected_f1 = expected_metrics
|
||||
pred_recall, pred_precision, pred_f1 = calculate_table_detection_metrics(
|
||||
matched_indices=matched_indices, ground_truth_tables_number=ground_truth_tables_number
|
||||
)
|
||||
|
||||
assert pred_recall == expected_recall
|
||||
assert pred_precision == expected_precision
|
||||
assert pred_f1 == pytest.approx(expected_f1, abs=0.01)
|
||||
@@ -0,0 +1,33 @@
|
||||
import pytest
|
||||
|
||||
from unstructured.metrics.table.table_formats import SimpleTableCell
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("row_nums", "column_nums", "x", "y", "w", "h"),
|
||||
[
|
||||
([3, 2, 1], [6, 7], 6, 1, 2, 3),
|
||||
([2], [6, 7], 6, 2, 2, 1),
|
||||
([1, 2, 3], [20], 20, 1, 1, 3),
|
||||
([5], [5], 5, 5, 1, 1),
|
||||
],
|
||||
)
|
||||
def test_simple_table_cell_parsing_from_table_transformer_when_expected_input(
|
||||
row_nums, column_nums, x, y, w, h
|
||||
):
|
||||
table_transformer_cell = {"row_nums": row_nums, "column_nums": column_nums, "cell text": "text"}
|
||||
transformed_cell = SimpleTableCell.from_table_transformer_cell(table_transformer_cell)
|
||||
expected_cell = SimpleTableCell(x=x, y=y, w=w, h=h, content="text")
|
||||
assert expected_cell == transformed_cell
|
||||
|
||||
|
||||
def test_simple_table_cell_parsing_from_table_transformer_when_missing_row_nums():
|
||||
cell = {"row_nums": [], "column_nums": [1], "cell text": "text"}
|
||||
with pytest.raises(ValueError, match='has missing values under "row_nums" key'):
|
||||
SimpleTableCell.from_table_transformer_cell(cell)
|
||||
|
||||
|
||||
def test_simple_table_cell_parsing_from_table_transformer_when_missing_column_nums():
|
||||
cell = {"row_nums": [1], "column_nums": [], "cell text": "text"}
|
||||
with pytest.raises(ValueError, match='has missing values under "column_nums" key'):
|
||||
SimpleTableCell.from_table_transformer_cell(cell)
|
||||
@@ -0,0 +1,700 @@
|
||||
from unittest import mock
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from test_unstructured.unit_utils import example_doc_path
|
||||
from unstructured.metrics.table.table_alignment import TableAlignment
|
||||
from unstructured.metrics.table.table_eval import TableEvalProcessor
|
||||
from unstructured.metrics.table_structure import (
|
||||
eval_table_transformer_for_file,
|
||||
image_or_pdf_to_dataframe,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"filename",
|
||||
[
|
||||
example_doc_path("img/table-multi-row-column-cells.png"),
|
||||
example_doc_path("pdf/table-multi-row-column-cells.pdf"),
|
||||
],
|
||||
)
|
||||
def test_image_or_pdf_to_dataframe(filename):
|
||||
df = image_or_pdf_to_dataframe(filename)
|
||||
assert ["Blind", "5", "1", "4", "34.5%, n=1", "1199 sec, n=1"] in df.values
|
||||
|
||||
|
||||
def test_eval_table_transformer_for_file():
|
||||
score = eval_table_transformer_for_file(
|
||||
example_doc_path("img/table-multi-row-column-cells.png"),
|
||||
example_doc_path("table-multi-row-column-cells-actual.csv"),
|
||||
)
|
||||
# avoid severe degradation of performance
|
||||
assert 0.8 < score < 1
|
||||
|
||||
|
||||
def test_table_eval_processor_simple():
|
||||
prediction = [
|
||||
{
|
||||
"type": "Table",
|
||||
"metadata": {
|
||||
"text_as_html": """<table><thead><tr><th>r1c1</th><th>r1c2</th></tr></thead>
|
||||
<tbody><tr><td>r2c1</td><td>r2c2</td></tr></tbody></table>"""
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
ground_truth = [
|
||||
{
|
||||
"type": "Table",
|
||||
"text": [
|
||||
{
|
||||
"id": "ee862c7a-d27e-4484-92de-4faa42a63f3b",
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r1c1",
|
||||
},
|
||||
{
|
||||
"id": "6237ac7b-bfc8-40d2-92f2-d138277205e2",
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c1",
|
||||
},
|
||||
{
|
||||
"id": "9d0933a9-5984-4cad-80d9-6752bf9bc4df",
|
||||
"x": 1,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r1c2",
|
||||
},
|
||||
{
|
||||
"id": "1152d043-5ead-4ab8-8b88-888d48831ac2",
|
||||
"x": 1,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c2",
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth)
|
||||
result = te_processor.process_file()
|
||||
assert result.total_tables == 1
|
||||
assert result.table_level_acc == 1.0
|
||||
assert result.element_row_level_index_acc == 1.0
|
||||
assert result.element_col_level_index_acc == 1.0
|
||||
assert result.element_row_level_content_acc == 1.0
|
||||
assert result.element_col_level_content_acc == 1.0
|
||||
|
||||
|
||||
def test_table_eval_processor_simple_when_input_as_cells():
|
||||
prediction = [
|
||||
{
|
||||
"type": "Table",
|
||||
"metadata": {
|
||||
"table_as_cells": [
|
||||
{
|
||||
"x": 1,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c2",
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r1c1",
|
||||
},
|
||||
{
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c1",
|
||||
},
|
||||
{
|
||||
"x": 1,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r1c2",
|
||||
},
|
||||
]
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
ground_truth = [
|
||||
{
|
||||
"type": "Table",
|
||||
"text": [
|
||||
{
|
||||
"id": "ee862c7a-d27e-4484-92de-4faa42a63f3b",
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r1c1",
|
||||
},
|
||||
{
|
||||
"id": "6237ac7b-bfc8-40d2-92f2-d138277205e2",
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c1",
|
||||
},
|
||||
{
|
||||
"id": "9d0933a9-5984-4cad-80d9-6752bf9bc4df",
|
||||
"x": 1,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r1c2",
|
||||
},
|
||||
{
|
||||
"id": "1152d043-5ead-4ab8-8b88-888d48831ac2",
|
||||
"x": 1,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c2",
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth, source_type="cells")
|
||||
result = te_processor.process_file()
|
||||
assert result.total_tables == 1
|
||||
assert result.table_level_acc == 1.0
|
||||
assert result.element_row_level_index_acc == 1.0
|
||||
assert result.element_col_level_index_acc == 1.0
|
||||
assert result.element_row_level_content_acc == 1.0
|
||||
assert result.element_col_level_content_acc == 1.0
|
||||
|
||||
|
||||
def test_table_eval_processor_when_wrong_source_type():
|
||||
prediction = [
|
||||
{
|
||||
"type": "Table",
|
||||
"metadata": {"table_as_cells": []},
|
||||
}
|
||||
]
|
||||
|
||||
ground_truth = [
|
||||
{
|
||||
"type": "Table",
|
||||
"text": [],
|
||||
}
|
||||
]
|
||||
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth, source_type="wrong_type")
|
||||
with pytest.raises(ValueError):
|
||||
te_processor.process_file()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"text_as_html",
|
||||
[
|
||||
"""
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>r1c1</th>
|
||||
<th>r1c2</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>r2c1</td>
|
||||
<td>r2c2</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>r3c1</td>
|
||||
<td>r3c2</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
""",
|
||||
"""
|
||||
<table>
|
||||
<tr>
|
||||
<th>r1c1</th>
|
||||
<th>r1c2</th>
|
||||
</tr>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>r2c1</td>
|
||||
<td>r2c2</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>r3c1</td>
|
||||
<td>r3c2</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
""",
|
||||
"""
|
||||
<table>
|
||||
</tbody>
|
||||
<tr>
|
||||
<td>r1c1</td>
|
||||
<td>r1c2</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>r2c1</td>
|
||||
<td>r2c2</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>r3c1</td>
|
||||
<td>r3c2</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
""",
|
||||
],
|
||||
)
|
||||
def test_table_eval_processor_various_table_html_structures(text_as_html):
|
||||
prediction = [{"type": "Table", "metadata": {"text_as_html": text_as_html}}]
|
||||
|
||||
ground_truth = [
|
||||
{
|
||||
"type": "Table",
|
||||
"text": [
|
||||
{
|
||||
"id": "ee862c7a-d27e-4484-92de-4faa42a63f3b",
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r1c1",
|
||||
},
|
||||
{
|
||||
"id": "6237ac7b-bfc8-40d2-92f2-d138277205e2",
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c1",
|
||||
},
|
||||
{
|
||||
"id": "9d0933a9-5984-4cad-80d9-6752bf9bc4df",
|
||||
"x": 1,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r1c2",
|
||||
},
|
||||
{
|
||||
"id": "1152d043-5ead-4ab8-8b88-888d48831ac2",
|
||||
"x": 1,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c2",
|
||||
},
|
||||
{
|
||||
"id": "364f4a17-2979-4506-ae77-e8adf8e3f554",
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r3c1",
|
||||
},
|
||||
{
|
||||
"id": "30f87503-ac1f-4db1-b924-b316af585702",
|
||||
"x": 1,
|
||||
"y": 2,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r3c2",
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth)
|
||||
result = te_processor.process_file()
|
||||
assert result.total_tables == 1
|
||||
assert result.table_level_acc == 1.0
|
||||
assert result.element_row_level_index_acc == 1.0
|
||||
assert result.element_col_level_index_acc == 1.0
|
||||
assert result.element_row_level_content_acc == 1.0
|
||||
assert result.element_col_level_content_acc == 1.0
|
||||
|
||||
|
||||
def test_table_eval_processor_non_str_values_in_table():
|
||||
prediction = [
|
||||
{
|
||||
"type": "Table",
|
||||
"metadata": {
|
||||
"text_as_html": """
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>11</th>
|
||||
<th>12</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>21</td>
|
||||
<td>22</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>"""
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
ground_truth = [
|
||||
{
|
||||
"type": "Table",
|
||||
"text": [
|
||||
{
|
||||
"id": "ee862c7a-d27e-4484-92de-4faa42a63f3b",
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "11",
|
||||
},
|
||||
{
|
||||
"id": "6237ac7b-bfc8-40d2-92f2-d138277205e2",
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "21",
|
||||
},
|
||||
{
|
||||
"id": "9d0933a9-5984-4cad-80d9-6752bf9bc4df",
|
||||
"x": 1,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "12",
|
||||
},
|
||||
{
|
||||
"id": "1152d043-5ead-4ab8-8b88-888d48831ac2",
|
||||
"x": 1,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "22",
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth)
|
||||
result = te_processor.process_file()
|
||||
assert result.total_tables == 1
|
||||
assert result.table_level_acc == 1.0
|
||||
assert result.element_row_level_index_acc == 1.0
|
||||
assert result.element_col_level_index_acc == 1.0
|
||||
assert result.element_row_level_content_acc == 1.0
|
||||
assert result.element_col_level_content_acc == 1.0
|
||||
|
||||
|
||||
def test_table_eval_processor_merged_cells():
|
||||
prediction = [
|
||||
{
|
||||
"type": "Table",
|
||||
"metadata": {
|
||||
"text_as_html": """
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th rowspan="2">r1c1</th>
|
||||
<th>r1c2</th>
|
||||
<th colspan="2">r1c3</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<th>r2c2</th>
|
||||
<th>r2c3</th>
|
||||
<th>r2c4</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>r3c1</td>
|
||||
<td>r3c2</td>
|
||||
<td colspan="2" rowspan="2">r3c3</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>r4c1</td>
|
||||
<td>r4c2</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
"""
|
||||
},
|
||||
}
|
||||
]
|
||||
|
||||
ground_truth = [
|
||||
{
|
||||
"type": "Table",
|
||||
"text": [
|
||||
{
|
||||
"id": "f399ef57-5b88-4509-8971-9cb63246866e",
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 2,
|
||||
"content": "r1c1",
|
||||
},
|
||||
{
|
||||
"id": "2dfdec2f-e8f3-4be7-a6ac-8ff21c4e8556",
|
||||
"x": 0,
|
||||
"y": 2,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r3c1",
|
||||
},
|
||||
{
|
||||
"id": "9c771c58-88c7-49d8-9c12-85d0e44b920e",
|
||||
"x": 0,
|
||||
"y": 3,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r4c1",
|
||||
},
|
||||
{
|
||||
"id": "5bd6f3f0-34c5-495b-8a28-c4ac96989ef8",
|
||||
"x": 1,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r1c2",
|
||||
},
|
||||
{
|
||||
"id": "7b8e6bc2-a310-4dd6-997c-313f951e7f96",
|
||||
"x": 1,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c2",
|
||||
},
|
||||
{
|
||||
"id": "1c152ad4-12fa-4a7b-90de-a992aa6410a4",
|
||||
"x": 1,
|
||||
"y": 2,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r3c2",
|
||||
},
|
||||
{
|
||||
"id": "55063f64-0003-4217-b6ca-aff5914793ff",
|
||||
"x": 1,
|
||||
"y": 3,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r4c2",
|
||||
},
|
||||
{
|
||||
"id": "22852e86-0e22-4d32-b63a-9ba7dd4118a2",
|
||||
"x": 2,
|
||||
"y": 0,
|
||||
"w": 2,
|
||||
"h": 1,
|
||||
"content": "r1c3",
|
||||
},
|
||||
{
|
||||
"id": "eae013c5-5597-4a8b-9771-82e28c5c5cba",
|
||||
"x": 2,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c3",
|
||||
},
|
||||
{
|
||||
"id": "0dea3a42-8523-4d6e-9e70-d65cc2314678",
|
||||
"x": 2,
|
||||
"y": 2,
|
||||
"w": 2,
|
||||
"h": 2,
|
||||
"content": "r3c3",
|
||||
},
|
||||
{
|
||||
"id": "60093e2c-d3e2-4146-92b5-97a2fc16c061",
|
||||
"x": 3,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r2c4",
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth)
|
||||
result = te_processor.process_file()
|
||||
assert result.total_tables == 1
|
||||
assert result.table_level_acc == 1.0
|
||||
assert result.element_row_level_index_acc == 1.0
|
||||
assert result.element_col_level_index_acc == 1.0
|
||||
assert result.element_row_level_content_acc == 1.0
|
||||
assert result.element_col_level_content_acc == 1.0
|
||||
|
||||
|
||||
def test_table_eval_processor_when_no_match_with_pred():
|
||||
prediction = [
|
||||
{
|
||||
"type": "Table",
|
||||
"metadata": {"text_as_html": """<table><tr><td>Some cell</td></tr></table>"""},
|
||||
}
|
||||
]
|
||||
|
||||
ground_truth = [
|
||||
{
|
||||
"type": "Table",
|
||||
"text": [
|
||||
{
|
||||
"id": "ee862c7a-d27e-4484-92de-4faa42a63f3b",
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "11",
|
||||
},
|
||||
{
|
||||
"id": "6237ac7b-bfc8-40d2-92f2-d138277205e2",
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "21",
|
||||
},
|
||||
{
|
||||
"id": "9d0933a9-5984-4cad-80d9-6752bf9bc4df",
|
||||
"x": 1,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "12",
|
||||
},
|
||||
{
|
||||
"id": "1152d043-5ead-4ab8-8b88-888d48831ac2",
|
||||
"x": 1,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "22",
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
with mock.patch.object(TableAlignment, "get_table_level_alignment") as align_fn:
|
||||
align_fn.return_value = [-1]
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth)
|
||||
result = te_processor.process_file()
|
||||
|
||||
assert result.total_tables == 1
|
||||
assert result.table_level_acc == 0
|
||||
assert result.element_row_level_index_acc == 0
|
||||
assert result.element_col_level_index_acc == 0
|
||||
assert result.element_row_level_content_acc == 0
|
||||
assert result.element_col_level_content_acc == 0
|
||||
|
||||
|
||||
def test_table_eval_processor_when_no_tables():
|
||||
prediction = [{}]
|
||||
|
||||
ground_truth = [{}]
|
||||
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth)
|
||||
result = te_processor.process_file()
|
||||
assert result.total_tables == 0
|
||||
assert result.table_level_acc == 1
|
||||
assert np.isnan(result.element_row_level_index_acc)
|
||||
assert np.isnan(result.element_col_level_index_acc)
|
||||
assert np.isnan(result.element_row_level_content_acc)
|
||||
assert np.isnan(result.element_col_level_content_acc)
|
||||
|
||||
|
||||
def test_table_eval_processor_when_only_gt():
|
||||
prediction = []
|
||||
|
||||
ground_truth = [
|
||||
{
|
||||
"type": "Table",
|
||||
"text": [
|
||||
{
|
||||
"id": "ee862c7a-d27e-4484-92de-4faa42a63f3b",
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "11",
|
||||
},
|
||||
{
|
||||
"id": "6237ac7b-bfc8-40d2-92f2-d138277205e2",
|
||||
"x": 0,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "21",
|
||||
},
|
||||
{
|
||||
"id": "9d0933a9-5984-4cad-80d9-6752bf9bc4df",
|
||||
"x": 1,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "12",
|
||||
},
|
||||
{
|
||||
"id": "1152d043-5ead-4ab8-8b88-888d48831ac2",
|
||||
"x": 1,
|
||||
"y": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "22",
|
||||
},
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth)
|
||||
result = te_processor.process_file()
|
||||
|
||||
assert result.total_tables == 1
|
||||
assert result.table_level_acc == 0
|
||||
assert result.element_row_level_index_acc == 0
|
||||
assert result.element_col_level_index_acc == 0
|
||||
assert result.element_row_level_content_acc == 0
|
||||
assert result.element_col_level_content_acc == 0
|
||||
|
||||
|
||||
def test_table_eval_processor_when_only_pred():
|
||||
prediction = [
|
||||
{
|
||||
"type": "Table",
|
||||
"metadata": {"text_as_html": """<table><tr><td>Some cell</td></tr></table>"""},
|
||||
}
|
||||
]
|
||||
|
||||
ground_truth = [{}]
|
||||
|
||||
te_processor = TableEvalProcessor(prediction, ground_truth)
|
||||
result = te_processor.process_file()
|
||||
|
||||
assert result.total_tables == 0
|
||||
assert result.table_level_acc == 0
|
||||
assert result.element_row_level_index_acc == 0
|
||||
assert result.element_col_level_index_acc == 0
|
||||
assert result.element_row_level_content_acc == 0
|
||||
assert result.element_col_level_content_acc == 0
|
||||
@@ -0,0 +1,859 @@
|
||||
import re
|
||||
|
||||
import pytest
|
||||
|
||||
from unstructured.metrics import text_extraction
|
||||
from unstructured.metrics.table.table_extraction import (
|
||||
deckerd_table_to_html,
|
||||
extract_cells_from_table_as_cells,
|
||||
extract_cells_from_text_as_html,
|
||||
html_table_to_deckerd,
|
||||
)
|
||||
from unstructured.partition.auto import partition
|
||||
|
||||
|
||||
def test_calculate_edit_distance():
|
||||
source_cct = "I like pizza. I like bagels."
|
||||
source_cct_word_space = "I like p i z z a . I like bagles."
|
||||
source_cct_spaces = re.sub(r"\s+", " ", " ".join(source_cct))
|
||||
source_cct_no_space = source_cct.replace(" ", "")
|
||||
source_cct_one_sentence = "I like pizza."
|
||||
source_cct_missing_word = "I like pizza. I like ."
|
||||
source_cct_addn_char = "I like pizza. I like beagles."
|
||||
source_cct_dup_word = "I like pizza pizza. I like bagels."
|
||||
|
||||
assert (
|
||||
round(text_extraction.calculate_edit_distance(source_cct, source_cct, return_as="score"), 2)
|
||||
== 1.0
|
||||
)
|
||||
assert (
|
||||
round(
|
||||
text_extraction.calculate_edit_distance(
|
||||
source_cct_word_space,
|
||||
source_cct,
|
||||
return_as="score",
|
||||
),
|
||||
2,
|
||||
)
|
||||
== 0.75
|
||||
)
|
||||
assert (
|
||||
round(
|
||||
text_extraction.calculate_edit_distance(
|
||||
source_cct_spaces,
|
||||
source_cct,
|
||||
return_as="score",
|
||||
),
|
||||
2,
|
||||
)
|
||||
== 0.39
|
||||
)
|
||||
assert (
|
||||
round(
|
||||
text_extraction.calculate_edit_distance(
|
||||
source_cct_no_space,
|
||||
source_cct,
|
||||
return_as="score",
|
||||
),
|
||||
2,
|
||||
)
|
||||
== 0.64
|
||||
)
|
||||
assert (
|
||||
round(
|
||||
text_extraction.calculate_edit_distance(
|
||||
source_cct_one_sentence,
|
||||
source_cct,
|
||||
return_as="score",
|
||||
),
|
||||
2,
|
||||
)
|
||||
== 0.0
|
||||
)
|
||||
assert (
|
||||
round(
|
||||
text_extraction.calculate_edit_distance(
|
||||
source_cct_missing_word,
|
||||
source_cct,
|
||||
return_as="score",
|
||||
),
|
||||
2,
|
||||
)
|
||||
== 0.57
|
||||
)
|
||||
assert (
|
||||
round(
|
||||
text_extraction.calculate_edit_distance(
|
||||
source_cct_addn_char,
|
||||
source_cct,
|
||||
return_as="score",
|
||||
),
|
||||
2,
|
||||
)
|
||||
== 0.89
|
||||
)
|
||||
assert (
|
||||
round(
|
||||
text_extraction.calculate_edit_distance(
|
||||
source_cct_dup_word,
|
||||
source_cct,
|
||||
return_as="score",
|
||||
),
|
||||
2,
|
||||
)
|
||||
== 0.79
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filename", "standardize_whitespaces", "expected_score", "expected_distance"),
|
||||
[
|
||||
("fake-text.txt", False, 0.78, 38),
|
||||
("fake-text.txt", True, 0.92, 12),
|
||||
],
|
||||
)
|
||||
def test_calculate_edit_distance_with_filename(
|
||||
filename, standardize_whitespaces, expected_score, expected_distance
|
||||
):
|
||||
with open("example-docs/fake-text.txt") as f:
|
||||
source_cct = f.read()
|
||||
|
||||
elements = partition(filename=f"example-docs/{filename}")
|
||||
output_cct = "\n".join([str(el) for el in elements])
|
||||
|
||||
score = text_extraction.calculate_edit_distance(
|
||||
output_cct, source_cct, return_as="score", standardize_whitespaces=standardize_whitespaces
|
||||
)
|
||||
distance = text_extraction.calculate_edit_distance(
|
||||
output_cct,
|
||||
source_cct,
|
||||
return_as="distance",
|
||||
standardize_whitespaces=standardize_whitespaces,
|
||||
)
|
||||
|
||||
assert score >= 0
|
||||
assert score <= 1.0
|
||||
assert distance >= 0
|
||||
assert round(score, 2) == expected_score
|
||||
assert distance == expected_distance
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text1", "text2"),
|
||||
[
|
||||
(
|
||||
"The dog\rloved the cat, but\t\n the cat\tloved the\n cow",
|
||||
"The dog loved the cat, but the cat loved the cow",
|
||||
),
|
||||
(
|
||||
"Hello my\tname\tis H a r p e r, \nwhat's your\vname?",
|
||||
"Hello my name is H a r p e r, what's your name?",
|
||||
),
|
||||
(
|
||||
"I have a\t\n\tdog and a\tcat,\fI love my\n\n\n\ndog.",
|
||||
"I have a dog and a cat, I love my dog.",
|
||||
),
|
||||
(
|
||||
"""
|
||||
Name Age City Occupation
|
||||
Alice 30 New York Engineer
|
||||
Bob 25 Los Angeles Designer
|
||||
Charlie 35 Chicago Teacher
|
||||
David 40 San Francisco Developer
|
||||
""",
|
||||
"""
|
||||
Name\tAge\tCity\tOccupation
|
||||
Alice\t30\tNew York\tEngineer
|
||||
Bob\t25\tLos Angeles\tDesigner
|
||||
Charlie\t35\tChicago\tTeacher
|
||||
David\t40\tSan Francisco\tDeveloper
|
||||
""",
|
||||
),
|
||||
(
|
||||
"""
|
||||
Name\tAge\tCity\tOccupation
|
||||
Alice\t30\tNew York\tEngineer
|
||||
Bob\t25\tLos Angeles\tDesigner
|
||||
Charlie\t35\tChicago\tTeacher
|
||||
David\t40\tSan Francisco\tDeveloper
|
||||
""",
|
||||
"Name\tAge\tCity\tOccupation\n\n \nAlice\t30\tNew York\tEngineer\nBob\t25\tLos Angeles\tDesigner\nCharlie\t35\tChicago\tTeacher\nDavid\t40\tSan Francisco\tDeveloper", # noqa: E501
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_calculate_edit_distance_with_various_whitespace_1(text1, text2):
|
||||
assert (
|
||||
text_extraction.calculate_edit_distance(
|
||||
text1, text2, return_as="score", standardize_whitespaces=True
|
||||
)
|
||||
== 1.0
|
||||
)
|
||||
assert (
|
||||
text_extraction.calculate_edit_distance(
|
||||
text1, text2, return_as="distance", standardize_whitespaces=True
|
||||
)
|
||||
== 0
|
||||
)
|
||||
assert (
|
||||
text_extraction.calculate_edit_distance(
|
||||
text1, text2, return_as="score", standardize_whitespaces=False
|
||||
)
|
||||
< 1.0
|
||||
)
|
||||
assert (
|
||||
text_extraction.calculate_edit_distance(
|
||||
text1, text2, return_as="distance", standardize_whitespaces=False
|
||||
)
|
||||
> 0
|
||||
)
|
||||
|
||||
|
||||
def test_calculate_edit_distance_with_various_whitespace_2():
|
||||
source_cct_tabs = """
|
||||
Name\tAge\tCity\tOccupation
|
||||
Alice\t30\tNew York\tEngineer
|
||||
Bob\t25\tLos Angeles\tDesigner
|
||||
Charlie\t35\tChicago\tTeacher
|
||||
David\t40\tSan Francisco\tDeveloper
|
||||
"""
|
||||
source_cct_with_borders = """
|
||||
|
||||
| Name | Age | City | Occupation |
|
||||
|---------|-----|--------------|----------------|
|
||||
| Alice | 30 | New York | Engineer |
|
||||
| Bob | 25 | Los Angeles | Designer |
|
||||
| Charlie | 35 | Chicago | Teacher |
|
||||
| David | 40 | San Francisco| Developer |
|
||||
|
||||
"""
|
||||
assert text_extraction.calculate_edit_distance(
|
||||
source_cct_tabs, source_cct_with_borders, return_as="score", standardize_whitespaces=True
|
||||
) > text_extraction.calculate_edit_distance(
|
||||
source_cct_tabs, source_cct_with_borders, return_as="score", standardize_whitespaces=False
|
||||
)
|
||||
assert text_extraction.calculate_edit_distance(
|
||||
source_cct_tabs, source_cct_with_borders, return_as="distance", standardize_whitespaces=True
|
||||
) < text_extraction.calculate_edit_distance(
|
||||
source_cct_tabs,
|
||||
source_cct_with_borders,
|
||||
return_as="distance",
|
||||
standardize_whitespaces=False,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
(
|
||||
"The dog loved the cat, but the cat loved the cow",
|
||||
{"the": 4, "cat": 2, "loved": 2, "dog": 1, "but": 1, "cow": 1},
|
||||
),
|
||||
(
|
||||
"Hello my name is H a r p e r, what's your name?",
|
||||
{"hello": 1, "my": 1, "name": 2, "is": 1, "what's": 1, "your": 1},
|
||||
),
|
||||
(
|
||||
"I have a dog and a cat, I love my dog.",
|
||||
{"i": 2, "have": 1, "a": 2, "dog": 2, "and": 1, "cat": 1, "love": 1, "my": 1},
|
||||
),
|
||||
(
|
||||
"My dog's hair is red, but the dogs' houses are blue.",
|
||||
{
|
||||
"my": 1,
|
||||
"dog's": 1,
|
||||
"hair": 1,
|
||||
"is": 1,
|
||||
"red": 1,
|
||||
"but": 1,
|
||||
"the": 1,
|
||||
"dogs'": 1,
|
||||
"houses": 1,
|
||||
"are": 1,
|
||||
"blue": 1,
|
||||
},
|
||||
),
|
||||
(
|
||||
"""Sometimes sentences have a dash - like this one!
|
||||
A hyphen connects 2 words with no gap: easy-peasy.""",
|
||||
{
|
||||
"sometimes": 1,
|
||||
"sentences": 1,
|
||||
"have": 1,
|
||||
"a": 2,
|
||||
"dash": 1,
|
||||
"like": 1,
|
||||
"this": 1,
|
||||
"one": 1,
|
||||
"hyphen": 1,
|
||||
"connects": 1,
|
||||
"2": 1,
|
||||
"words": 1,
|
||||
"with": 1,
|
||||
"no": 1,
|
||||
"gap": 1,
|
||||
"easy-peasy": 1,
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_bag_of_words(text, expected):
|
||||
assert text_extraction.bag_of_words(text) == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text", "expected"),
|
||||
[
|
||||
(
|
||||
"The dog\rloved the cat, but\t\n the cat\tloved the\n cow\n\n",
|
||||
"The dog loved the cat, but the cat loved the cow",
|
||||
),
|
||||
(
|
||||
"\n\nHello my\tname\tis H a r p e r, \nwhat's your\vname?",
|
||||
"Hello my name is H a r p e r, what's your name?",
|
||||
),
|
||||
(
|
||||
"I have a\t\n\tdog and a\tcat,\fI love my\n\n\n\ndog.",
|
||||
"I have a dog and a cat, I love my dog.",
|
||||
),
|
||||
(
|
||||
"""L is for the way you look at me
|
||||
O is for the only one I see
|
||||
V is very, very extraordinary
|
||||
E is even more than anyone that you adore can""",
|
||||
"L is for the way you look at me O is for the only one I see V is very, very extraordinary E is even more than anyone that you adore can", # noqa: E501
|
||||
),
|
||||
(
|
||||
"""
|
||||
| Name | Age | City | Occupation |
|
||||
|---------|-----|--------------|----------------|
|
||||
| Alice | 30 | New York | Engineer |
|
||||
| Bob | 25 | Los Angeles | Designer |
|
||||
| Charlie | 35 | Chicago | Teacher |
|
||||
| David | 40 | San Francisco| Developer |
|
||||
""",
|
||||
"| Name | Age | City | Occupation | |---------|-----|--------------|----------------| | Alice | 30 | New York | Engineer | | Bob | 25 | Los Angeles | Designer | | Charlie | 35 | Chicago | Teacher | | David | 40 | San Francisco| Developer |", # noqa: E501
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_prepare_string(text, expected):
|
||||
assert text_extraction.prepare_str(text, standardize_whitespaces=True) == expected
|
||||
assert text_extraction.prepare_str(text) == text
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("input_text", "expected_output"),
|
||||
[
|
||||
# Mixed quotes in longer sentences
|
||||
(
|
||||
"She said \"Hello\" and then whispered 'Goodbye' before leaving.",
|
||||
"She said \"Hello\" and then whispered 'Goodbye' before leaving.",
|
||||
),
|
||||
# Double low-9 quotes with complex content
|
||||
(
|
||||
"„To be, or not to be, that is the question\" - Shakespeare's famous quote.",
|
||||
'"To be, or not to be, that is the question" - Shakespeare\'s famous quote.',
|
||||
),
|
||||
# Angle quotes with nested quotes
|
||||
(
|
||||
'«When he said "life is beautiful," I believed him» wrote Maria.',
|
||||
'"When he said "life is beautiful," I believed him" wrote Maria.',
|
||||
),
|
||||
# Heavy ornament quotes in dialogue
|
||||
(
|
||||
"❝Do you remember when we first met?❞ she asked with a smile.",
|
||||
'"Do you remember when we first met?" she asked with a smile.',
|
||||
),
|
||||
# Double prime quotes with punctuation
|
||||
(
|
||||
"〝The meeting starts at 10:00, don't be late!〟 announced the manager.",
|
||||
'"The meeting starts at 10:00, don\'t be late!" announced the manager.',
|
||||
),
|
||||
# Corner brackets with nested quotes
|
||||
(
|
||||
'「He told me "This is important" yesterday」, she explained.',
|
||||
"'He told me \"This is important\" yesterday', she explained.",
|
||||
),
|
||||
# White corner brackets with multiple sentences
|
||||
(
|
||||
"『The sun was setting. The birds were singing. It was peaceful.』",
|
||||
"'The sun was setting. The birds were singing. It was peaceful.'",
|
||||
),
|
||||
# Vertical corner brackets with numbers and special characters
|
||||
("﹂Meeting #123 @ 15:00 - Don't forget!﹁", "'Meeting #123 @ 15:00 - Don't forget!'"),
|
||||
# Complex mixed quote types
|
||||
(
|
||||
'「Hello」, ❝World❞, "Test", \'Example\', „Quote", «Final»',
|
||||
'\'Hello\', "World", "Test", \'Example\', "Quote", "Final"',
|
||||
),
|
||||
# Quotes with multiple apostrophes
|
||||
("It's John's book, isn't it?", "It's John's book, isn't it?"),
|
||||
# Single angle quotes with nested content
|
||||
(
|
||||
'‹Testing the system\'s capability for "quoted" text›',
|
||||
"'Testing the system's capability for \"quoted\" text'",
|
||||
),
|
||||
# Heavy single ornament quotes with multiple sentences
|
||||
(
|
||||
"❛First sentence. Second sentence. Third sentence.❜",
|
||||
"'First sentence. Second sentence. Third sentence.'",
|
||||
),
|
||||
# Mix of various quote types in complex text
|
||||
(
|
||||
'「Chapter 1」: ❝The Beginning❞ - „A new story" begins «today».',
|
||||
'\'Chapter 1\': "The Beginning" - "A new story" begins "today".',
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_standardize_quotes(input_text, expected_output):
|
||||
assert text_extraction.standardize_quotes(input_text) == expected_output
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("output_text", "source_text", "expected_percentage"),
|
||||
[
|
||||
(
|
||||
"extra",
|
||||
"",
|
||||
0,
|
||||
),
|
||||
(
|
||||
"",
|
||||
"Source text has a sentence.",
|
||||
1,
|
||||
),
|
||||
(
|
||||
"The original s e n t e n c e is normal.",
|
||||
"The original sentence is normal...",
|
||||
0.2,
|
||||
),
|
||||
(
|
||||
"We saw 23% improvement in this quarter.",
|
||||
"We saw 23% improvement in sales this quarter.",
|
||||
0.125,
|
||||
),
|
||||
(
|
||||
"no",
|
||||
"Is it possible to have more than everything missing?",
|
||||
1,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_calculate_percent_missing_text(output_text, source_text, expected_percentage):
|
||||
assert (
|
||||
text_extraction.calculate_percent_missing_text(output_text, source_text)
|
||||
== expected_percentage
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("table_as_cells", "expected_extraction"),
|
||||
[
|
||||
pytest.param(
|
||||
[
|
||||
{"x": 0, "y": 0, "w": 1, "h": 1, "content": "Month A."},
|
||||
{"x": 0, "y": 1, "w": 1, "h": 1, "content": "22"},
|
||||
],
|
||||
[
|
||||
{"row_index": 0, "col_index": 0, "content": "Month A."},
|
||||
{"row_index": 1, "col_index": 0, "content": "22"},
|
||||
],
|
||||
id="Simple table, 1 head cell, 1 body cell, no spans",
|
||||
),
|
||||
pytest.param(
|
||||
[
|
||||
{"x": 0, "y": 0, "w": 1, "h": 1, "content": "Month A."},
|
||||
{"x": 1, "y": 0, "w": 1, "h": 1, "content": "Month B."},
|
||||
{"x": 2, "y": 0, "w": 1, "h": 1, "content": "Month C."},
|
||||
{"x": 0, "y": 1, "w": 1, "h": 1, "content": "11"},
|
||||
{"x": 1, "y": 1, "w": 1, "h": 1, "content": "12"},
|
||||
{"x": 2, "y": 1, "w": 1, "h": 1, "content": "13"},
|
||||
{"x": 0, "y": 2, "w": 1, "h": 1, "content": "21"},
|
||||
{"x": 1, "y": 2, "w": 1, "h": 1, "content": "22"},
|
||||
{"x": 2, "y": 2, "w": 1, "h": 1, "content": "23"},
|
||||
],
|
||||
[
|
||||
{"row_index": 0, "col_index": 0, "content": "Month A."},
|
||||
{"row_index": 0, "col_index": 1, "content": "Month B."},
|
||||
{"row_index": 0, "col_index": 2, "content": "Month C."},
|
||||
{"row_index": 1, "col_index": 0, "content": "11"},
|
||||
{"row_index": 1, "col_index": 1, "content": "12"},
|
||||
{"row_index": 1, "col_index": 2, "content": "13"},
|
||||
{"row_index": 2, "col_index": 0, "content": "21"},
|
||||
{"row_index": 2, "col_index": 1, "content": "22"},
|
||||
{"row_index": 2, "col_index": 2, "content": "23"},
|
||||
],
|
||||
id="Simple table, 3 head cell, 5 body cell, no spans",
|
||||
),
|
||||
# +----------+---------------------+----------+
|
||||
# | | h1col23 | h1col4 |
|
||||
# | h12col1 |----------+----------+----------|
|
||||
# | | h2col2 | h2col34 |
|
||||
# |----------|----------+----------+----------+
|
||||
# | r3col1 | r3col2 | |
|
||||
# |----------+----------| r34col34 |
|
||||
# | r4col12 | |
|
||||
# +----------+----------+----------+----------+
|
||||
pytest.param(
|
||||
[
|
||||
{
|
||||
"y": 0,
|
||||
"x": 0,
|
||||
"w": 2,
|
||||
"h": 1,
|
||||
"content": "h12col1",
|
||||
},
|
||||
{
|
||||
"y": 0,
|
||||
"x": 1,
|
||||
"w": 1,
|
||||
"h": 2,
|
||||
"content": "h1col23",
|
||||
},
|
||||
{
|
||||
"y": 0,
|
||||
"x": 3,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "h1col4",
|
||||
},
|
||||
{
|
||||
"y": 1,
|
||||
"x": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "h2col2",
|
||||
},
|
||||
{
|
||||
"y": 1,
|
||||
"x": 2,
|
||||
"w": 1,
|
||||
"h": 2,
|
||||
"content": "h2col34",
|
||||
},
|
||||
{
|
||||
"y": 2,
|
||||
"x": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r3col1",
|
||||
},
|
||||
{
|
||||
"y": 2,
|
||||
"x": 1,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "r3col2",
|
||||
},
|
||||
{
|
||||
"y": 2,
|
||||
"x": 2,
|
||||
"w": 2,
|
||||
"h": 2,
|
||||
"content": "r34col34",
|
||||
},
|
||||
{
|
||||
"y": 3,
|
||||
"x": 0,
|
||||
"w": 1,
|
||||
"h": 2,
|
||||
"content": "r4col12",
|
||||
},
|
||||
],
|
||||
[
|
||||
{
|
||||
"row_index": 0,
|
||||
"col_index": 0,
|
||||
"content": "h12col1",
|
||||
},
|
||||
{
|
||||
"row_index": 0,
|
||||
"col_index": 1,
|
||||
"content": "h1col23",
|
||||
},
|
||||
{
|
||||
"row_index": 0,
|
||||
"col_index": 3,
|
||||
"content": "h1col4",
|
||||
},
|
||||
{
|
||||
"row_index": 1,
|
||||
"col_index": 1,
|
||||
"content": "h2col2",
|
||||
},
|
||||
{
|
||||
"row_index": 1,
|
||||
"col_index": 2,
|
||||
"content": "h2col34",
|
||||
},
|
||||
{
|
||||
"row_index": 2,
|
||||
"col_index": 0,
|
||||
"content": "r3col1",
|
||||
},
|
||||
{
|
||||
"row_index": 2,
|
||||
"col_index": 1,
|
||||
"content": "r3col2",
|
||||
},
|
||||
{
|
||||
"row_index": 2,
|
||||
"col_index": 2,
|
||||
"content": "r34col34",
|
||||
},
|
||||
{
|
||||
"row_index": 3,
|
||||
"col_index": 0,
|
||||
"content": "r4col12",
|
||||
},
|
||||
],
|
||||
id="various spans, with 2 row header",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_cells_table_extraction_from_prediction(table_as_cells, expected_extraction):
|
||||
example_element = {
|
||||
"type": "Table",
|
||||
"metadata": {"table_as_cells": table_as_cells},
|
||||
}
|
||||
assert extract_cells_from_table_as_cells(example_element) == expected_extraction
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("text_as_html", "expected_extraction"),
|
||||
[
|
||||
pytest.param(
|
||||
"""
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Month A.</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>22</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>"
|
||||
""",
|
||||
[
|
||||
{"row_index": 0, "col_index": 0, "content": "Month A."},
|
||||
{"row_index": 1, "col_index": 0, "content": "22"},
|
||||
],
|
||||
id="Simple table, 1 head cell, 1 body cell, no spans",
|
||||
),
|
||||
pytest.param(
|
||||
"""
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Month A.</th>
|
||||
<th>Month B.</th>
|
||||
<th>Month C.</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>11</td>
|
||||
<td>12</td>
|
||||
<td>13</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>21</td>
|
||||
<td>22</td>
|
||||
<td>23</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>"
|
||||
""",
|
||||
[
|
||||
{"row_index": 0, "col_index": 0, "content": "Month A."},
|
||||
{"row_index": 0, "col_index": 1, "content": "Month B."},
|
||||
{"row_index": 0, "col_index": 2, "content": "Month C."},
|
||||
{"row_index": 1, "col_index": 0, "content": "11"},
|
||||
{"row_index": 1, "col_index": 1, "content": "12"},
|
||||
{"row_index": 1, "col_index": 2, "content": "13"},
|
||||
{"row_index": 2, "col_index": 0, "content": "21"},
|
||||
{"row_index": 2, "col_index": 1, "content": "22"},
|
||||
{"row_index": 2, "col_index": 2, "content": "23"},
|
||||
],
|
||||
id="Simple table, 3 head cell, 5 body cell, no spans",
|
||||
),
|
||||
# +----------+---------------------+----------+
|
||||
# | | h1col23 | h1col4 |
|
||||
# | h12col1 |----------+----------+----------|
|
||||
# | | h2col2 | h2col34 |
|
||||
# |----------|----------+----------+----------+
|
||||
# | r3col1 | r3col2 | |
|
||||
# |----------+----------| r34col34 |
|
||||
# | r4col12 | |
|
||||
# +----------+----------+----------+----------+
|
||||
pytest.param(
|
||||
"""
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th rowspan="2">h12col1</th>
|
||||
<th colspan="2">h1col23</th>
|
||||
<th>h1col4</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<th>h2col2</th>
|
||||
<th colspan="2">h2col34</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>r3col1</td>
|
||||
<td>r3col2</td>
|
||||
<td colspan="2" rowspan="2">r34col34</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td colspan="2">r4col12</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
""",
|
||||
[
|
||||
{
|
||||
"row_index": 0,
|
||||
"col_index": 0,
|
||||
"content": "h12col1",
|
||||
},
|
||||
{
|
||||
"row_index": 0,
|
||||
"col_index": 1,
|
||||
"content": "h1col23",
|
||||
},
|
||||
{
|
||||
"row_index": 0,
|
||||
"col_index": 3,
|
||||
"content": "h1col4",
|
||||
},
|
||||
{
|
||||
"row_index": 1,
|
||||
"col_index": 1,
|
||||
"content": "h2col2",
|
||||
},
|
||||
{
|
||||
"row_index": 1,
|
||||
"col_index": 2,
|
||||
"content": "h2col34",
|
||||
},
|
||||
{
|
||||
"row_index": 2,
|
||||
"col_index": 0,
|
||||
"content": "r3col1",
|
||||
},
|
||||
{
|
||||
"row_index": 2,
|
||||
"col_index": 1,
|
||||
"content": "r3col2",
|
||||
},
|
||||
{
|
||||
"row_index": 2,
|
||||
"col_index": 2,
|
||||
"content": "r34col34",
|
||||
},
|
||||
{
|
||||
"row_index": 3,
|
||||
"col_index": 0,
|
||||
"content": "r4col12",
|
||||
},
|
||||
],
|
||||
id="various spans, with 2 row header",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_html_table_extraction_from_prediction(text_as_html, expected_extraction):
|
||||
example_element = {
|
||||
"type": "Table",
|
||||
"metadata": {
|
||||
"text_as_html": text_as_html,
|
||||
},
|
||||
}
|
||||
assert extract_cells_from_text_as_html(example_element) == expected_extraction
|
||||
|
||||
|
||||
def test_cells_extraction_from_prediction_when_missing_prediction():
|
||||
example_element = {"type": "Table", "metadata": {"text_as_html": "", "table_as_cells": []}}
|
||||
assert extract_cells_from_text_as_html(example_element) is None
|
||||
assert extract_cells_from_table_as_cells(example_element) is None
|
||||
|
||||
|
||||
def _trim_html(html: str) -> str:
|
||||
html_lines = [line.strip() for line in html.split("\n") if line]
|
||||
return "".join(html_lines)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"html_to_test",
|
||||
[
|
||||
"""
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Month A.</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>22</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
""",
|
||||
"""
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Month A.</th>
|
||||
<th>Month B.</th>
|
||||
<th>Month C.</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>11</td>
|
||||
<td>12</td>
|
||||
<td>13</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td>21</td>
|
||||
<td>22</td>
|
||||
<td>23</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
""",
|
||||
"""
|
||||
<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th rowspan="2">h12col1</th>
|
||||
<th colspan="2">h1col23</th>
|
||||
<th>h1col4</th>
|
||||
</tr>
|
||||
<tr>
|
||||
<th>h2col2</th>
|
||||
<th colspan="2">h2col34</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
<tr>
|
||||
<td>r3col1</td>
|
||||
<td>r3col2</td>
|
||||
<td colspan="2" rowspan="2">r34col34</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td colspan="2">r4col12</td>
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>
|
||||
""",
|
||||
],
|
||||
)
|
||||
def test_deckerd_html_converter(html_to_test):
|
||||
deckerd_table = html_table_to_deckerd(html_to_test)
|
||||
html_table = deckerd_table_to_html(deckerd_table)
|
||||
assert _trim_html(html_to_test) == html_table
|
||||
@@ -0,0 +1,35 @@
|
||||
import pytest
|
||||
|
||||
from unstructured.metrics.utils import (
|
||||
_mean,
|
||||
_pstdev,
|
||||
_stdev,
|
||||
_uniquity_file,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("numbers", "expected_mean", "expected_stdev", "expected_pstdev"),
|
||||
[
|
||||
([2, 5, 6, 7], 5, 2.16, 1.871),
|
||||
([1, 100], 50.5, 70.004, 49.5),
|
||||
([1], 1, None, None),
|
||||
([], None, None, None),
|
||||
],
|
||||
)
|
||||
def test_stats(numbers, expected_mean, expected_stdev, expected_pstdev):
|
||||
mean = _mean(numbers)
|
||||
stdev = _stdev(numbers)
|
||||
pstdev = _pstdev(numbers)
|
||||
assert mean == expected_mean
|
||||
assert stdev == expected_stdev
|
||||
assert pstdev == expected_pstdev
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("filenames"),
|
||||
[("filename.ext", "filename (1).ext", "randomfile.ext", "filename.txt", "filename (5).txt")],
|
||||
)
|
||||
def test_uniquity_file(filenames):
|
||||
final_filename = _uniquity_file(filenames, "filename.ext")
|
||||
assert final_filename == "filename (2).ext"
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,21 @@
|
||||
from typing import List, Tuple
|
||||
|
||||
|
||||
def mock_sent_tokenize(text: str) -> List[str]:
|
||||
sentences = text.split(".")
|
||||
return sentences[:-1] if text.endswith(".") else sentences
|
||||
|
||||
|
||||
def mock_word_tokenize(text: str) -> List[str]:
|
||||
return text.split(" ")
|
||||
|
||||
|
||||
def mock_pos_tag(text: str) -> List[Tuple[str, str]]:
|
||||
tokens = mock_word_tokenize(text)
|
||||
pos_tags: List[Tuple[str, str]] = []
|
||||
for token in tokens:
|
||||
if token.lower() == "ask":
|
||||
pos_tags.append((token, "VB"))
|
||||
else:
|
||||
pos_tags.append((token, ""))
|
||||
return pos_tags
|
||||
@@ -0,0 +1 @@
|
||||
# flake8: noqa
|
||||
@@ -0,0 +1,59 @@
|
||||
from typing import List, Tuple
|
||||
|
||||
from test_unstructured.nlp.mock_nltk import mock_sent_tokenize, mock_word_tokenize
|
||||
from unstructured.nlp import tokenize
|
||||
|
||||
|
||||
def mock_pos_tag(tokens: List[str]) -> List[Tuple[str, str]]:
|
||||
pos_tags: List[Tuple[str, str]] = []
|
||||
for token in tokens:
|
||||
if token.lower() == "ask":
|
||||
pos_tags.append((token, "VB"))
|
||||
else:
|
||||
pos_tags.append((token, ""))
|
||||
return pos_tags
|
||||
|
||||
|
||||
def test_pos_tag():
|
||||
parts_of_speech = tokenize.pos_tag("ITEM 2A. PROPERTIES")
|
||||
assert parts_of_speech == [
|
||||
("ITEM", "NNP"),
|
||||
("2A", "CD"),
|
||||
(".", "."),
|
||||
("PROPERTIES", "NN"),
|
||||
]
|
||||
|
||||
|
||||
def test_word_tokenize_caches(monkeypatch):
|
||||
monkeypatch.setattr(tokenize, "_word_tokenize", mock_word_tokenize)
|
||||
monkeypatch.setattr(tokenize, "_pos_tag", mock_pos_tag)
|
||||
tokenize.word_tokenize.cache_clear()
|
||||
assert tokenize.word_tokenize.cache_info().currsize == 0
|
||||
tokenize.word_tokenize("Greetings! I am from outer space.")
|
||||
assert tokenize.word_tokenize.cache_info().currsize == 1
|
||||
|
||||
|
||||
def test_sent_tokenize_caches(monkeypatch):
|
||||
monkeypatch.setattr(tokenize, "_sent_tokenize", mock_sent_tokenize)
|
||||
monkeypatch.setattr(tokenize, "_word_tokenize", mock_word_tokenize)
|
||||
monkeypatch.setattr(tokenize, "_pos_tag", mock_pos_tag)
|
||||
tokenize._tokenize_for_cache.cache_clear()
|
||||
assert tokenize._tokenize_for_cache.cache_info().currsize == 0
|
||||
tokenize._tokenize_for_cache("Greetings! I am from outer space.")
|
||||
assert tokenize._tokenize_for_cache.cache_info().currsize == 1
|
||||
|
||||
|
||||
def test_pos_tag_caches(monkeypatch):
|
||||
monkeypatch.setattr(tokenize, "_word_tokenize", mock_word_tokenize)
|
||||
monkeypatch.setattr(tokenize, "_pos_tag", mock_pos_tag)
|
||||
tokenize.pos_tag.cache_clear()
|
||||
assert tokenize.pos_tag.cache_info().currsize == 0
|
||||
tokenize.pos_tag("Greetings! I am from outer space.")
|
||||
assert tokenize.pos_tag.cache_info().currsize == 1
|
||||
|
||||
|
||||
def test_tokenizers_functions_run():
|
||||
sentence = "I am a big brown bear. What are you?"
|
||||
tokenize.sent_tokenize(sentence)
|
||||
tokenize.word_tokenize(sentence)
|
||||
tokenize.pos_tag(sentence)
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user