修改为东南天坐标系
This commit is contained in:
@@ -0,0 +1,4 @@
|
||||
from .partition.utils.config import env_config
|
||||
|
||||
# init env_config
|
||||
env_config
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1 @@
|
||||
__version__ = "0.18.27" # pragma: no cover
|
||||
@@ -0,0 +1,22 @@
|
||||
"""Chunking module initializer.
|
||||
|
||||
Publishes the public aspects of the chunking sub-package interface.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from unstructured.chunking.base import CHUNK_MAX_CHARS_DEFAULT, CHUNK_MULTI_PAGE_DEFAULT
|
||||
from unstructured.chunking.dispatch import (
|
||||
Chunker,
|
||||
add_chunking_strategy,
|
||||
register_chunking_strategy,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"CHUNK_MAX_CHARS_DEFAULT",
|
||||
"CHUNK_MULTI_PAGE_DEFAULT",
|
||||
"add_chunking_strategy",
|
||||
# -- these must be published to allow pluggable chunkers in other code-bases --
|
||||
"Chunker",
|
||||
"register_chunking_strategy",
|
||||
]
|
||||
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,91 @@
|
||||
"""Implementation of baseline chunking.
|
||||
|
||||
This is the "plain-vanilla" chunking strategy. All the fundamental chunking behaviors are present in
|
||||
this strategy and also in all other strategies. Those are:
|
||||
|
||||
- Maximally fill each chunk with sequential elements.
|
||||
- Isolate oversized elements and divide (only) those chunks by text-splitting.
|
||||
- Overlap when requested.
|
||||
|
||||
"Fancier" strategies add higher-level semantic-unit boundaries to be respected. For example, in the
|
||||
by-title strategy, section boundaries are respected, meaning a chunk never contains text from two
|
||||
different sections. When a new section is detected the current chunk is closed and a new one
|
||||
started.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Iterable, Optional
|
||||
|
||||
from unstructured.chunking.base import ChunkingOptions, PreChunker
|
||||
from unstructured.documents.elements import Element
|
||||
|
||||
|
||||
def chunk_elements(
|
||||
elements: Iterable[Element],
|
||||
*,
|
||||
include_orig_elements: Optional[bool] = None,
|
||||
max_characters: Optional[int] = None,
|
||||
new_after_n_chars: Optional[int] = None,
|
||||
overlap: Optional[int] = None,
|
||||
overlap_all: Optional[bool] = None,
|
||||
) -> list[Element]:
|
||||
"""Combine sequential `elements` into chunks, respecting specified text-length limits.
|
||||
|
||||
Produces a sequence of `CompositeElement`, `Table`, and `TableChunk` elements (chunks).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
elements
|
||||
A list of unstructured elements. Usually the output of a partition function.
|
||||
include_orig_elements
|
||||
When `True` (default), add elements from pre-chunk to the `.metadata.orig_elements` field
|
||||
of the chunk(s) formed from that pre-chunk. Among other things, this allows access to
|
||||
original-element metadata that cannot be consolidated and is dropped in the course of
|
||||
chunking.
|
||||
max_characters
|
||||
Hard maximum chunk length. No chunk will exceed this length. A single element that exceeds
|
||||
this length will be divided into two or more chunks using text-splitting.
|
||||
new_after_n_chars
|
||||
A chunk that of this length or greater is not extended to include the next element, even if
|
||||
that element would fit without exceeding `max_characters`. A "soft max" length that can be
|
||||
used in conjunction with `max_characters` to limit most chunks to a preferred length while
|
||||
still allowing larger elements to be included in a single chunk without resorting to
|
||||
text-splitting. Defaults to `max_characters` when not specified, which effectively disables
|
||||
any soft window. Specifying 0 for this argument causes each element to appear in a chunk by
|
||||
itself (although an element with text longer than `max_characters` will be still be split
|
||||
into two or more chunks).
|
||||
overlap
|
||||
Specifies the length of a string ("tail") to be drawn from each chunk and prefixed to the
|
||||
next chunk as a context-preserving mechanism. By default, this only applies to split-chunks
|
||||
where an oversized element is divided into multiple chunks by text-splitting.
|
||||
overlap_all
|
||||
Default: `False`. When `True`, apply overlap between "normal" chunks formed from whole
|
||||
elements and not subject to text-splitting. Use this with caution as it produces a certain
|
||||
level of "pollution" of otherwise clean semantic chunk boundaries.
|
||||
"""
|
||||
# -- raises ValueError on invalid parameters --
|
||||
opts = _BasicChunkingOptions.new(
|
||||
include_orig_elements=include_orig_elements,
|
||||
max_characters=max_characters,
|
||||
new_after_n_chars=new_after_n_chars,
|
||||
overlap=overlap,
|
||||
overlap_all=overlap_all,
|
||||
)
|
||||
|
||||
return _chunk_elements(elements, opts)
|
||||
|
||||
|
||||
def _chunk_elements(elements: Iterable[Element], opts: _BasicChunkingOptions) -> list[Element]:
|
||||
"""Implementation of actual basic chunking."""
|
||||
# -- Note(scanny): it might seem like over-abstraction for this to be a separate function but
|
||||
# -- it eases overriding or adding individual chunking options when customizing a stock chunker.
|
||||
return [
|
||||
chunk
|
||||
for pre_chunk in PreChunker.iter_pre_chunks(elements, opts)
|
||||
for chunk in pre_chunk.iter_chunks()
|
||||
]
|
||||
|
||||
|
||||
class _BasicChunkingOptions(ChunkingOptions):
|
||||
"""Options for `basic` chunking."""
|
||||
@@ -0,0 +1,129 @@
|
||||
"""Handles dispatch of elements to a chunking-strategy by name.
|
||||
|
||||
Also provides the `@add_chunking_strategy` decorator which is the chief current user of "by-name"
|
||||
chunking dispatch.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import dataclasses as dc
|
||||
import functools
|
||||
import inspect
|
||||
from typing import Any, Callable, Iterable, Optional, Protocol
|
||||
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
from unstructured.chunking.basic import chunk_elements
|
||||
from unstructured.chunking.title import chunk_by_title
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.utils import get_call_args_applying_defaults, lazyproperty
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
|
||||
class Chunker(Protocol):
|
||||
"""Abstract interface for chunking functions."""
|
||||
|
||||
def __call__(
|
||||
self, elements: Iterable[Element], *, max_characters: Optional[int]
|
||||
) -> list[Element]:
|
||||
"""A chunking function must have this signature.
|
||||
|
||||
In particular it must minimally have an `elements` parameter and all chunkers will have a
|
||||
`max_characters` parameter (doesn't need to follow `elements` directly). All others can
|
||||
vary by chunker.
|
||||
"""
|
||||
...
|
||||
|
||||
|
||||
def add_chunking_strategy(func: Callable[_P, list[Element]]) -> Callable[_P, list[Element]]:
|
||||
"""Decorator for chunking text.
|
||||
|
||||
Chunks the element sequence produced by the partitioner it decorates when a `chunking_strategy`
|
||||
argument is present in the partitioner call and it names an available chunking strategy.
|
||||
"""
|
||||
# -- Patch the docstring of the decorated function to add chunking strategy and
|
||||
# -- chunking-related argument documentation. This only applies when `chunking_strategy`
|
||||
# -- is an explicit argument of the decorated function and "chunking_strategy" is not
|
||||
# -- already mentioned in the docstring.
|
||||
if func.__doc__ and (
|
||||
"chunking_strategy" in func.__code__.co_varnames and "chunking_strategy" not in func.__doc__
|
||||
):
|
||||
func.__doc__ += (
|
||||
"\nchunking_strategy"
|
||||
+ "\n\tStrategy used for chunking text into larger or smaller elements."
|
||||
+ "\n\tDefaults to `None` with optional arg of 'basic' or 'by_title'."
|
||||
+ "\n\tAdditional Parameters:"
|
||||
+ "\n\t\tmultipage_sections"
|
||||
+ "\n\t\t\tIf True, sections can span multiple pages. Defaults to True."
|
||||
+ "\n\t\tcombine_text_under_n_chars"
|
||||
+ "\n\t\t\tCombines elements (for example a series of titles) until a section"
|
||||
+ "\n\t\t\treaches a length of n characters. Only applies to 'by_title' strategy."
|
||||
+ "\n\t\tnew_after_n_chars"
|
||||
+ "\n\t\t\tCuts off chunks once they reach a length of n characters; a soft max."
|
||||
+ "\n\t\tmax_characters"
|
||||
+ "\n\t\t\tChunks elements text and text_as_html (if present) into chunks"
|
||||
+ "\n\t\t\tof length n characters, a hard max."
|
||||
)
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> list[Element]:
|
||||
"""The decorated function is replaced with this one."""
|
||||
|
||||
# -- call the partitioning function to get the elements --
|
||||
elements = func(*args, **kwargs)
|
||||
|
||||
# -- look for a chunking-strategy argument --
|
||||
call_args = get_call_args_applying_defaults(func, *args, **kwargs)
|
||||
chunking_strategy = call_args.pop("chunking_strategy", None)
|
||||
|
||||
# -- no chunking-strategy means no chunking --
|
||||
if chunking_strategy is None:
|
||||
return elements
|
||||
|
||||
# -- otherwise, chunk away :) --
|
||||
return chunk(elements, chunking_strategy, **call_args)
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def chunk(elements: Iterable[Element], chunking_strategy: str, **kwargs: Any) -> list[Element]:
|
||||
"""Dispatch chunking of `elements` to the chunking function for `chunking_strategy`."""
|
||||
chunker_spec = _chunker_registry.get(chunking_strategy)
|
||||
|
||||
if chunker_spec is None:
|
||||
raise ValueError(f"unrecognized chunking strategy {repr(chunking_strategy)}")
|
||||
|
||||
# -- `kwargs` will in general be an omnibus dict of all keyword arguments to the partitioner;
|
||||
# -- pick out and use only those supported by this chunker.
|
||||
chunking_kwargs = {k: v for k, v in kwargs.items() if k in chunker_spec.kw_arg_names}
|
||||
|
||||
return chunker_spec.chunker(elements, **chunking_kwargs)
|
||||
|
||||
|
||||
def register_chunking_strategy(name: str, chunker: Chunker) -> None:
|
||||
"""Make chunker available by using `name` as `chunking_strategy` arg in partitioner call."""
|
||||
_chunker_registry[name] = _ChunkerSpec(chunker)
|
||||
|
||||
|
||||
@dc.dataclass(frozen=True)
|
||||
class _ChunkerSpec:
|
||||
"""A registry entry for a chunker."""
|
||||
|
||||
chunker: Chunker
|
||||
"""The "chunk_by_{x}() function that implements this chunking strategy."""
|
||||
|
||||
@lazyproperty
|
||||
def kw_arg_names(self) -> tuple[str, ...]:
|
||||
"""Keyword arguments supported by this chunker.
|
||||
|
||||
These are all arguments other than the required `elements: list[Element]` first parameter.
|
||||
"""
|
||||
sig = inspect.signature(self.chunker)
|
||||
return tuple(key for key in sig.parameters if key != "elements")
|
||||
|
||||
|
||||
_chunker_registry: dict[str, _ChunkerSpec] = {
|
||||
"basic": _ChunkerSpec(chunk_elements),
|
||||
"by_title": _ChunkerSpec(chunk_by_title),
|
||||
}
|
||||
@@ -0,0 +1,167 @@
|
||||
"""Implementation of chunking by title.
|
||||
|
||||
Main entry point is the `@add_chunking_strategy()` decorator.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Iterable, Iterator, Optional
|
||||
|
||||
from unstructured.chunking.base import (
|
||||
CHUNK_MULTI_PAGE_DEFAULT,
|
||||
BoundaryPredicate,
|
||||
ChunkingOptions,
|
||||
PreChunkCombiner,
|
||||
PreChunker,
|
||||
is_on_next_page,
|
||||
is_title,
|
||||
)
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.utils import lazyproperty
|
||||
|
||||
|
||||
def chunk_by_title(
|
||||
elements: Iterable[Element],
|
||||
*,
|
||||
combine_text_under_n_chars: Optional[int] = None,
|
||||
include_orig_elements: Optional[bool] = None,
|
||||
max_characters: Optional[int] = None,
|
||||
multipage_sections: Optional[bool] = None,
|
||||
new_after_n_chars: Optional[int] = None,
|
||||
overlap: Optional[int] = None,
|
||||
overlap_all: Optional[bool] = None,
|
||||
) -> list[Element]:
|
||||
"""Uses title elements to identify sections within the document for chunking.
|
||||
|
||||
Splits off into a new CompositeElement when a title is detected or if metadata changes, which
|
||||
happens when page numbers or sections change. Cuts off sections once they have exceeded a
|
||||
character length of max_characters.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
elements
|
||||
A list of unstructured elements. Usually the output of a partition function.
|
||||
combine_text_under_n_chars
|
||||
Combines elements (for example a series of titles) until a section reaches a length of
|
||||
n characters. Defaults to `max_characters` which combines chunks whenever space allows.
|
||||
Specifying 0 for this argument suppresses combining of small chunks. Note this value is
|
||||
"capped" at the `new_after_n_chars` value since a value higher than that would not change
|
||||
this parameter's effect.
|
||||
include_orig_elements
|
||||
When `True` (default), add elements from pre-chunk to the `.metadata.orig_elements` field
|
||||
of the chunk(s) formed from that pre-chunk. Among other things, this allows access to
|
||||
original-element metadata that cannot be consolidated and is dropped in the course of
|
||||
chunking.
|
||||
max_characters
|
||||
Chunks elements text and text_as_html (if present) into chunks of length
|
||||
n characters (hard max)
|
||||
multipage_sections
|
||||
If True, sections can span multiple pages. Defaults to True.
|
||||
new_after_n_chars
|
||||
Cuts off new sections once they reach a length of n characters (soft max). Defaults to
|
||||
`max_characters` when not specified, which effectively disables any soft window.
|
||||
Specifying 0 for this argument causes each element to appear in a chunk by itself (although
|
||||
an element with text longer than `max_characters` will be still be split into two or more
|
||||
chunks).
|
||||
overlap
|
||||
Specifies the length of a string ("tail") to be drawn from each chunk and prefixed to the
|
||||
next chunk as a context-preserving mechanism. By default, this only applies to split-chunks
|
||||
where an oversized element is divided into multiple chunks by text-splitting.
|
||||
overlap_all
|
||||
Default: `False`. When `True`, apply overlap between "normal" chunks formed from whole
|
||||
elements and not subject to text-splitting. Use this with caution as it entails a certain
|
||||
level of "pollution" of otherwise clean semantic chunk boundaries.
|
||||
"""
|
||||
opts = _ByTitleChunkingOptions.new(
|
||||
combine_text_under_n_chars=combine_text_under_n_chars,
|
||||
include_orig_elements=include_orig_elements,
|
||||
max_characters=max_characters,
|
||||
multipage_sections=multipage_sections,
|
||||
new_after_n_chars=new_after_n_chars,
|
||||
overlap=overlap,
|
||||
overlap_all=overlap_all,
|
||||
)
|
||||
return _chunk_by_title(elements, opts)
|
||||
|
||||
|
||||
def _chunk_by_title(elements: Iterable[Element], opts: _ByTitleChunkingOptions) -> list[Element]:
|
||||
"""Implementation of actual "by-title" chunking."""
|
||||
# -- Note(scanny): it might seem like over-abstraction for this to be a separate function but
|
||||
# -- it eases overriding or adding individual chunking options when customizing a stock chunker.
|
||||
pre_chunks = PreChunkCombiner(
|
||||
PreChunker.iter_pre_chunks(elements, opts), opts=opts
|
||||
).iter_combined_pre_chunks()
|
||||
|
||||
return [chunk for pre_chunk in pre_chunks for chunk in pre_chunk.iter_chunks()]
|
||||
|
||||
|
||||
class _ByTitleChunkingOptions(ChunkingOptions):
|
||||
"""Adds the by-title-specific chunking options to the base case.
|
||||
|
||||
`by_title`-specific options:
|
||||
|
||||
combine_text_under_n_chars
|
||||
A remedy to over-chunking caused by elements mis-identified as Title elements.
|
||||
Every Title element would start a new chunk and this setting mitigates that, at the
|
||||
expense of sometimes violating legitimate semantic boundaries.
|
||||
multipage_sections
|
||||
Indicates that page-boundaries should not be respected while chunking, i.e. elements
|
||||
appearing on two different pages can appear in the same chunk.
|
||||
"""
|
||||
|
||||
@lazyproperty
|
||||
def boundary_predicates(self) -> tuple[BoundaryPredicate, ...]:
|
||||
"""The semantic-boundary detectors to be applied to break pre-chunks.
|
||||
|
||||
For the `by_title` strategy these are sections indicated by a title (section-heading), an
|
||||
explicit section metadata item (only present for certain document types), and optionally
|
||||
page boundaries.
|
||||
"""
|
||||
|
||||
def iter_boundary_predicates() -> Iterator[BoundaryPredicate]:
|
||||
yield is_title
|
||||
if not self.multipage_sections:
|
||||
yield is_on_next_page()
|
||||
|
||||
return tuple(iter_boundary_predicates())
|
||||
|
||||
@lazyproperty
|
||||
def combine_text_under_n_chars(self) -> int:
|
||||
"""Combine consecutive text pre-chunks if former is smaller than this and both will fit.
|
||||
|
||||
- Does not combine text chunks if together they would exceed the chunking window.
|
||||
- Defaults to `max_characters` when not specified.
|
||||
- Is reduced to `new_after_n_chars` when it exceeds that value.
|
||||
"""
|
||||
# -- `combine_text_under_n_chars` defaults to `max_characters` when not specified --
|
||||
arg_value = self._kwargs.get("combine_text_under_n_chars")
|
||||
return self.hard_max if arg_value is None else arg_value
|
||||
|
||||
@lazyproperty
|
||||
def multipage_sections(self) -> bool:
|
||||
"""When False, break pre-chunks on page-boundaries."""
|
||||
arg_value = self._kwargs.get("multipage_sections")
|
||||
return CHUNK_MULTI_PAGE_DEFAULT if arg_value is None else bool(arg_value)
|
||||
|
||||
def _validate(self) -> None:
|
||||
"""Raise ValueError if request option-set is invalid."""
|
||||
# -- start with base-class validations --
|
||||
super()._validate()
|
||||
|
||||
# -- `combine_text_under_n_chars == 0` is valid (suppresses chunk combination)
|
||||
# -- but a negative value is not
|
||||
if self.combine_text_under_n_chars < 0:
|
||||
raise ValueError(
|
||||
f"'combine_text_under_n_chars' argument must be >= 0,"
|
||||
f" got {self.combine_text_under_n_chars}"
|
||||
)
|
||||
|
||||
# -- `combine_text_under_n_chars` > `max_characters` can produce behavior confusing to
|
||||
# -- users. The chunking behavior would be no different than when
|
||||
# -- `combine_text_under_n_chars == max_characters`, but if `max_characters` is left to
|
||||
# -- default (500) then it can look like chunk-combining isn't working.
|
||||
if self.combine_text_under_n_chars > self.hard_max:
|
||||
raise ValueError(
|
||||
f"'combine_text_under_n_chars' argument must not exceed `max_characters`"
|
||||
f" value, got {self.combine_text_under_n_chars} > {self.hard_max}"
|
||||
)
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,480 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import quopri
|
||||
import re
|
||||
import sys
|
||||
import unicodedata
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from unstructured.file_utils.encoding import (
|
||||
format_encoding_str,
|
||||
)
|
||||
from unstructured.nlp.patterns import (
|
||||
DOUBLE_PARAGRAPH_PATTERN_RE,
|
||||
E_BULLET_PATTERN,
|
||||
LINE_BREAK_RE,
|
||||
PARAGRAPH_PATTERN,
|
||||
PARAGRAPH_PATTERN_RE,
|
||||
UNICODE_BULLETS_RE,
|
||||
UNICODE_BULLETS_RE_0W,
|
||||
)
|
||||
|
||||
|
||||
def clean_non_ascii_chars(text) -> str:
|
||||
"""Cleans non-ascii characters from unicode string.
|
||||
|
||||
Example
|
||||
-------
|
||||
\x88This text contains non-ascii characters!\x88
|
||||
-> This text contains non-ascii characters!
|
||||
"""
|
||||
en = text.encode("ascii", "ignore")
|
||||
return en.decode()
|
||||
|
||||
|
||||
def clean_bullets(text: str) -> str:
|
||||
"""Cleans unicode bullets from a section of text.
|
||||
|
||||
Example
|
||||
-------
|
||||
● This is an excellent point! -> This is an excellent point!
|
||||
"""
|
||||
search = UNICODE_BULLETS_RE.match(text)
|
||||
if search is None:
|
||||
return text
|
||||
|
||||
cleaned_text = UNICODE_BULLETS_RE.sub("", text, 1)
|
||||
return cleaned_text.strip()
|
||||
|
||||
|
||||
def clean_ordered_bullets(text) -> str:
|
||||
"""Cleans the start of bulleted text sections up to three “sub-section”
|
||||
bullets accounting numeric and alphanumeric types.
|
||||
|
||||
Example
|
||||
-------
|
||||
1.1 This is a very important point -> This is a very important point
|
||||
a.b This is a very important point -> This is a very important point
|
||||
"""
|
||||
text_sp = text.split()
|
||||
text_cl = " ".join(text_sp[1:])
|
||||
if any(["." not in text_sp[0], ".." in text_sp[0]]):
|
||||
return text
|
||||
|
||||
bullet = re.split(pattern=r"[\.]", string=text_sp[0])
|
||||
if not bullet[-1]:
|
||||
del bullet[-1]
|
||||
|
||||
if len(bullet[0]) > 2:
|
||||
return text
|
||||
|
||||
return text_cl
|
||||
|
||||
|
||||
def clean_ligatures(text) -> str:
|
||||
"""Replaces ligatures with their most likely equivalent characters.
|
||||
|
||||
Example
|
||||
-------
|
||||
The benefits -> The benefits
|
||||
High quality financial -> High quality financial
|
||||
"""
|
||||
ligatures_map = {
|
||||
"æ": "ae",
|
||||
"Æ": "AE",
|
||||
"ff": "ff",
|
||||
"fi": "fi",
|
||||
"fl": "fl",
|
||||
"ffi": "ffi",
|
||||
"ffl": "ffl",
|
||||
"ſt": "ft",
|
||||
"ʪ": "ls",
|
||||
"œ": "oe",
|
||||
"Œ": "OE",
|
||||
"ȹ": "qp",
|
||||
"st": "st",
|
||||
"ʦ": "ts",
|
||||
}
|
||||
cleaned_text: str = text
|
||||
for k, v in ligatures_map.items():
|
||||
cleaned_text = cleaned_text.replace(k, v)
|
||||
|
||||
return cleaned_text
|
||||
|
||||
|
||||
def group_bullet_paragraph(paragraph: str) -> list:
|
||||
"""Groups paragraphs with bullets that have line breaks for visual/formatting purposes.
|
||||
For example:
|
||||
|
||||
'''○ The big red fox
|
||||
is walking down the lane.
|
||||
|
||||
○ At the end of the lane
|
||||
the fox met a friendly bear.'''
|
||||
|
||||
Gets converted to
|
||||
|
||||
'''○ The big red fox is walking down the lane.
|
||||
○ At the end of the land the fox met a bear.'''
|
||||
"""
|
||||
paragraph_pattern_re = re.compile(PARAGRAPH_PATTERN)
|
||||
|
||||
# pytesseract converts some bullet points to standalone "e" characters.
|
||||
# Substitute "e" with bullets since they are later used in partition_text
|
||||
# to determine list element type.
|
||||
paragraph = E_BULLET_PATTERN.sub("·", paragraph).strip()
|
||||
|
||||
bullet_paras = UNICODE_BULLETS_RE_0W.split(paragraph)
|
||||
clean_paragraphs = []
|
||||
for bullet in bullet_paras:
|
||||
if bullet:
|
||||
clean_paragraphs.append(paragraph_pattern_re.sub(" ", bullet))
|
||||
return clean_paragraphs
|
||||
|
||||
|
||||
def group_broken_paragraphs(
|
||||
text: str,
|
||||
line_split: re.Pattern[str] = PARAGRAPH_PATTERN_RE,
|
||||
paragraph_split: re.Pattern[str] = DOUBLE_PARAGRAPH_PATTERN_RE,
|
||||
) -> str:
|
||||
"""Groups paragraphs that have line breaks for visual/formatting purposes.
|
||||
For example:
|
||||
|
||||
'''The big red fox
|
||||
is walking down the lane.
|
||||
|
||||
At the end of the lane
|
||||
the fox met a bear.'''
|
||||
|
||||
Gets converted to
|
||||
|
||||
'''The big red fox is walking down the lane.
|
||||
At the end of the land the fox met a bear.'''
|
||||
"""
|
||||
paragraph_pattern_re = (
|
||||
PARAGRAPH_PATTERN
|
||||
if isinstance(PARAGRAPH_PATTERN, re.Pattern)
|
||||
else re.compile(PARAGRAPH_PATTERN)
|
||||
)
|
||||
|
||||
paragraphs = paragraph_split.split(text)
|
||||
clean_paragraphs = []
|
||||
for paragraph in paragraphs:
|
||||
stripped_par = paragraph.strip()
|
||||
if not stripped_par:
|
||||
continue
|
||||
|
||||
if UNICODE_BULLETS_RE.match(stripped_par) or E_BULLET_PATTERN.match(stripped_par):
|
||||
clean_paragraphs.extend(group_bullet_paragraph(paragraph))
|
||||
continue
|
||||
# NOTE(robinson) - This block is to account for lines like the following that shouldn't be
|
||||
# grouped together, but aren't separated by a double line break.
|
||||
# Apache License
|
||||
# Version 2.0, January 2004
|
||||
# http://www.apache.org/licenses/
|
||||
para_split = line_split.split(paragraph)
|
||||
all_lines_short = all(len(line.strip().split(" ")) < 5 for line in para_split)
|
||||
if all_lines_short:
|
||||
clean_paragraphs.extend(line for line in para_split if line.strip())
|
||||
else:
|
||||
clean_paragraphs.append(paragraph_pattern_re.sub(" ", paragraph))
|
||||
|
||||
return "\n\n".join(clean_paragraphs)
|
||||
|
||||
|
||||
def new_line_grouper(
|
||||
text: str,
|
||||
paragraph_split: re.Pattern[str] = LINE_BREAK_RE,
|
||||
) -> str:
|
||||
"""
|
||||
Concatenates text document that has one-line paragraph break pattern
|
||||
|
||||
For example,
|
||||
|
||||
Iwan Roberts
|
||||
Roberts celebrating after scoring a goal for Norwich City
|
||||
in 2004
|
||||
|
||||
Will be returned as:
|
||||
|
||||
Iwan Roberts\n\nRoberts celebrating after scoring a goal for Norwich City\n\nin 2004
|
||||
"""
|
||||
paragraphs = paragraph_split.split(text)
|
||||
clean_paragraphs = []
|
||||
for paragraph in paragraphs:
|
||||
if not paragraph.strip():
|
||||
continue
|
||||
clean_paragraphs.append(paragraph)
|
||||
return "\n\n".join(clean_paragraphs)
|
||||
|
||||
|
||||
def blank_line_grouper(
|
||||
text: str,
|
||||
paragraph_split: re.Pattern = DOUBLE_PARAGRAPH_PATTERN_RE,
|
||||
) -> str:
|
||||
"""
|
||||
Concatenates text document that has blank-line paragraph break pattern
|
||||
|
||||
For example,
|
||||
|
||||
Vestibulum auctor dapibus neque.
|
||||
|
||||
Nunc dignissim risus id metus.
|
||||
|
||||
Will be returned as:
|
||||
|
||||
Vestibulum auctor dapibus neque.\n\nNunc dignissim risus id metus.\n\n
|
||||
|
||||
"""
|
||||
return group_broken_paragraphs(text)
|
||||
|
||||
|
||||
def auto_paragraph_grouper(
|
||||
text: str,
|
||||
line_split: re.Pattern[str] = LINE_BREAK_RE,
|
||||
max_line_count: int = 2000,
|
||||
threshold: float = 0.1,
|
||||
) -> str:
|
||||
"""
|
||||
Checks the ratio of new line (\n) over the total max_line_count
|
||||
|
||||
If the ratio of new line is less than the threshold,
|
||||
the document is considered a new-line grouping type
|
||||
and return the original text
|
||||
|
||||
If the ratio of new line is greater than or equal to the threshold,
|
||||
the document is considered a blank-line grouping type
|
||||
and passed on to blank_line_grouper function
|
||||
"""
|
||||
lines = line_split.split(text)
|
||||
max_line_count = min(len(lines), max_line_count)
|
||||
line_count, empty_line_count = 0, 0
|
||||
for line in lines[:max_line_count]:
|
||||
line_count += 1
|
||||
if not line.strip():
|
||||
empty_line_count += 1
|
||||
ratio = empty_line_count / line_count
|
||||
|
||||
# NOTE(klaijan) - for ratio < threshold, we pass to new-line grouper,
|
||||
# otherwise to blank-line grouper
|
||||
if ratio < threshold:
|
||||
return new_line_grouper(text)
|
||||
else:
|
||||
return blank_line_grouper(text)
|
||||
|
||||
|
||||
# TODO(robinson) - There's likely a cleaner was to accomplish this and get all of the
|
||||
# unicode characters instead of just the quotes. Doing this for now since quotes are
|
||||
# an issue that are popping up in the SEC filings tests
|
||||
def replace_unicode_quotes(text: str) -> str:
|
||||
"""Replaces unicode bullets in text with the expected character
|
||||
|
||||
Example
|
||||
-------
|
||||
\x93What a lovely quote!\x94 -> “What a lovely quote!”
|
||||
"""
|
||||
# NOTE(robinson) - We should probably make this something more sane like a regex
|
||||
# instead of a whole big series of replaces
|
||||
text = text.replace("\x91", "‘")
|
||||
text = text.replace("\x92", "’")
|
||||
text = text.replace("\x93", "“")
|
||||
text = text.replace("\x94", "”")
|
||||
text = text.replace("'", "'")
|
||||
text = text.replace("â\x80\x99", "'")
|
||||
text = text.replace("â\x80“", "—")
|
||||
text = text.replace("â\x80”", "–")
|
||||
text = text.replace("â\x80˜", "‘")
|
||||
text = text.replace("â\x80¦", "…")
|
||||
text = text.replace("â\x80™", "’")
|
||||
text = text.replace("â\x80œ", "“")
|
||||
text = text.replace("â\x80?", "”")
|
||||
text = text.replace("â\x80ť", "”")
|
||||
text = text.replace("â\x80ś", "“")
|
||||
text = text.replace("â\x80¨", "—")
|
||||
text = text.replace("â\x80ł", "″")
|
||||
text = text.replace("â\x80Ž", "")
|
||||
text = text.replace("â\x80‚", "")
|
||||
text = text.replace("â\x80‰", "")
|
||||
text = text.replace("â\x80‹", "")
|
||||
text = text.replace("â\x80", "")
|
||||
text = text.replace("â\x80s'", "")
|
||||
return text
|
||||
|
||||
|
||||
tbl = dict.fromkeys(
|
||||
i for i in range(sys.maxunicode) if unicodedata.category(chr(i)).startswith("P")
|
||||
)
|
||||
|
||||
|
||||
def remove_punctuation(s: str) -> str:
|
||||
"""Removes punctuation from a given string."""
|
||||
return s.translate(tbl)
|
||||
|
||||
|
||||
def remove_sentence_punctuation(s: str, exclude_punctuation: Optional[list]) -> str:
|
||||
tbl_new = tbl.copy()
|
||||
if exclude_punctuation:
|
||||
for punct in exclude_punctuation:
|
||||
del tbl_new[ord(punct)]
|
||||
s = s.translate(tbl_new)
|
||||
return s
|
||||
|
||||
|
||||
def clean_extra_whitespace(text: str) -> str:
|
||||
"""Cleans extra whitespace characters that appear between words.
|
||||
|
||||
Example
|
||||
-------
|
||||
ITEM 1. BUSINESS -> ITEM 1. BUSINESS
|
||||
"""
|
||||
cleaned_text = re.sub(r"[\xa0\n]", " ", text)
|
||||
cleaned_text = re.sub(r"([ ]{2,})", " ", cleaned_text)
|
||||
return cleaned_text.strip()
|
||||
|
||||
|
||||
def clean_dashes(text: str) -> str:
|
||||
"""Cleans dash characters in text.
|
||||
|
||||
Example
|
||||
-------
|
||||
ITEM 1. -BUSINESS -> ITEM 1. BUSINESS
|
||||
"""
|
||||
# NOTE(Yuming): '\u2013' is the unicode string of 'EN DASH', a variation of "-"
|
||||
return re.sub(r"[-\u2013]", " ", text).strip()
|
||||
|
||||
|
||||
def clean_trailing_punctuation(text: str) -> str:
|
||||
"""Clean all trailing punctuation in text
|
||||
|
||||
Example
|
||||
-------
|
||||
ITEM 1. BUSINESS. -> ITEM 1. BUSINESS
|
||||
"""
|
||||
return text.strip().rstrip(".,:;")
|
||||
|
||||
|
||||
def replace_mime_encodings(text: str, encoding: str = "utf-8") -> str:
|
||||
"""Replaces MIME encodings with their equivalent characters in the specified encoding.
|
||||
|
||||
Example
|
||||
-------
|
||||
5 w=E2=80-99s -> 5 w’s
|
||||
"""
|
||||
formatted_encoding = format_encoding_str(encoding)
|
||||
return quopri.decodestring(text.encode(formatted_encoding)).decode(formatted_encoding)
|
||||
|
||||
|
||||
def clean_prefix(text: str, pattern: str, ignore_case: bool = False, strip: bool = True) -> str:
|
||||
"""Removes prefixes from a string according to the specified pattern. Strips leading
|
||||
whitespace if the strip parameter is set to True.
|
||||
|
||||
Input
|
||||
-----
|
||||
text: The text to clean
|
||||
pattern: The pattern for the prefix. Can be a simple string or a regex pattern
|
||||
ignore_case: If True, ignores case in the pattern
|
||||
strip: If True, removes leading whitespace from the cleaned string.
|
||||
"""
|
||||
flags = re.IGNORECASE if ignore_case else 0
|
||||
clean_text = re.sub(rf"^{pattern}", "", text, flags=flags)
|
||||
clean_text = clean_text.lstrip() if strip else clean_text
|
||||
return clean_text
|
||||
|
||||
|
||||
def clean_postfix(text: str, pattern: str, ignore_case: bool = False, strip: bool = True) -> str:
|
||||
"""Removes postfixes from a string according to the specified pattern. Strips trailing
|
||||
whitespace if the strip parameters is set to True.
|
||||
|
||||
Input
|
||||
-----
|
||||
text: The text to clean
|
||||
pattern: The pattern for the postfix. Can be a simple string or a regex pattern
|
||||
ignore_case: If True, ignores case in the pattern
|
||||
strip: If True, removes trailing whitespace from the cleaned string.
|
||||
"""
|
||||
flags = re.IGNORECASE if ignore_case else 0
|
||||
clean_text = re.sub(rf"{pattern}$", "", text, flags=flags)
|
||||
clean_text = clean_text.rstrip() if strip else clean_text
|
||||
return clean_text
|
||||
|
||||
|
||||
def clean(
|
||||
text: str,
|
||||
extra_whitespace: bool = False,
|
||||
dashes: bool = False,
|
||||
bullets: bool = False,
|
||||
trailing_punctuation: bool = False,
|
||||
lowercase: bool = False,
|
||||
) -> str:
|
||||
"""Cleans text.
|
||||
|
||||
Input
|
||||
-----
|
||||
extra_whitespace: Whether to clean extra whitespace characters in text.
|
||||
dashes: Whether to clean dash characters in text.
|
||||
bullets: Whether to clean unicode bullets from a section of text.
|
||||
trailing_punctuation: Whether to clean all trailing punctuation in text.
|
||||
lowercase: Whether to return lowercase text.
|
||||
"""
|
||||
|
||||
cleaned_text = text.lower() if lowercase else text
|
||||
cleaned_text = (
|
||||
clean_trailing_punctuation(cleaned_text) if trailing_punctuation else cleaned_text
|
||||
)
|
||||
cleaned_text = clean_dashes(cleaned_text) if dashes else cleaned_text
|
||||
cleaned_text = clean_extra_whitespace(cleaned_text) if extra_whitespace else cleaned_text
|
||||
cleaned_text = clean_bullets(cleaned_text) if bullets else cleaned_text
|
||||
return cleaned_text.strip()
|
||||
|
||||
|
||||
def bytes_string_to_string(text: str, encoding: str = "utf-8"):
|
||||
"""Converts a string representation of a byte string to a regular string using the
|
||||
specified encoding."""
|
||||
text_bytes = bytes([ord(char) for char in text])
|
||||
formatted_encoding = format_encoding_str(encoding)
|
||||
return text_bytes.decode(formatted_encoding)
|
||||
|
||||
|
||||
def clean_extra_whitespace_with_index_run(text: str) -> Tuple[str, np.ndarray]:
|
||||
"""Cleans extra whitespace characters that appear between words.
|
||||
Calculate distance between characters of original text and cleaned text.
|
||||
|
||||
Returns cleaned text along with array of indices it has moved from original.
|
||||
|
||||
Example
|
||||
-------
|
||||
ITEM 1. BUSINESS -> ITEM 1. BUSINESS
|
||||
array([0., 0., 0., 0., 0., 0., 0., 0., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4., 4.]))
|
||||
"""
|
||||
|
||||
cleaned_text = re.sub(r"[\xa0\n]", " ", text)
|
||||
cleaned_text = re.sub(r"([ ]{2,})", " ", cleaned_text)
|
||||
|
||||
cleaned_text = cleaned_text.strip()
|
||||
|
||||
moved_indices = np.zeros(len(text))
|
||||
|
||||
distance, original_index, cleaned_index = 0, 0, 0
|
||||
while cleaned_index < len(cleaned_text):
|
||||
if text[original_index] == cleaned_text[cleaned_index] or (
|
||||
bool(re.match("[\xa0\n]", text[original_index]))
|
||||
and bool(re.match(" ", cleaned_text[cleaned_index]))
|
||||
):
|
||||
moved_indices[cleaned_index] = distance
|
||||
original_index += 1
|
||||
cleaned_index += 1
|
||||
continue
|
||||
|
||||
distance += 1
|
||||
moved_indices[cleaned_index] = distance
|
||||
original_index += 1
|
||||
|
||||
moved_indices[cleaned_index:] = distance
|
||||
|
||||
return cleaned_text, moved_indices
|
||||
|
||||
|
||||
def index_adjustment_after_clean_extra_whitespace(index, moved_indices) -> int:
|
||||
return int(index - moved_indices[index])
|
||||
@@ -0,0 +1,143 @@
|
||||
import datetime
|
||||
import re
|
||||
from typing import List, Optional
|
||||
|
||||
from unstructured.nlp.patterns import (
|
||||
EMAIL_ADDRESS_PATTERN,
|
||||
EMAIL_DATETIMETZ_PATTERN,
|
||||
IMAGE_URL_PATTERN,
|
||||
IP_ADDRESS_NAME_PATTERN,
|
||||
IP_ADDRESS_PATTERN_RE,
|
||||
MAPI_ID_PATTERN,
|
||||
US_PHONE_NUMBERS_RE,
|
||||
)
|
||||
|
||||
|
||||
def _get_indexed_match(text: str, pattern: str, index: int = 0) -> re.Match:
|
||||
if not isinstance(index, int) or index < 0:
|
||||
raise ValueError(f"The index is {index}. Index must be a non-negative integer.")
|
||||
|
||||
regex_match = None
|
||||
for i, result in enumerate(re.finditer(pattern, text)):
|
||||
if i == index:
|
||||
regex_match = result
|
||||
|
||||
if regex_match is None:
|
||||
raise ValueError(f"Result with index {index} was not found. The largest index was {i}.")
|
||||
|
||||
return regex_match
|
||||
|
||||
|
||||
def extract_text_before(text: str, pattern: str, index: int = 0, strip: bool = True) -> str:
|
||||
"""Extracts texts that occurs before the specified pattern. By default, it will use
|
||||
the first occurrence of the pattern (index 0). Use the index kwarg to choose a different
|
||||
index.
|
||||
|
||||
Input
|
||||
-----
|
||||
strip: If True, removes trailing whitespace from the extracted string
|
||||
"""
|
||||
regex_match = _get_indexed_match(text, pattern, index)
|
||||
start, _ = regex_match.span()
|
||||
before_text = text[:start]
|
||||
return before_text.rstrip() if strip else before_text
|
||||
|
||||
|
||||
def extract_text_after(text: str, pattern: str, index: int = 0, strip: bool = True) -> str:
|
||||
"""Extracts texts that occurs before the specified pattern. By default, it will use
|
||||
the first occurrence of the pattern (index 0). Use the index kwarg to choose a different
|
||||
index.
|
||||
|
||||
Input
|
||||
-----
|
||||
strip: If True, removes leading whitespace from the extracted string
|
||||
"""
|
||||
regex_match = _get_indexed_match(text, pattern, index)
|
||||
_, end = regex_match.span()
|
||||
before_text = text[end:]
|
||||
return before_text.lstrip() if strip else before_text
|
||||
|
||||
|
||||
def extract_email_address(text: str) -> List[str]:
|
||||
return re.findall(EMAIL_ADDRESS_PATTERN, text.lower())
|
||||
|
||||
|
||||
def extract_ip_address(text: str) -> List[str]:
|
||||
return re.findall(IP_ADDRESS_PATTERN_RE, text)
|
||||
|
||||
|
||||
def extract_ip_address_name(text: str) -> List[str]:
|
||||
return re.findall(IP_ADDRESS_NAME_PATTERN, text)
|
||||
|
||||
|
||||
def extract_mapi_id(text: str) -> List[str]:
|
||||
mapi_ids = re.findall(MAPI_ID_PATTERN, text)
|
||||
mapi_ids = [mid.replace(";", "") for mid in mapi_ids]
|
||||
return mapi_ids
|
||||
|
||||
|
||||
def extract_datetimetz(text: str) -> Optional[datetime.datetime]:
|
||||
date_extractions = re.findall(EMAIL_DATETIMETZ_PATTERN, text)
|
||||
if len(date_extractions) > 0:
|
||||
return datetime.datetime.strptime(date_extractions[0], "%a, %d %b %Y %H:%M:%S %z")
|
||||
else:
|
||||
return None
|
||||
|
||||
|
||||
def extract_us_phone_number(text: str):
|
||||
"""Extracts a US phone number from a section of text that includes a phone number. If there
|
||||
is no phone number present, the result will be an empty string.
|
||||
|
||||
Example
|
||||
-------
|
||||
extract_phone_number("Phone Number: 215-867-5309") -> "215-867-5309"
|
||||
"""
|
||||
regex_match = US_PHONE_NUMBERS_RE.search(text)
|
||||
if regex_match is None:
|
||||
return ""
|
||||
|
||||
start, end = regex_match.span()
|
||||
phone_number = text[start:end]
|
||||
return phone_number.strip()
|
||||
|
||||
|
||||
def extract_ordered_bullets(text) -> tuple:
|
||||
"""Extracts the start of bulleted text sections bullets
|
||||
accounting numeric and alphanumeric types.
|
||||
|
||||
Output
|
||||
-----
|
||||
tuple(section, sub_section, sub_sub_section): Each bullet partition
|
||||
is a string or None if not present.
|
||||
|
||||
Example
|
||||
-------
|
||||
This is a very important point -> (None, None, None)
|
||||
1.1 This is a very important point -> ("1", "1", None)
|
||||
a.1 This is a very important point -> ("a", "1", None)
|
||||
"""
|
||||
a, b, c, temp = None, None, None, None
|
||||
text_sp = text.split()
|
||||
if any(["." not in text_sp[0], ".." in text_sp[0]]):
|
||||
return a, b, c
|
||||
|
||||
bullet = re.split(pattern=r"[\.]", string=text_sp[0])
|
||||
if not bullet[-1]:
|
||||
del bullet[-1]
|
||||
|
||||
if len(bullet[0]) > 2:
|
||||
return a, b, c
|
||||
|
||||
a, *temp = bullet
|
||||
if temp:
|
||||
try:
|
||||
b, c, *_ = temp
|
||||
except ValueError:
|
||||
b = temp
|
||||
b = "".join(b)
|
||||
c = "".join(c) if c else None
|
||||
return a, b, c
|
||||
|
||||
|
||||
def extract_image_urls_from_html(text: str) -> List[str]:
|
||||
return re.findall(IMAGE_URL_PATTERN, text)
|
||||
@@ -0,0 +1,87 @@
|
||||
import warnings
|
||||
from typing import List, Optional
|
||||
|
||||
import langdetect
|
||||
from transformers import MarianMTModel, MarianTokenizer
|
||||
|
||||
from unstructured.nlp.tokenize import sent_tokenize
|
||||
from unstructured.staging.huggingface import chunk_by_attention_window
|
||||
|
||||
|
||||
def _get_opus_mt_model_name(source_lang: str, target_lang: str):
|
||||
"""Constructs the name of the MarianMT machine translation model based on the
|
||||
source and target language."""
|
||||
return f"Helsinki-NLP/opus-mt-{source_lang}-{target_lang}"
|
||||
|
||||
|
||||
def _validate_language_code(language_code: str):
|
||||
if not isinstance(language_code, str) or len(language_code) != 2:
|
||||
raise ValueError(
|
||||
f"Invalid language code: {language_code}. Language codes must be two letter strings.",
|
||||
)
|
||||
|
||||
|
||||
def translate_text(text: str, source_lang: Optional[str] = None, target_lang: str = "en") -> str:
|
||||
"""Translates the foreign language text. If the source language is not specified, the
|
||||
function will attempt to detect it using langdetect.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
text: str
|
||||
The text to translate
|
||||
target_lang: str
|
||||
The two letter language code for the target langague. Defaults to "en".
|
||||
source_lang: Optional[str]
|
||||
The two letter language code for the language of the input text. If source_lang is
|
||||
not provided, the function will try to detect it.
|
||||
"""
|
||||
if text.strip() == "":
|
||||
return text
|
||||
|
||||
_source_lang: str = source_lang if source_lang is not None else langdetect.detect(text)
|
||||
# NOTE(robinson) - Chinese gets detected with codes zh-cn, zh-tw, zh-hk for various
|
||||
# Chinese variants. We normalizes these because there is a single model for Chinese
|
||||
# machine translation
|
||||
if _source_lang.startswith("zh"):
|
||||
_source_lang = "zh"
|
||||
|
||||
_validate_language_code(target_lang)
|
||||
_validate_language_code(_source_lang)
|
||||
|
||||
if target_lang == _source_lang:
|
||||
return text
|
||||
|
||||
model_name = _get_opus_mt_model_name(_source_lang, target_lang)
|
||||
print(f"Using model: {model_name}")
|
||||
|
||||
try:
|
||||
tokenizer = MarianTokenizer.from_pretrained(model_name)
|
||||
model = MarianMTModel.from_pretrained(model_name)
|
||||
except OSError:
|
||||
raise ValueError(
|
||||
f"Transformers could not find the translation model {model_name}. "
|
||||
"The requested source/target language combo is not supported.",
|
||||
)
|
||||
|
||||
chunks: List[str] = chunk_by_attention_window(text, tokenizer, split_function=sent_tokenize)
|
||||
|
||||
translated_chunks: List[str] = []
|
||||
for chunk in chunks:
|
||||
translated_chunks.append(_translate_text(text, model, tokenizer))
|
||||
|
||||
return " ".join(translated_chunks)
|
||||
|
||||
|
||||
def _translate_text(text, model, tokenizer):
|
||||
"""Translates text using the specified model and tokenizer."""
|
||||
# NOTE(robinson) - Suppresses the HuggingFace UserWarning resulting from the "max_length"
|
||||
# key in the MarianMT config. The warning states that "max_length" will be deprecated
|
||||
# in transformers v5
|
||||
with warnings.catch_warnings():
|
||||
warnings.simplefilter("ignore")
|
||||
translated = model.generate(
|
||||
**tokenizer([text], return_tensors="pt", padding=True, truncation=True),
|
||||
)
|
||||
return [tokenizer.decode(t, max_new_tokens=512, skip_special_tokens=True) for t in translated][
|
||||
0
|
||||
]
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,165 @@
|
||||
"""Provides operations related to the HTML table stored in `.metadata.text_as_html`.
|
||||
|
||||
Used during partitioning as well as chunking.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import html
|
||||
from typing import TYPE_CHECKING, Iterator, Sequence, cast
|
||||
|
||||
from lxml import etree
|
||||
from lxml.html import fragment_fromstring
|
||||
|
||||
from unstructured.utils import lazyproperty
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from lxml.html import HtmlElement
|
||||
|
||||
|
||||
def htmlify_matrix_of_cell_texts(matrix: Sequence[Sequence[str]]) -> str:
|
||||
"""Form an HTML table from "rows" and "columns" of `matrix`.
|
||||
|
||||
Character overhead is minimized:
|
||||
- No whitespace padding is added for human readability
|
||||
- No newlines ("\n") are added
|
||||
- No `<thead>`, `<tbody>`, or `<tfoot>` elements are used; we can't tell where those might be
|
||||
semantically appropriate anyway so at best they would consume unnecessary space and at worst
|
||||
would be misleading.
|
||||
"""
|
||||
|
||||
def iter_trs(rows_of_cell_strs: Sequence[Sequence[str]]) -> Iterator[str]:
|
||||
for row_cell_strs in rows_of_cell_strs:
|
||||
# -- suppress emission of rows with no cells --
|
||||
if not row_cell_strs:
|
||||
continue
|
||||
yield f"<tr>{''.join(iter_tds(row_cell_strs))}</tr>"
|
||||
|
||||
def iter_tds(row_cell_strs: Sequence[str]) -> Iterator[str]:
|
||||
for s in row_cell_strs:
|
||||
# -- take care of things like '<' and '>' in the text --
|
||||
s = html.escape(s)
|
||||
# -- substitute <br/> elements for line-feeds in the text --
|
||||
s = "<br/>".join(s.split("\n"))
|
||||
# -- normalize whitespace in cell --
|
||||
cell_text = " ".join(s.split())
|
||||
# -- emit void `<td/>` when cell text is empty string --
|
||||
yield f"<td>{cell_text}</td>" if cell_text else "<td/>"
|
||||
|
||||
return f"<table>{''.join(iter_trs(matrix))}</table>" if matrix else ""
|
||||
|
||||
|
||||
class HtmlTable:
|
||||
"""A `<table>` element."""
|
||||
|
||||
def __init__(self, table: HtmlElement):
|
||||
self._table = table
|
||||
|
||||
@classmethod
|
||||
def from_html_text(cls, html_text: str) -> HtmlTable:
|
||||
# -- root is always a `<table>` element so far but let's be robust --
|
||||
root = fragment_fromstring(html_text)
|
||||
tables = root.xpath("//table")
|
||||
if not tables:
|
||||
raise ValueError("`html_text` contains no `<table>` element")
|
||||
table = tables[0]
|
||||
|
||||
# -- remove `<thead>`, `<tbody>`, and `<tfoot>` noise elements when present --
|
||||
noise_elements = table.xpath(".//thead | .//tbody | .//tfoot")
|
||||
for e in noise_elements:
|
||||
e.drop_tag()
|
||||
|
||||
# -- normalize and compactify the HTML --
|
||||
for e in table.iter():
|
||||
# -- Strip all attributes from elements, like border="1", class="dataframe" added
|
||||
# -- by pandas.DataFrame.to_html(), style="text-align: right;", etc.
|
||||
e.attrib.clear()
|
||||
|
||||
# -- change any `<th>` elements to `<td>` so all cells have the same tag --
|
||||
if e.tag == "th":
|
||||
e.tag = "td"
|
||||
|
||||
# -- normalize whitespace in element text; this removes indent whitespace before nested
|
||||
# -- elements and reduces whitespace between words to a single space.
|
||||
if e.text:
|
||||
e.text = " ".join(e.text.split())
|
||||
|
||||
# -- remove all tails, those are newline + indent if anything --
|
||||
if e.tail:
|
||||
e.tail = None
|
||||
|
||||
return cls(table)
|
||||
|
||||
@lazyproperty
|
||||
def html(self) -> str:
|
||||
"""The HTML-fragment for this `<table>` element, all on one line.
|
||||
|
||||
Like: `<table><tr><td>foo</td></tr><tr><td>bar</td></tr></table>`
|
||||
|
||||
The HTML contains no human-readability whitespace, attributes, or `<thead>`, `<tbody>`, or
|
||||
`<tfoot>` tags. It is made as compact as possible to maximize the semantic content in a
|
||||
given space. This is particularly important for chunking.
|
||||
"""
|
||||
return etree.tostring(self._table, encoding=str)
|
||||
|
||||
def iter_rows(self) -> Iterator[HtmlRow]:
|
||||
yield from (HtmlRow(tr) for tr in cast("list[HtmlElement]", self._table.xpath("./tr")))
|
||||
|
||||
@lazyproperty
|
||||
def text(self) -> str:
|
||||
"""The clean, concatenated, text for this table."""
|
||||
table_text = " ".join(self._table.itertext())
|
||||
# -- blank cells will introduce extra whitespace, so normalize after accumulating --
|
||||
return " ".join(table_text.split())
|
||||
|
||||
|
||||
class HtmlRow:
|
||||
"""A `<tr>` element."""
|
||||
|
||||
def __init__(self, tr: HtmlElement):
|
||||
self._tr = tr
|
||||
|
||||
@lazyproperty
|
||||
def html(self) -> str:
|
||||
"""Like "<tr><td>foo</td><td>bar</td></tr>"."""
|
||||
return etree.tostring(self._tr, encoding=str)
|
||||
|
||||
def iter_cells(self) -> Iterator[HtmlCell]:
|
||||
for td in self._tr:
|
||||
yield HtmlCell(td)
|
||||
|
||||
def iter_cell_texts(self) -> Iterator[str]:
|
||||
"""Generate contents of each cell of this row as a separate string.
|
||||
|
||||
A cell that is empty or contains only whitespace does not generate a string.
|
||||
"""
|
||||
for td in self._tr:
|
||||
if (text := td.text) is None:
|
||||
continue
|
||||
if not text:
|
||||
continue
|
||||
yield text
|
||||
|
||||
@lazyproperty
|
||||
def text_len(self) -> int:
|
||||
"""Length of the normalized text, as it would appear in `element.text`."""
|
||||
return len(" ".join(self.iter_cell_texts()))
|
||||
|
||||
|
||||
class HtmlCell:
|
||||
"""A `<td>` element."""
|
||||
|
||||
def __init__(self, td: HtmlElement):
|
||||
self._td = td
|
||||
|
||||
@lazyproperty
|
||||
def html(self) -> str:
|
||||
"""Like "<td>foo bar baz</td>"."""
|
||||
return etree.tostring(self._td, encoding=str) if self.text else "<td/>"
|
||||
|
||||
@lazyproperty
|
||||
def text(self) -> str:
|
||||
"""Text inside `<td>` element, empty string when no text."""
|
||||
if (text := self._td.text) is None:
|
||||
return ""
|
||||
return " ".join(text.strip().split())
|
||||
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
|
||||
|
||||
from enum import Enum
|
||||
from typing import Any, Dict, Sequence, Tuple, Union
|
||||
|
||||
|
||||
class Orientation(Enum):
|
||||
SCREEN = (1, -1) # Origin in top left, y increases in the down direction
|
||||
CARTESIAN = (1, 1) # Origin in bottom left, y increases in upward direction
|
||||
|
||||
|
||||
def convert_coordinate(old_t, old_t_max, new_t_max, t_orientation):
|
||||
"""Convert a coordinate into another system along an axis using a linear transformation"""
|
||||
return (
|
||||
(1 - old_t / old_t_max) * (1 - t_orientation) / 2
|
||||
+ old_t / old_t_max * (1 + t_orientation) / 2
|
||||
) * new_t_max
|
||||
|
||||
|
||||
class CoordinateSystem:
|
||||
"""A finite coordinate plane with given width and height."""
|
||||
|
||||
orientation: Orientation
|
||||
|
||||
def __init__(self, width: Union[int, float], height: Union[int, float]):
|
||||
self.width = width
|
||||
self.height = height
|
||||
|
||||
def __eq__(self, other: object):
|
||||
if not isinstance(other, CoordinateSystem):
|
||||
return False
|
||||
return (
|
||||
str(self.__class__.__name__) == str(other.__class__.__name__)
|
||||
and self.width == other.width
|
||||
and self.height == other.height
|
||||
and self.orientation == other.orientation
|
||||
)
|
||||
|
||||
def convert_from_relative(
|
||||
self,
|
||||
x: Union[float, int],
|
||||
y: Union[float, int],
|
||||
) -> Tuple[Union[float, int], Union[float, int]]:
|
||||
"""Convert to this coordinate system from a relative coordinate system."""
|
||||
x_orientation, y_orientation = self.orientation.value
|
||||
new_x = convert_coordinate(x, 1, self.width, x_orientation)
|
||||
new_y = convert_coordinate(y, 1, self.height, y_orientation)
|
||||
return new_x, new_y
|
||||
|
||||
def convert_to_relative(
|
||||
self,
|
||||
x: Union[float, int],
|
||||
y: Union[float, int],
|
||||
) -> Tuple[Union[float, int], Union[float, int]]:
|
||||
"""Convert from this coordinate system to a relative coordinate system."""
|
||||
x_orientation, y_orientation = self.orientation.value
|
||||
new_x = convert_coordinate(x, self.width, 1, x_orientation)
|
||||
new_y = convert_coordinate(y, self.height, 1, y_orientation)
|
||||
return new_x, new_y
|
||||
|
||||
def convert_coordinates_to_new_system(
|
||||
self,
|
||||
new_system: CoordinateSystem,
|
||||
x: Union[float, int],
|
||||
y: Union[float, int],
|
||||
) -> Tuple[Union[float, int], Union[float, int]]:
|
||||
"""Convert from this coordinate system to another given coordinate system."""
|
||||
rel_x, rel_y = self.convert_to_relative(x, y)
|
||||
return new_system.convert_from_relative(rel_x, rel_y)
|
||||
|
||||
def convert_multiple_coordinates_to_new_system(
|
||||
self,
|
||||
new_system: CoordinateSystem,
|
||||
coordinates: Sequence[Tuple[Union[float, int], Union[float, int]]],
|
||||
) -> Tuple[Tuple[Union[float, int], Union[float, int]], ...]:
|
||||
"""Convert (x, y) coordinates from current system to another coordinate system."""
|
||||
new_system_coordinates = []
|
||||
for x, y in coordinates:
|
||||
new_system_coordinates.append(
|
||||
self.convert_coordinates_to_new_system(new_system=new_system, x=x, y=y),
|
||||
)
|
||||
return tuple(new_system_coordinates)
|
||||
|
||||
|
||||
class RelativeCoordinateSystem(CoordinateSystem):
|
||||
"""Relative coordinate system where x and y are on a scale from 0 to 1."""
|
||||
|
||||
orientation = Orientation.CARTESIAN
|
||||
|
||||
def __init__(self):
|
||||
self.width = 1
|
||||
self.height = 1
|
||||
|
||||
|
||||
class PixelSpace(CoordinateSystem):
|
||||
"""Coordinate system representing a pixel space, such as an image. The origin is at the top
|
||||
left."""
|
||||
|
||||
orientation = Orientation.SCREEN
|
||||
|
||||
|
||||
class PointSpace(CoordinateSystem):
|
||||
"""Coordinate system representing a point space, such as a pdf. The origin is at the bottom
|
||||
left."""
|
||||
|
||||
orientation = Orientation.CARTESIAN
|
||||
|
||||
|
||||
TYPE_TO_COORDINATE_SYSTEM_MAP: Dict[str, Any] = {
|
||||
"PixelSpace": PixelSpace,
|
||||
"PointSpace": PointSpace,
|
||||
"CoordinateSystem": CoordinateSystem,
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,195 @@
|
||||
"""
|
||||
This module contains mapping between:
|
||||
HTML Tags <-> Elements Ontology <-> Unstructured Element classes
|
||||
They are used to simplify transformations between different representations
|
||||
of parsed documents
|
||||
"""
|
||||
|
||||
from typing import Dict, Type
|
||||
|
||||
from unstructured.documents import elements, ontology
|
||||
from unstructured.documents.elements import Element
|
||||
|
||||
|
||||
def get_all_subclasses(cls: type) -> list[type]:
|
||||
"""
|
||||
Recursively find all subclasses of a given class.
|
||||
|
||||
Parameters:
|
||||
cls (type): The class for which to find all subclasses.
|
||||
|
||||
Returns:
|
||||
list[type]: A list of all subclasses of the given class.
|
||||
"""
|
||||
subclasses = cls.__subclasses__()
|
||||
all_subclasses = subclasses.copy()
|
||||
|
||||
for subclass in subclasses:
|
||||
all_subclasses.extend(get_all_subclasses(subclass))
|
||||
|
||||
return all_subclasses
|
||||
|
||||
|
||||
def get_ontology_to_unstructured_type_mapping() -> (
|
||||
dict[Type[ontology.OntologyElement], Type[Element]]
|
||||
):
|
||||
"""
|
||||
Get a mapping of ontology element to unstructured type.
|
||||
|
||||
The dictionary here was created base on ontology mapping json
|
||||
Can be generated via the following code:
|
||||
```
|
||||
ontology_elements_list = json.loads(
|
||||
Path("unstructured_element_ontology.json").read_text()
|
||||
)
|
||||
ontology_to_unstructured_class_mapping = {
|
||||
ontology_element["name"]: ontology_element["ontologyV1Mapping"]
|
||||
for ontology_element in ontology_elements_list
|
||||
}
|
||||
```
|
||||
|
||||
Returns:
|
||||
dict: A dictionary where keys are ontology element classes
|
||||
and values are unstructured types.
|
||||
"""
|
||||
ontology_to_unstructured_class_mapping: Dict[Type[ontology.OntologyElement], Type[Element]] = {
|
||||
ontology.Document: elements.Text,
|
||||
ontology.Section: elements.Text,
|
||||
ontology.Page: elements.Text,
|
||||
ontology.Column: elements.Text,
|
||||
ontology.Paragraph: elements.NarrativeText,
|
||||
ontology.Header: elements.Header,
|
||||
ontology.Footer: elements.Footer,
|
||||
ontology.Sidebar: elements.Text,
|
||||
ontology.PageBreak: elements.PageBreak,
|
||||
ontology.Title: elements.Title,
|
||||
ontology.Subtitle: elements.Title,
|
||||
ontology.Heading: elements.Title,
|
||||
ontology.NarrativeText: elements.NarrativeText,
|
||||
ontology.Quote: elements.NarrativeText,
|
||||
ontology.Footnote: elements.Text,
|
||||
ontology.Caption: elements.FigureCaption,
|
||||
ontology.PageNumber: elements.PageNumber,
|
||||
ontology.UncategorizedText: elements.Text,
|
||||
ontology.OrderedList: elements.Text,
|
||||
ontology.UnorderedList: elements.Text,
|
||||
ontology.DefinitionList: elements.Text,
|
||||
ontology.ListItem: elements.ListItem,
|
||||
ontology.Table: elements.Table,
|
||||
ontology.TableRow: elements.Table,
|
||||
ontology.TableCell: elements.Table,
|
||||
ontology.TableCellHeader: elements.Table,
|
||||
ontology.TableBody: elements.Table,
|
||||
ontology.TableHeader: elements.Table,
|
||||
ontology.Image: elements.Image,
|
||||
ontology.Figure: elements.Image,
|
||||
ontology.Video: elements.Text,
|
||||
ontology.Audio: elements.Text,
|
||||
ontology.Barcode: elements.Image,
|
||||
ontology.QRCode: elements.Image,
|
||||
ontology.Logo: elements.Image,
|
||||
ontology.CodeBlock: elements.CodeSnippet,
|
||||
ontology.InlineCode: elements.CodeSnippet,
|
||||
ontology.Formula: elements.Formula,
|
||||
ontology.Equation: elements.Formula,
|
||||
ontology.FootnoteReference: elements.Text,
|
||||
ontology.Citation: elements.Text,
|
||||
ontology.Bibliography: elements.Text,
|
||||
ontology.Glossary: elements.Text,
|
||||
ontology.Author: elements.Text,
|
||||
ontology.MetaDate: elements.Text,
|
||||
ontology.Keywords: elements.Text,
|
||||
ontology.Abstract: elements.NarrativeText,
|
||||
ontology.Hyperlink: elements.Text,
|
||||
ontology.TableOfContents: elements.Table,
|
||||
ontology.Index: elements.Text,
|
||||
ontology.Form: elements.Text,
|
||||
ontology.FormField: elements.Text,
|
||||
ontology.FormFieldValue: elements.Text,
|
||||
ontology.Checkbox: elements.Text,
|
||||
ontology.RadioButton: elements.Text,
|
||||
ontology.Button: elements.Text,
|
||||
ontology.Comment: elements.Text,
|
||||
ontology.Highlight: elements.Text,
|
||||
ontology.RevisionInsertion: elements.Text,
|
||||
ontology.RevisionDeletion: elements.Text,
|
||||
ontology.Address: elements.Address,
|
||||
ontology.EmailAddress: elements.EmailAddress,
|
||||
ontology.PhoneNumber: elements.Text,
|
||||
ontology.CalendarDate: elements.Text,
|
||||
ontology.Time: elements.Text,
|
||||
ontology.Currency: elements.Text,
|
||||
ontology.Measurement: elements.Text,
|
||||
ontology.Letterhead: elements.Header,
|
||||
ontology.Signature: elements.Text,
|
||||
ontology.Watermark: elements.Text,
|
||||
ontology.Stamp: elements.Text,
|
||||
}
|
||||
|
||||
return ontology_to_unstructured_class_mapping
|
||||
|
||||
|
||||
ALL_ONTOLOGY_ELEMENT_TYPES = get_all_subclasses(ontology.OntologyElement)
|
||||
HTML_TAG_AND_CSS_NAME_TO_ELEMENT_TYPE_MAP: Dict[tuple[str, str], Type[ontology.OntologyElement]] = {
|
||||
(tag, element_type().css_class_name): element_type
|
||||
for element_type in ALL_ONTOLOGY_ELEMENT_TYPES
|
||||
for tag in element_type().allowed_tags
|
||||
}
|
||||
CSS_CLASS_TO_ELEMENT_TYPE_MAP: Dict[str, Type[ontology.OntologyElement]] = {
|
||||
element_type().css_class_name: element_type for element_type in ALL_ONTOLOGY_ELEMENT_TYPES
|
||||
}
|
||||
|
||||
HTML_TAG_TO_DEFAULT_ELEMENT_TYPE_MAP: Dict[str, Type[ontology.OntologyElement]] = {
|
||||
"a": ontology.Hyperlink,
|
||||
"address": ontology.Address,
|
||||
"aside": ontology.Sidebar,
|
||||
"audio": ontology.Audio,
|
||||
"blockquote": ontology.Quote,
|
||||
"body": ontology.Document,
|
||||
"button": ontology.Button,
|
||||
"cite": ontology.Citation,
|
||||
"code": ontology.CodeBlock,
|
||||
"del": ontology.RevisionDeletion,
|
||||
"div": ontology.UncategorizedText,
|
||||
"dl": ontology.DefinitionList,
|
||||
"figcaption": ontology.Caption,
|
||||
"figure": ontology.Figure,
|
||||
"footer": ontology.Footer,
|
||||
"form": ontology.Form,
|
||||
"h1": ontology.Title,
|
||||
"h2": ontology.Subtitle,
|
||||
"h3": ontology.Heading,
|
||||
"h4": ontology.Heading,
|
||||
"h5": ontology.Heading,
|
||||
"h6": ontology.Heading,
|
||||
"header": ontology.Header,
|
||||
"hr": ontology.PageBreak,
|
||||
"img": ontology.Image,
|
||||
"input": ontology.Checkbox,
|
||||
"ins": ontology.RevisionInsertion,
|
||||
"label": ontology.FormField,
|
||||
"li": ontology.ListItem,
|
||||
"mark": ontology.Highlight,
|
||||
"math": ontology.Equation,
|
||||
"meta": ontology.Keywords,
|
||||
"nav": ontology.Index,
|
||||
"ol": ontology.OrderedList,
|
||||
"p": ontology.Paragraph,
|
||||
"pre": ontology.CodeBlock,
|
||||
"section": ontology.Section,
|
||||
"span": ontology.UncategorizedText,
|
||||
"sub": ontology.FootnoteReference,
|
||||
"svg": ontology.Signature,
|
||||
"table": ontology.Table,
|
||||
"tbody": ontology.TableBody,
|
||||
"td": ontology.TableCell,
|
||||
"th": ontology.TableCellHeader,
|
||||
"thead": ontology.TableHeader,
|
||||
"time": ontology.Time,
|
||||
"tr": ontology.TableRow,
|
||||
"ul": ontology.UnorderedList,
|
||||
"video": ontology.Video,
|
||||
}
|
||||
|
||||
|
||||
ONTOLOGY_CLASS_TO_UNSTRUCTURED_ELEMENT_TYPE = get_ontology_to_unstructured_type_mapping()
|
||||
@@ -0,0 +1,622 @@
|
||||
"""
|
||||
This file contains all classes allowed in the ontology V2.
|
||||
This Type is used as intermediate representation between HTML
|
||||
and Unstructured Elements.
|
||||
All the processing could be done without the intermediate representation,
|
||||
but it simplifies the process.
|
||||
It needs to be decide whether we keep it or not.
|
||||
|
||||
The classes are represented as pydantic models to mimic Unstructured Elements V1 solutions.
|
||||
However it results in lots of code that could be strongly simplified.
|
||||
|
||||
TODO (Pluto): OntologyElement is the only needed class. It could contains data about
|
||||
allowed html tags, css classes and descriptions as metadata.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from copy import copy
|
||||
from enum import Enum
|
||||
from typing import Any, List, Optional
|
||||
|
||||
from bs4 import BeautifulSoup
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class ElementTypeEnum(str, Enum):
|
||||
layout = "Layout"
|
||||
text = "Text"
|
||||
list = "List"
|
||||
table = "Table"
|
||||
media = "Media"
|
||||
code = "Code"
|
||||
mathematical = "Mathematical"
|
||||
reference = "Reference"
|
||||
metadata = "Metadata"
|
||||
navigation = "Navigation"
|
||||
form = "Form"
|
||||
annotation = "Annotation"
|
||||
specialized_text = "Specialized Text"
|
||||
document_specific = "Document-Specific"
|
||||
|
||||
|
||||
class OntologyElement(BaseModel):
|
||||
text: Optional[str] = Field("", description="Text content of the element")
|
||||
css_class_name: Optional[str] = Field(
|
||||
default_factory=lambda: "", description="CSS class associated with the element"
|
||||
)
|
||||
html_tag_name: Optional[str] = Field(
|
||||
default_factory=lambda: "", description="HTML Tag name associated with the element"
|
||||
)
|
||||
elementType: ElementTypeEnum = Field(description="Type of the element")
|
||||
children: list["OntologyElement"] = Field(
|
||||
default_factory=list, description="List of child elements"
|
||||
)
|
||||
|
||||
description: str = Field(description="Description of the element")
|
||||
allowed_tags: list[str] = Field(description="HTML tags associated with the element")
|
||||
|
||||
additional_attributes: dict[str, Any] = Field(
|
||||
default_factory=dict, description="Optional HTML attributes or CSS properties"
|
||||
)
|
||||
|
||||
def __init__(self, **kwargs: dict[str, Any]):
|
||||
super().__init__(**kwargs)
|
||||
if self.css_class_name == "": # if None, then do not set
|
||||
self.css_class_name = self.__class__.__name__
|
||||
if self.html_tag_name == "":
|
||||
self.html_tag_name = self.allowed_tags[0]
|
||||
if "id" not in self.additional_attributes:
|
||||
self.additional_attributes["id"] = self.generate_unique_id()
|
||||
|
||||
@staticmethod
|
||||
def generate_unique_id() -> str:
|
||||
return str(uuid.uuid4()).replace("-", "")
|
||||
|
||||
def to_html(self, add_children: bool = True) -> str:
|
||||
additional_attrs = copy(self.additional_attributes)
|
||||
additional_attrs.pop("class", None)
|
||||
additional_attrs.pop("id", None)
|
||||
|
||||
attr_str = self._construct_attribute_string(additional_attrs)
|
||||
class_attr = f'class="{self.css_class_name}"' if self.css_class_name else ""
|
||||
|
||||
combined_attr_str = f"{class_attr} {attr_str}".strip()
|
||||
|
||||
children_html = self._generate_children_html(add_children)
|
||||
|
||||
result_html = self._generate_final_html(combined_attr_str, children_html)
|
||||
|
||||
return result_html
|
||||
|
||||
def to_text(self, add_children: bool = True, add_img_alt_text: bool = True) -> str:
|
||||
"""
|
||||
Returns the text representation of the element.
|
||||
|
||||
Args:
|
||||
add_children: If True, the text of the children will be included.
|
||||
Otherwise, element is represented as single self-closing tag.
|
||||
add_img_alt_text: If True, the alt text of the image will be included.
|
||||
"""
|
||||
if self.children and add_children:
|
||||
children_text = " ".join(
|
||||
child.to_text(add_children, add_img_alt_text).strip() for child in self.children
|
||||
)
|
||||
return children_text
|
||||
|
||||
text = BeautifulSoup(self.to_html(), "html.parser").get_text().strip()
|
||||
|
||||
if add_img_alt_text and self.html_tag_name == "img" and "alt" in self.additional_attributes:
|
||||
text += f" {self.additional_attributes.get('alt', '')}"
|
||||
|
||||
return text.strip()
|
||||
|
||||
def _construct_attribute_string(self, attributes: dict[str, str]) -> str:
|
||||
return " ".join(
|
||||
f'{key}="{value}"' if value else f"{key}" for key, value in attributes.items()
|
||||
)
|
||||
|
||||
def _generate_children_html(self, add_children: bool) -> str:
|
||||
if not add_children or not self.children:
|
||||
return ""
|
||||
return "".join(child.to_html() for child in self.children)
|
||||
|
||||
def _generate_final_html(self, attr_str: str, children_html: str) -> str:
|
||||
text = self.text or ""
|
||||
|
||||
if text or children_html:
|
||||
inside_tag_text = f"{text} {children_html}".strip()
|
||||
return f"<{self.html_tag_name} {attr_str}>{inside_tag_text}</{self.html_tag_name}>"
|
||||
else:
|
||||
return f"<{self.html_tag_name} {attr_str} />"
|
||||
|
||||
@property
|
||||
def id(self) -> str | None:
|
||||
return self.additional_attributes.get("id", None)
|
||||
|
||||
@property
|
||||
def page_number(self) -> int | None:
|
||||
if "data-page-number" in self.additional_attributes:
|
||||
try:
|
||||
page_attr = self.additional_attributes.get("data-page-number")
|
||||
if page_attr is not None:
|
||||
return int(page_attr)
|
||||
except ValueError:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def remove_ids_and_class_from_table(
|
||||
soup: BeautifulSoup, class_attr_to_keep: list[str] = ["img", "input"]
|
||||
) -> BeautifulSoup:
|
||||
"""
|
||||
Remove id and class attributes from tags inside tables,
|
||||
except preserve class attributes for selected tags.
|
||||
|
||||
Args:
|
||||
soup: BeautifulSoup object containing the HTML
|
||||
class_attr_to_keep: a list of tag names whose class attr will be kept
|
||||
|
||||
Returns:
|
||||
BeautifulSoup: Modified soup with attributes removed
|
||||
"""
|
||||
for tag in soup.find_all(True):
|
||||
if tag.name.lower() == "table": # type: ignore
|
||||
continue # We keep table tag
|
||||
tag.attrs.pop("id", None) # type: ignore
|
||||
if tag.name.lower() not in class_attr_to_keep: # type: ignore
|
||||
tag.attrs.pop("class", None) # type: ignore
|
||||
return soup
|
||||
|
||||
|
||||
# Define specific elements
|
||||
class Document(OntologyElement):
|
||||
description: str = Field("Root element of the document", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.layout, frozen=True)
|
||||
allowed_tags: List[str] = Field(["body"], frozen=True)
|
||||
|
||||
|
||||
class Section(OntologyElement):
|
||||
description: str = Field("A distinct part or subdivision of a document", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.layout, frozen=True)
|
||||
allowed_tags: List[str] = Field(["section"], frozen=True)
|
||||
|
||||
|
||||
class Page(OntologyElement):
|
||||
description: str = Field("A single side of a paper in a document", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.layout, frozen=True)
|
||||
allowed_tags: List[str] = Field(["div"], frozen=True)
|
||||
|
||||
|
||||
class Column(OntologyElement):
|
||||
description: str = Field("A vertical section of a page", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.layout, frozen=True)
|
||||
allowed_tags: List[str] = Field(["div"], frozen=True)
|
||||
|
||||
|
||||
class Paragraph(OntologyElement):
|
||||
description: str = Field("A self-contained unit of discourse in writing", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["p"], frozen=True)
|
||||
|
||||
|
||||
class Header(OntologyElement):
|
||||
description: str = Field("The top section of a page", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["header"], frozen=True)
|
||||
|
||||
|
||||
class Footer(OntologyElement):
|
||||
description: str = Field("The bottom section of a page", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["footer"], frozen=True)
|
||||
|
||||
|
||||
class Sidebar(OntologyElement):
|
||||
description: str = Field("A side section of a page", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.layout, frozen=True)
|
||||
allowed_tags: List[str] = Field(["aside"], frozen=True)
|
||||
|
||||
|
||||
class PageBreak(OntologyElement):
|
||||
description: str = Field("A break between pages", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.layout, frozen=True)
|
||||
allowed_tags: List[str] = Field(["hr"], frozen=True)
|
||||
|
||||
|
||||
class Title(OntologyElement):
|
||||
description: str = Field("Main heading of a document or section", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["h1"], frozen=True)
|
||||
|
||||
|
||||
class Subtitle(OntologyElement):
|
||||
description: str = Field("Secondary title of a document or section", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["h2"], frozen=True)
|
||||
|
||||
|
||||
class Heading(OntologyElement):
|
||||
description: str = Field("Section headings (levels 1-6)", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["h1", "h2", "h3", "h4", "h5", "h6"], frozen=True)
|
||||
|
||||
|
||||
class NarrativeText(OntologyElement):
|
||||
description: str = Field("Main content text", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["p"], frozen=True)
|
||||
|
||||
|
||||
class Quote(OntologyElement):
|
||||
description: str = Field("A repetition of someone else's statement", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["blockquote"], frozen=True)
|
||||
|
||||
|
||||
class Footnote(OntologyElement):
|
||||
description: str = Field("A note at the bottom of a page", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["div"], frozen=True)
|
||||
|
||||
|
||||
class Caption(OntologyElement):
|
||||
description: str = Field("Text describing a figure or image", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["figcaption"], frozen=True)
|
||||
|
||||
|
||||
class PageNumber(OntologyElement):
|
||||
description: str = Field("The number of a page", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["span"], frozen=True)
|
||||
|
||||
|
||||
class UncategorizedText(OntologyElement):
|
||||
description: str = Field("Miscellaneous text", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["span"], frozen=True)
|
||||
|
||||
|
||||
class OrderedList(OntologyElement):
|
||||
description: str = Field("A list with a specific sequence", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.list, frozen=True)
|
||||
allowed_tags: List[str] = Field(["ol"], frozen=True)
|
||||
|
||||
|
||||
class UnorderedList(OntologyElement):
|
||||
description: str = Field("A list without a specific sequence", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.list, frozen=True)
|
||||
allowed_tags: List[str] = Field(["ul"], frozen=True)
|
||||
|
||||
|
||||
class DefinitionList(OntologyElement):
|
||||
description: str = Field("A list of terms and their definitions", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.list, frozen=True)
|
||||
allowed_tags: List[str] = Field(["dl"], frozen=True)
|
||||
|
||||
|
||||
class ListItem(OntologyElement):
|
||||
description: str = Field("An item in a list", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.list, frozen=True)
|
||||
allowed_tags: List[str] = Field(["li"], frozen=True)
|
||||
|
||||
|
||||
class Table(OntologyElement):
|
||||
description: str = Field("A structured set of data", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.table, frozen=True)
|
||||
allowed_tags: List[str] = Field(["table"], frozen=True)
|
||||
|
||||
def to_html(self, add_children: bool = True) -> str:
|
||||
soup = BeautifulSoup(super().to_html(add_children), "html.parser")
|
||||
soup = remove_ids_and_class_from_table(soup)
|
||||
return str(soup)
|
||||
|
||||
|
||||
class TableBody(OntologyElement):
|
||||
description: str = Field("A body of the table", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.table, frozen=True)
|
||||
allowed_tags: List[str] = Field(["tbody"], frozen=True)
|
||||
|
||||
|
||||
class TableHeader(OntologyElement):
|
||||
description: str = Field("A header of the table", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.table, frozen=True)
|
||||
allowed_tags: List[str] = Field(["thead"], frozen=True)
|
||||
|
||||
|
||||
class TableRow(OntologyElement):
|
||||
description: str = Field("A row in a table", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.table, frozen=True)
|
||||
allowed_tags: List[str] = Field(["tr"], frozen=True)
|
||||
|
||||
|
||||
class TableCell(OntologyElement):
|
||||
description: str = Field("A cell in a table", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.table, frozen=True)
|
||||
allowed_tags: List[str] = Field(["td"], frozen=True)
|
||||
|
||||
|
||||
# Note(Pluto): Renamed from TableCellHeader to TableHeaderCell to be consistent with TableCell
|
||||
class TableCellHeader(OntologyElement):
|
||||
description: str = Field("A header cell in a table", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.table, frozen=True)
|
||||
allowed_tags: List[str] = Field(["th"], frozen=True)
|
||||
|
||||
|
||||
class Image(OntologyElement):
|
||||
description: str = Field("A visual representation", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.media, frozen=True)
|
||||
allowed_tags: List[str] = Field(["img"], frozen=True)
|
||||
|
||||
|
||||
class Figure(OntologyElement):
|
||||
description: str = Field("An illustration or diagram in a document", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.media, frozen=True)
|
||||
allowed_tags: List[str] = Field(["figure"], frozen=True)
|
||||
|
||||
|
||||
class Video(OntologyElement):
|
||||
description: str = Field("A moving visual media element", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.media, frozen=True)
|
||||
allowed_tags: List[str] = Field(["video"], frozen=True)
|
||||
|
||||
|
||||
class Audio(OntologyElement):
|
||||
description: str = Field("A sound or music element", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.media, frozen=True)
|
||||
allowed_tags: List[str] = Field(["audio"], frozen=True)
|
||||
|
||||
|
||||
class Barcode(OntologyElement):
|
||||
description: str = Field("A machine-readable representation of data", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.media, frozen=True)
|
||||
allowed_tags: List[str] = Field(["img"], frozen=True)
|
||||
|
||||
|
||||
class QRCode(OntologyElement):
|
||||
description: str = Field("A two-dimensional barcode", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.media, frozen=True)
|
||||
allowed_tags: List[str] = Field(["img"], frozen=True)
|
||||
|
||||
|
||||
class Logo(OntologyElement):
|
||||
description: str = Field("A graphical representation of a company or brand", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.media, frozen=True)
|
||||
allowed_tags: List[str] = Field(["img"], frozen=True)
|
||||
|
||||
|
||||
class CodeBlock(OntologyElement):
|
||||
description: str = Field("A block of programming code", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.code, frozen=True)
|
||||
allowed_tags: List[str] = Field(["pre", "code"], frozen=True)
|
||||
|
||||
|
||||
class InlineCode(OntologyElement):
|
||||
description: str = Field("Code within a line of text", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.code, frozen=True)
|
||||
allowed_tags: List[str] = Field(["code"], frozen=True)
|
||||
|
||||
|
||||
class Formula(OntologyElement):
|
||||
description: str = Field("A mathematical formula", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.mathematical, frozen=True)
|
||||
allowed_tags: List[str] = Field(["math"], frozen=True)
|
||||
|
||||
|
||||
class Equation(OntologyElement):
|
||||
description: str = Field("A mathematical equation", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.mathematical, frozen=True)
|
||||
allowed_tags: List[str] = Field(["math"], frozen=True)
|
||||
|
||||
|
||||
class FootnoteReference(OntologyElement):
|
||||
description: str = Field(
|
||||
"A subscripted reference to a note at the bottom of a page", frozen=True
|
||||
)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.reference, frozen=True)
|
||||
allowed_tags: List[str] = Field(["sub"], frozen=True)
|
||||
|
||||
|
||||
class Citation(OntologyElement):
|
||||
description: str = Field("A reference to a source", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.reference, frozen=True)
|
||||
allowed_tags: List[str] = Field(["cite"], frozen=True)
|
||||
|
||||
|
||||
class Bibliography(OntologyElement):
|
||||
description: str = Field("A list of sources", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.reference, frozen=True)
|
||||
allowed_tags: List[str] = Field(["ul"], frozen=True)
|
||||
|
||||
|
||||
class Glossary(OntologyElement):
|
||||
description: str = Field("A list of terms and their definitions", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.reference, frozen=True)
|
||||
allowed_tags: List[str] = Field(["dl"], frozen=True)
|
||||
|
||||
|
||||
class Author(OntologyElement):
|
||||
description: str = Field("The creator of the document", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.metadata, frozen=True)
|
||||
allowed_tags: List[str] = Field(["meta"], frozen=True)
|
||||
|
||||
|
||||
class MetaDate(OntologyElement):
|
||||
description: str = Field("The date associated with the document", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.metadata, frozen=True)
|
||||
allowed_tags: List[str] = Field(["meta"], frozen=True)
|
||||
|
||||
|
||||
class Keywords(OntologyElement):
|
||||
description: str = Field("Key terms associated with the document", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.metadata, frozen=True)
|
||||
allowed_tags: List[str] = Field(["meta"], frozen=True)
|
||||
|
||||
|
||||
class Abstract(OntologyElement):
|
||||
description: str = Field("A summary of the document", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.metadata, frozen=True)
|
||||
allowed_tags: List[str] = Field(["section"], frozen=True)
|
||||
|
||||
|
||||
class Hyperlink(OntologyElement):
|
||||
description: str = Field("A reference to data that can be directly followed", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.navigation, frozen=True)
|
||||
allowed_tags: List[str] = Field(["a"], frozen=True)
|
||||
|
||||
|
||||
class TableOfContents(OntologyElement):
|
||||
description: str = Field(
|
||||
"A list of the document's contents. Total table columns will be "
|
||||
"equal to the degree of hierarchy (n) plus 1 for the target value. "
|
||||
"Header Row: L1,L2,...Ln,Value",
|
||||
frozen=True,
|
||||
)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.table, frozen=True)
|
||||
allowed_tags: List[str] = Field(["table"], frozen=True)
|
||||
|
||||
def to_html(self, add_children: bool = True) -> str:
|
||||
soup = BeautifulSoup(super().to_html(add_children), "html.parser")
|
||||
soup = remove_ids_and_class_from_table(soup)
|
||||
return str(soup)
|
||||
|
||||
|
||||
class Index(OntologyElement):
|
||||
description: str = Field("An alphabetical list of terms and their page numbers", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.navigation, frozen=True)
|
||||
allowed_tags: List[str] = Field(["nav"], frozen=True)
|
||||
|
||||
|
||||
class Form(OntologyElement):
|
||||
description: str = Field("A document section with interactive controls", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.form, frozen=True)
|
||||
allowed_tags: List[str] = Field(["form"], frozen=True)
|
||||
|
||||
|
||||
class FormField(OntologyElement):
|
||||
description: str = Field("A property value of a form", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.form, frozen=True)
|
||||
allowed_tags: List[str] = Field(["label"], frozen=True)
|
||||
|
||||
|
||||
class FormFieldValue(OntologyElement):
|
||||
description: str = Field("A field for user input", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.form, frozen=True)
|
||||
allowed_tags: List[str] = Field(["input"], frozen=True)
|
||||
|
||||
def to_text(self, add_children: bool = True, add_img_alt_text: bool = True) -> str:
|
||||
text = super().to_text(add_children, add_img_alt_text)
|
||||
value = self.additional_attributes.get("value", "")
|
||||
if not value:
|
||||
return text
|
||||
return f"{text} {value}".strip()
|
||||
|
||||
|
||||
class Checkbox(OntologyElement):
|
||||
description: str = Field("A small box that can be checked or unchecked", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.form, frozen=True)
|
||||
allowed_tags: List[str] = Field(["input"], frozen=True)
|
||||
|
||||
|
||||
class RadioButton(OntologyElement):
|
||||
description: str = Field("A circular button that can be selected", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.form, frozen=True)
|
||||
allowed_tags: List[str] = Field(["input"], frozen=True)
|
||||
|
||||
|
||||
class Button(OntologyElement):
|
||||
description: str = Field("An interactive button element", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.form, frozen=True)
|
||||
allowed_tags: List[str] = Field(["button"], frozen=True)
|
||||
|
||||
|
||||
class Comment(OntologyElement):
|
||||
description: str = Field("A note or remark", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.annotation, frozen=True)
|
||||
allowed_tags: List[str] = Field(["span"], frozen=True)
|
||||
|
||||
|
||||
class Highlight(OntologyElement):
|
||||
description: str = Field("Emphasized text or section", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.annotation, frozen=True)
|
||||
allowed_tags: List[str] = Field(["mark"], frozen=True)
|
||||
|
||||
|
||||
class RevisionInsertion(OntologyElement):
|
||||
description: str = Field("A changed or edited element", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.annotation, frozen=True)
|
||||
allowed_tags: List[str] = Field(["ins"], frozen=True)
|
||||
|
||||
|
||||
class RevisionDeletion(OntologyElement):
|
||||
description: str = Field("A changed or edited element", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.annotation, frozen=True)
|
||||
allowed_tags: List[str] = Field(["del"], frozen=True)
|
||||
|
||||
|
||||
class Address(OntologyElement):
|
||||
description: str = Field("A physical location", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.specialized_text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["address"], frozen=True)
|
||||
|
||||
|
||||
class EmailAddress(OntologyElement):
|
||||
description: str = Field("An email address", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.specialized_text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["a"], frozen=True)
|
||||
|
||||
|
||||
class PhoneNumber(OntologyElement):
|
||||
description: str = Field("A telephone number", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.specialized_text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["span"], frozen=True)
|
||||
|
||||
|
||||
class CalendarDate(OntologyElement):
|
||||
description: str = Field("A calendar date", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.specialized_text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["time"], frozen=True)
|
||||
|
||||
|
||||
class Time(OntologyElement):
|
||||
description: str = Field("A specific time", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.specialized_text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["time"], frozen=True)
|
||||
|
||||
|
||||
class Currency(OntologyElement):
|
||||
description: str = Field("A monetary value", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.specialized_text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["span"], frozen=True)
|
||||
|
||||
|
||||
class Measurement(OntologyElement):
|
||||
description: str = Field("A quantitative value with units", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.specialized_text, frozen=True)
|
||||
allowed_tags: List[str] = Field(["span"], frozen=True)
|
||||
|
||||
|
||||
class Letterhead(OntologyElement):
|
||||
description: str = Field("The heading at the top of a letter", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.document_specific, frozen=True)
|
||||
allowed_tags: List[str] = Field(["header"], frozen=True)
|
||||
|
||||
|
||||
class Signature(OntologyElement):
|
||||
description: str = Field("A person's name written in a distinctive way", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.document_specific, frozen=True)
|
||||
allowed_tags: List[str] = Field(["img", "svg"], frozen=True)
|
||||
|
||||
|
||||
class Watermark(OntologyElement):
|
||||
description: str = Field("A faint design made in paper during manufacture", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.document_specific, frozen=True)
|
||||
allowed_tags: List[str] = Field(["div"], frozen=True)
|
||||
|
||||
|
||||
class Stamp(OntologyElement):
|
||||
description: str = Field("An official mark or seal", frozen=True)
|
||||
elementType: ElementTypeEnum = Field(ElementTypeEnum.document_specific, frozen=True)
|
||||
allowed_tags: List[str] = Field(["img", "svg"], frozen=True)
|
||||
@@ -0,0 +1,27 @@
|
||||
import warnings
|
||||
|
||||
from unstructured.embed.bedrock import BedrockEmbeddingEncoder
|
||||
from unstructured.embed.huggingface import HuggingFaceEmbeddingEncoder
|
||||
from unstructured.embed.mixedbreadai import MixedbreadAIEmbeddingEncoder
|
||||
from unstructured.embed.octoai import OctoAIEmbeddingEncoder
|
||||
from unstructured.embed.openai import OpenAIEmbeddingEncoder
|
||||
from unstructured.embed.vertexai import VertexAIEmbeddingEncoder
|
||||
from unstructured.embed.voyageai import VoyageAIEmbeddingEncoder
|
||||
|
||||
EMBEDDING_PROVIDER_TO_CLASS_MAP = {
|
||||
"langchain-openai": OpenAIEmbeddingEncoder,
|
||||
"langchain-huggingface": HuggingFaceEmbeddingEncoder,
|
||||
"langchain-aws-bedrock": BedrockEmbeddingEncoder,
|
||||
"langchain-vertexai": VertexAIEmbeddingEncoder,
|
||||
"voyageai": VoyageAIEmbeddingEncoder,
|
||||
"mixedbread-ai": MixedbreadAIEmbeddingEncoder,
|
||||
"octoai": OctoAIEmbeddingEncoder,
|
||||
}
|
||||
|
||||
|
||||
warnings.warn(
|
||||
"unstructured.ingest will be removed in a future version. "
|
||||
"Functionality moved to the unstructured-ingest project.",
|
||||
DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
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,76 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, List
|
||||
|
||||
import numpy as np
|
||||
from pydantic import SecretStr
|
||||
|
||||
from unstructured.documents.elements import (
|
||||
Element,
|
||||
)
|
||||
from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_community.embeddings import BedrockEmbeddings
|
||||
|
||||
|
||||
class BedrockEmbeddingConfig(EmbeddingConfig):
|
||||
aws_access_key_id: SecretStr
|
||||
aws_secret_access_key: SecretStr
|
||||
region_name: str = "us-west-2"
|
||||
|
||||
@requires_dependencies(
|
||||
["boto3", "numpy", "langchain_community"],
|
||||
extras="bedrock",
|
||||
)
|
||||
def get_client(self) -> "BedrockEmbeddings":
|
||||
# delay import only when needed
|
||||
import boto3
|
||||
from langchain_community.embeddings import BedrockEmbeddings
|
||||
|
||||
bedrock_runtime = boto3.client(
|
||||
service_name="bedrock-runtime",
|
||||
aws_access_key_id=self.aws_access_key_id.get_secret_value(),
|
||||
aws_secret_access_key=self.aws_secret_access_key.get_secret_value(),
|
||||
region_name=self.region_name,
|
||||
)
|
||||
|
||||
bedrock_client = BedrockEmbeddings(client=bedrock_runtime)
|
||||
return bedrock_client
|
||||
|
||||
|
||||
@dataclass
|
||||
class BedrockEmbeddingEncoder(BaseEmbeddingEncoder):
|
||||
config: BedrockEmbeddingConfig
|
||||
|
||||
def get_exemplary_embedding(self) -> List[float]:
|
||||
return self.embed_query(query="Q")
|
||||
|
||||
def __post_init__(self):
|
||||
self.initialize()
|
||||
|
||||
def num_of_dimensions(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.shape(exemplary_embedding)
|
||||
|
||||
def is_unit_vector(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
|
||||
|
||||
def embed_query(self, query):
|
||||
bedrock_client = self.config.get_client()
|
||||
return np.array(bedrock_client.embed_query(query))
|
||||
|
||||
def embed_documents(self, elements: List[Element]) -> List[Element]:
|
||||
bedrock_client = self.config.get_client()
|
||||
embeddings = bedrock_client.embed_documents([str(e) for e in elements])
|
||||
elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
|
||||
return elements_with_embeddings
|
||||
|
||||
def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
|
||||
assert len(elements) == len(embeddings)
|
||||
elements_w_embedding = []
|
||||
for i, element in enumerate(elements):
|
||||
element.embeddings = embeddings[i]
|
||||
elements_w_embedding.append(element)
|
||||
return elements
|
||||
@@ -0,0 +1,67 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from pydantic import Field
|
||||
|
||||
from unstructured.documents.elements import (
|
||||
Element,
|
||||
)
|
||||
from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
|
||||
|
||||
|
||||
class HuggingFaceEmbeddingConfig(EmbeddingConfig):
|
||||
model_name: Optional[str] = Field(default="sentence-transformers/all-MiniLM-L6-v2")
|
||||
model_kwargs: Optional[dict] = Field(default_factory=lambda: {"device": "cpu"})
|
||||
encode_kwargs: Optional[dict] = Field(default_factory=lambda: {"normalize_embeddings": False})
|
||||
cache_folder: Optional[dict] = Field(default=None)
|
||||
|
||||
@requires_dependencies(
|
||||
["langchain_huggingface"],
|
||||
extras="embed-huggingface",
|
||||
)
|
||||
def get_client(self) -> "HuggingFaceEmbeddings":
|
||||
"""Creates a langchain Huggingface python client to embed elements."""
|
||||
from langchain_huggingface.embeddings import HuggingFaceEmbeddings
|
||||
|
||||
client = HuggingFaceEmbeddings(**self.dict())
|
||||
return client
|
||||
|
||||
|
||||
@dataclass
|
||||
class HuggingFaceEmbeddingEncoder(BaseEmbeddingEncoder):
|
||||
config: HuggingFaceEmbeddingConfig
|
||||
|
||||
def get_exemplary_embedding(self) -> List[float]:
|
||||
return self.embed_query(query="Q")
|
||||
|
||||
def num_of_dimensions(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.shape(exemplary_embedding)
|
||||
|
||||
def is_unit_vector(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
|
||||
|
||||
def embed_query(self, query):
|
||||
client = self.config.get_client()
|
||||
return client.embed_query(str(query))
|
||||
|
||||
def embed_documents(self, elements: List[Element]) -> List[Element]:
|
||||
client = self.config.get_client()
|
||||
embeddings = client.embed_documents([str(e) for e in elements])
|
||||
elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
|
||||
return elements_with_embeddings
|
||||
|
||||
def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
|
||||
assert len(elements) == len(embeddings)
|
||||
elements_w_embedding = []
|
||||
|
||||
for i, element in enumerate(elements):
|
||||
element.embeddings = embeddings[i]
|
||||
elements_w_embedding.append(element)
|
||||
return elements
|
||||
@@ -0,0 +1,39 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Tuple
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
|
||||
|
||||
class EmbeddingConfig(BaseModel):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class BaseEmbeddingEncoder(ABC):
|
||||
config: EmbeddingConfig
|
||||
|
||||
@abstractmethod
|
||||
def initialize(self):
|
||||
"""Initializes the embedding encoder class. Should also validate the instance
|
||||
is properly configured: e.g., embed a single a element"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def num_of_dimensions(self) -> Tuple[int]:
|
||||
"""Number of dimensions for the embedding vector."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def is_unit_vector(self) -> bool:
|
||||
"""Denotes if the embedding vector is a unit vector."""
|
||||
|
||||
@abstractmethod
|
||||
def embed_documents(self, elements: List[Element]) -> List[Element]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def embed_query(self, query: str) -> List[float]:
|
||||
pass
|
||||
@@ -0,0 +1,178 @@
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from pydantic import Field, SecretStr
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
USER_AGENT = "@mixedbread-ai/unstructured"
|
||||
BATCH_SIZE = 128
|
||||
TIMEOUT = 60
|
||||
MAX_RETRIES = 3
|
||||
ENCODING_FORMAT = "float"
|
||||
TRUNCATION_STRATEGY = "end"
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from mixedbread_ai.client import MixedbreadAI
|
||||
from mixedbread_ai.core import RequestOptions
|
||||
|
||||
|
||||
class MixedbreadAIEmbeddingConfig(EmbeddingConfig):
|
||||
"""
|
||||
Configuration class for Mixedbread AI Embedding Encoder.
|
||||
|
||||
Attributes:
|
||||
api_key (str): API key for accessing Mixedbread AI..
|
||||
model_name (str): Name of the model to use for embeddings.
|
||||
"""
|
||||
|
||||
api_key: SecretStr = Field(
|
||||
default_factory=lambda: SecretStr(os.environ.get("MXBAI_API_KEY")),
|
||||
)
|
||||
|
||||
model_name: str = Field(
|
||||
default="mixedbread-ai/mxbai-embed-large-v1",
|
||||
)
|
||||
|
||||
@requires_dependencies(
|
||||
["mixedbread_ai"],
|
||||
extras="embed-mixedbreadai",
|
||||
)
|
||||
def get_client(self) -> "MixedbreadAI":
|
||||
"""
|
||||
Create the Mixedbread AI client.
|
||||
|
||||
Returns:
|
||||
MixedbreadAI: Initialized client.
|
||||
"""
|
||||
from mixedbread_ai.client import MixedbreadAI
|
||||
|
||||
return MixedbreadAI(
|
||||
api_key=self.api_key.get_secret_value(),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MixedbreadAIEmbeddingEncoder(BaseEmbeddingEncoder):
|
||||
"""
|
||||
Embedding encoder for Mixedbread AI.
|
||||
|
||||
Attributes:
|
||||
config (MixedbreadAIEmbeddingConfig): Configuration for the embedding encoder.
|
||||
"""
|
||||
|
||||
config: MixedbreadAIEmbeddingConfig
|
||||
|
||||
_exemplary_embedding: Optional[List[float]] = field(init=False, default=None)
|
||||
_request_options: Optional["RequestOptions"] = field(init=False, default=None)
|
||||
|
||||
def get_exemplary_embedding(self) -> List[float]:
|
||||
"""Get an exemplary embedding to determine dimensions and unit vector status."""
|
||||
return self._embed(["Q"])[0]
|
||||
|
||||
def initialize(self):
|
||||
if self.config.api_key is None:
|
||||
raise ValueError(
|
||||
"The Mixedbread AI API key must be specified."
|
||||
+ "You either pass it in the constructor using 'api_key'"
|
||||
+ "or via the 'MXBAI_API_KEY' environment variable."
|
||||
)
|
||||
|
||||
from mixedbread_ai.core import RequestOptions
|
||||
|
||||
self._request_options = RequestOptions(
|
||||
max_retries=MAX_RETRIES,
|
||||
timeout_in_seconds=TIMEOUT,
|
||||
additional_headers={"User-Agent": USER_AGENT},
|
||||
)
|
||||
|
||||
@property
|
||||
def num_of_dimensions(self):
|
||||
"""Get the number of dimensions for the embeddings."""
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.shape(exemplary_embedding)
|
||||
|
||||
@property
|
||||
def is_unit_vector(self) -> bool:
|
||||
"""Check if the embedding is a unit vector."""
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
|
||||
|
||||
def _embed(self, texts: List[str]) -> List[List[float]]:
|
||||
"""
|
||||
Embed a list of texts using the Mixedbread AI API.
|
||||
|
||||
Args:
|
||||
texts (List[str]): List of texts to embed.
|
||||
|
||||
Returns:
|
||||
List[List[float]]: List of embeddings.
|
||||
"""
|
||||
batch_size = BATCH_SIZE
|
||||
batch_itr = range(0, len(texts), batch_size)
|
||||
|
||||
responses = []
|
||||
client = self.config.get_client()
|
||||
for i in batch_itr:
|
||||
batch = texts[i : i + batch_size]
|
||||
response = client.embeddings(
|
||||
model=self.config.model_name,
|
||||
normalized=True,
|
||||
encoding_format=ENCODING_FORMAT,
|
||||
truncation_strategy=TRUNCATION_STRATEGY,
|
||||
request_options=self._request_options,
|
||||
input=batch,
|
||||
)
|
||||
responses.append(response)
|
||||
return [item.embedding for response in responses for item in response.data]
|
||||
|
||||
@staticmethod
|
||||
def _add_embeddings_to_elements(
|
||||
elements: List[Element], embeddings: List[List[float]]
|
||||
) -> List[Element]:
|
||||
"""
|
||||
Add embeddings to elements.
|
||||
|
||||
Args:
|
||||
elements (List[Element]): List of elements.
|
||||
embeddings (List[List[float]]): List of embeddings.
|
||||
|
||||
Returns:
|
||||
List[Element]: Elements with embeddings added.
|
||||
"""
|
||||
assert len(elements) == len(embeddings)
|
||||
elements_w_embedding = []
|
||||
for i, element in enumerate(elements):
|
||||
element.embeddings = embeddings[i]
|
||||
elements_w_embedding.append(element)
|
||||
return elements
|
||||
|
||||
def embed_documents(self, elements: List[Element]) -> List[Element]:
|
||||
"""
|
||||
Embed a list of document elements.
|
||||
|
||||
Args:
|
||||
elements (List[Element]): List of document elements.
|
||||
|
||||
Returns:
|
||||
List[Element]: Elements with embeddings.
|
||||
"""
|
||||
embeddings = self._embed([str(e) for e in elements])
|
||||
return self._add_embeddings_to_elements(elements, embeddings)
|
||||
|
||||
def embed_query(self, query: str) -> List[float]:
|
||||
"""
|
||||
Embed a query string.
|
||||
|
||||
Args:
|
||||
query (str): Query string to embed.
|
||||
|
||||
Returns:
|
||||
List[float]: Embedding of the query.
|
||||
"""
|
||||
return self._embed([query])[0]
|
||||
@@ -0,0 +1,69 @@
|
||||
from dataclasses import dataclass, field
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from pydantic import Field, SecretStr
|
||||
|
||||
from unstructured.documents.elements import (
|
||||
Element,
|
||||
)
|
||||
from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from openai import OpenAI
|
||||
|
||||
|
||||
class OctoAiEmbeddingConfig(EmbeddingConfig):
|
||||
api_key: SecretStr
|
||||
model_name: str = Field(default="thenlper/gte-large")
|
||||
base_url: str = Field(default="https://text.octoai.run/v1")
|
||||
|
||||
@requires_dependencies(
|
||||
["openai", "tiktoken"],
|
||||
extras="embed-octoai",
|
||||
)
|
||||
def get_client(self) -> "OpenAI":
|
||||
"""Creates an OpenAI python client to embed elements. Uses the OpenAI SDK."""
|
||||
from openai import OpenAI
|
||||
|
||||
return OpenAI(api_key=self.api_key.get_secret_value(), base_url=self.base_url)
|
||||
|
||||
|
||||
@dataclass
|
||||
class OctoAIEmbeddingEncoder(BaseEmbeddingEncoder):
|
||||
config: OctoAiEmbeddingConfig
|
||||
# Uses the OpenAI SDK
|
||||
_exemplary_embedding: Optional[List[float]] = field(init=False, default=None)
|
||||
|
||||
def get_exemplary_embedding(self) -> List[float]:
|
||||
return self.embed_query("Q")
|
||||
|
||||
def initialize(self):
|
||||
pass
|
||||
|
||||
def num_of_dimensions(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.shape(exemplary_embedding)
|
||||
|
||||
def is_unit_vector(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
|
||||
|
||||
def embed_query(self, query):
|
||||
client = self.config.get_client()
|
||||
response = client.embeddings.create(input=str(query), model=self.config.model_name)
|
||||
return response.data[0].embedding
|
||||
|
||||
def embed_documents(self, elements: List[Element]) -> List[Element]:
|
||||
embeddings = [self.embed_query(e) for e in elements]
|
||||
elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
|
||||
return elements_with_embeddings
|
||||
|
||||
def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
|
||||
assert len(elements) == len(embeddings)
|
||||
elements_w_embedding = []
|
||||
for i, element in enumerate(elements):
|
||||
element.embeddings = embeddings[i]
|
||||
elements_w_embedding.append(element)
|
||||
return elements
|
||||
@@ -0,0 +1,67 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, List
|
||||
|
||||
import numpy as np
|
||||
from pydantic import Field, SecretStr
|
||||
|
||||
from unstructured.documents.elements import (
|
||||
Element,
|
||||
)
|
||||
from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_openai.embeddings import OpenAIEmbeddings
|
||||
|
||||
|
||||
class OpenAIEmbeddingConfig(EmbeddingConfig):
|
||||
api_key: SecretStr
|
||||
model_name: str = Field(default="text-embedding-ada-002")
|
||||
|
||||
@requires_dependencies(["langchain_openai"], extras="openai")
|
||||
def get_client(self) -> "OpenAIEmbeddings":
|
||||
"""Creates a langchain OpenAI python client to embed elements."""
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
|
||||
openai_client = OpenAIEmbeddings(
|
||||
openai_api_key=self.api_key.get_secret_value(),
|
||||
model=self.model_name, # type:ignore
|
||||
)
|
||||
return openai_client
|
||||
|
||||
|
||||
@dataclass
|
||||
class OpenAIEmbeddingEncoder(BaseEmbeddingEncoder):
|
||||
config: OpenAIEmbeddingConfig
|
||||
|
||||
def get_exemplary_embedding(self) -> List[float]:
|
||||
return self.embed_query(query="Q")
|
||||
|
||||
def initialize(self):
|
||||
pass
|
||||
|
||||
def num_of_dimensions(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.shape(exemplary_embedding)
|
||||
|
||||
def is_unit_vector(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
|
||||
|
||||
def embed_query(self, query):
|
||||
client = self.config.get_client()
|
||||
return client.embed_query(str(query))
|
||||
|
||||
def embed_documents(self, elements: List[Element]) -> List[Element]:
|
||||
client = self.config.get_client()
|
||||
embeddings = client.embed_documents([str(e) for e in elements])
|
||||
elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
|
||||
return elements_with_embeddings
|
||||
|
||||
def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
|
||||
assert len(elements) == len(embeddings)
|
||||
elements_w_embedding = []
|
||||
for i, element in enumerate(elements):
|
||||
element.embeddings = embeddings[i]
|
||||
elements_w_embedding.append(element)
|
||||
return elements
|
||||
@@ -0,0 +1,78 @@
|
||||
# type: ignore
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from pydantic import Field, SecretStr
|
||||
|
||||
from unstructured.documents.elements import (
|
||||
Element,
|
||||
)
|
||||
from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
|
||||
from unstructured.utils import FileHandler, requires_dependencies
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from langchain_google_vertexai import VertexAIEmbeddings
|
||||
|
||||
|
||||
class VertexAIEmbeddingConfig(EmbeddingConfig):
|
||||
api_key: SecretStr
|
||||
model_name: Optional[str] = Field(default="textembedding-gecko@001")
|
||||
|
||||
def register_application_credentials(self):
|
||||
application_credentials_path = os.path.join("/tmp", "google-vertex-app-credentials.json")
|
||||
credentials_file = FileHandler(application_credentials_path)
|
||||
credentials_file.write_file(json.dumps(json.loads(self.api_key.get_secret_value())))
|
||||
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = application_credentials_path
|
||||
|
||||
@requires_dependencies(
|
||||
["langchain", "langchain_google_vertexai"],
|
||||
extras="embed-vertexai",
|
||||
)
|
||||
def get_client(self) -> "VertexAIEmbeddings":
|
||||
"""Creates a Langchain VertexAI python client to embed elements."""
|
||||
from langchain_google_vertexai import VertexAIEmbeddings
|
||||
|
||||
self.register_application_credentials()
|
||||
vertexai_client = VertexAIEmbeddings(model_name=self.model_name)
|
||||
return vertexai_client
|
||||
|
||||
|
||||
@dataclass
|
||||
class VertexAIEmbeddingEncoder(BaseEmbeddingEncoder):
|
||||
config: VertexAIEmbeddingConfig
|
||||
|
||||
def get_exemplary_embedding(self) -> List[float]:
|
||||
return self.embed_query(query="A sample query.")
|
||||
|
||||
def initialize(self):
|
||||
pass
|
||||
|
||||
def num_of_dimensions(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.shape(exemplary_embedding)
|
||||
|
||||
def is_unit_vector(self):
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
|
||||
|
||||
def embed_query(self, query):
|
||||
client = self.config.get_client()
|
||||
result = client.embed_query(str(query))
|
||||
return result
|
||||
|
||||
def embed_documents(self, elements: List[Element]) -> List[Element]:
|
||||
client = self.config.get_client()
|
||||
embeddings = client.embed_documents([str(e) for e in elements])
|
||||
elements_with_embeddings = self._add_embeddings_to_elements(elements, embeddings)
|
||||
return elements_with_embeddings
|
||||
|
||||
def _add_embeddings_to_elements(self, elements, embeddings) -> List[Element]:
|
||||
assert len(elements) == len(embeddings)
|
||||
elements_w_embedding = []
|
||||
for i, element in enumerate(elements):
|
||||
element.embeddings = embeddings[i]
|
||||
elements_w_embedding.append(element)
|
||||
return elements
|
||||
@@ -0,0 +1,237 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import TYPE_CHECKING, Iterable, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from pydantic import Field, SecretStr
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.embed.interfaces import BaseEmbeddingEncoder, EmbeddingConfig
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from voyageai import Client
|
||||
|
||||
# Token limits for different VoyageAI models
|
||||
VOYAGE_TOTAL_TOKEN_LIMITS = {
|
||||
"voyage-context-3": 32_000,
|
||||
"voyage-3.5-lite": 1_000_000,
|
||||
"voyage-3.5": 320_000,
|
||||
"voyage-2": 320_000,
|
||||
"voyage-02": 320_000,
|
||||
"voyage-3-large": 120_000,
|
||||
"voyage-code-3": 120_000,
|
||||
"voyage-large-2-instruct": 120_000,
|
||||
"voyage-finance-2": 120_000,
|
||||
"voyage-multilingual-2": 120_000,
|
||||
"voyage-law-2": 120_000,
|
||||
"voyage-large-2": 120_000,
|
||||
"voyage-3": 120_000,
|
||||
"voyage-3-lite": 120_000,
|
||||
"voyage-code-2": 120_000,
|
||||
"voyage-3-m-exp": 120_000,
|
||||
"voyage-multimodal-3": 120_000,
|
||||
}
|
||||
|
||||
# Batch size for embedding requests (max documents per batch)
|
||||
MAX_BATCH_SIZE = 1000
|
||||
|
||||
|
||||
class VoyageAIEmbeddingConfig(EmbeddingConfig):
|
||||
api_key: SecretStr
|
||||
model_name: str
|
||||
show_progress_bar: bool = False
|
||||
batch_size: Optional[int] = Field(default=None)
|
||||
truncation: Optional[bool] = Field(default=None)
|
||||
output_dimension: Optional[int] = Field(default=None)
|
||||
|
||||
@requires_dependencies(
|
||||
["voyageai"],
|
||||
extras="embed-voyageai",
|
||||
)
|
||||
def get_client(self) -> "Client":
|
||||
"""Creates a VoyageAI python client to embed elements."""
|
||||
from voyageai import Client
|
||||
|
||||
return Client(
|
||||
api_key=self.api_key.get_secret_value(),
|
||||
)
|
||||
|
||||
def get_token_limit(self) -> int:
|
||||
"""Get the token limit for the current model."""
|
||||
return VOYAGE_TOTAL_TOKEN_LIMITS.get(self.model_name, 120_000)
|
||||
|
||||
|
||||
@dataclass
|
||||
class VoyageAIEmbeddingEncoder(BaseEmbeddingEncoder):
|
||||
config: VoyageAIEmbeddingConfig
|
||||
|
||||
def get_exemplary_embedding(self) -> List[float]:
|
||||
return self.embed_query(query="A sample query.")
|
||||
|
||||
def initialize(self):
|
||||
pass
|
||||
|
||||
@property
|
||||
def num_of_dimensions(self) -> tuple[int, ...]:
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.shape(exemplary_embedding)
|
||||
|
||||
@property
|
||||
def is_unit_vector(self) -> bool:
|
||||
exemplary_embedding = self.get_exemplary_embedding()
|
||||
return np.isclose(np.linalg.norm(exemplary_embedding), 1.0)
|
||||
|
||||
def _is_context_model(self) -> bool:
|
||||
"""Check if the model is a contextualized embedding model."""
|
||||
return "context" in self.config.model_name
|
||||
|
||||
def _build_batches(self, texts: List[str], client: "Client") -> Iterable[List[str]]:
|
||||
"""
|
||||
Generate batches of texts based on token limits.
|
||||
|
||||
Args:
|
||||
texts: List of texts to batch.
|
||||
client: VoyageAI client instance to use for tokenization.
|
||||
|
||||
Yields:
|
||||
Batches of texts as lists.
|
||||
"""
|
||||
if not texts:
|
||||
return
|
||||
|
||||
max_tokens_per_batch = self.config.get_token_limit()
|
||||
current_batch: List[str] = []
|
||||
current_batch_tokens = 0
|
||||
|
||||
# Tokenize all texts in one API call
|
||||
all_token_lists = client.tokenize(texts, model=self.config.model_name)
|
||||
token_counts = [len(tokens) for tokens in all_token_lists]
|
||||
|
||||
for i, text in enumerate(texts):
|
||||
n_tokens = token_counts[i]
|
||||
|
||||
# Check if adding this text would exceed limits
|
||||
if current_batch and (
|
||||
len(current_batch) >= MAX_BATCH_SIZE
|
||||
or (current_batch_tokens + n_tokens > max_tokens_per_batch)
|
||||
):
|
||||
# Yield the current batch and start a new one
|
||||
yield current_batch
|
||||
current_batch = []
|
||||
current_batch_tokens = 0
|
||||
|
||||
current_batch.append(text)
|
||||
current_batch_tokens += n_tokens
|
||||
|
||||
# Yield the last batch (always has at least one text)
|
||||
if current_batch:
|
||||
yield current_batch
|
||||
|
||||
def _embed_batch(
|
||||
self, batch: List[str], client: "Client", input_type: str = "document"
|
||||
) -> List[List[float]]:
|
||||
"""
|
||||
Embed a batch of texts using the appropriate method for the model.
|
||||
|
||||
Args:
|
||||
batch: List of texts to embed.
|
||||
client: VoyageAI client instance to use for embedding.
|
||||
input_type: Type of input ("document" or "query").
|
||||
|
||||
Returns:
|
||||
List of embedding vectors.
|
||||
"""
|
||||
if self._is_context_model():
|
||||
result = client.contextualized_embed(
|
||||
inputs=[batch],
|
||||
model=self.config.model_name,
|
||||
input_type=input_type,
|
||||
output_dimension=self.config.output_dimension,
|
||||
)
|
||||
return [list(emb) for emb in result.results[0].embeddings]
|
||||
else:
|
||||
result = client.embed(
|
||||
texts=batch,
|
||||
model=self.config.model_name,
|
||||
input_type=input_type,
|
||||
truncation=self.config.truncation,
|
||||
output_dimension=self.config.output_dimension,
|
||||
)
|
||||
return [list(emb) for emb in result.embeddings]
|
||||
|
||||
def embed_documents(self, elements: List[Element]) -> List[Element]:
|
||||
"""
|
||||
Embed documents with automatic batching based on token limits.
|
||||
|
||||
Args:
|
||||
elements: List of elements to embed.
|
||||
|
||||
Returns:
|
||||
List of elements with embeddings added.
|
||||
"""
|
||||
if not elements:
|
||||
return []
|
||||
|
||||
client = self.config.get_client()
|
||||
texts = [str(e) for e in elements]
|
||||
all_embeddings: List[List[float]] = []
|
||||
|
||||
# Process each batch
|
||||
batches = list(self._build_batches(texts, client))
|
||||
|
||||
if self.config.show_progress_bar:
|
||||
try:
|
||||
from tqdm.auto import tqdm # type: ignore
|
||||
|
||||
batches = tqdm(batches, desc="Embedding batches")
|
||||
except ImportError as e:
|
||||
raise ImportError(
|
||||
"Must have tqdm installed if `show_progress_bar` is set to True. "
|
||||
"Please install with `pip install tqdm`."
|
||||
) from e
|
||||
|
||||
for batch in batches:
|
||||
batch_embeddings = self._embed_batch(batch, client, input_type="document")
|
||||
all_embeddings.extend(batch_embeddings)
|
||||
|
||||
return self._add_embeddings_to_elements(elements, all_embeddings)
|
||||
|
||||
def embed_query(self, query: str) -> List[float]:
|
||||
"""
|
||||
Embed a single query string.
|
||||
|
||||
Args:
|
||||
query: Query string to embed.
|
||||
|
||||
Returns:
|
||||
Embedding vector.
|
||||
"""
|
||||
client = self.config.get_client()
|
||||
batch_embeddings = self._embed_batch([query], client, input_type="query")
|
||||
return batch_embeddings[0]
|
||||
|
||||
def count_tokens(self, texts: List[str]) -> List[int]:
|
||||
"""
|
||||
Count tokens for the given texts.
|
||||
|
||||
Args:
|
||||
texts: List of texts to count tokens for.
|
||||
|
||||
Returns:
|
||||
List of token counts for each text.
|
||||
"""
|
||||
if not texts:
|
||||
return []
|
||||
|
||||
client = self.config.get_client()
|
||||
token_lists = client.tokenize(texts, model=self.config.model_name)
|
||||
return [len(token_list) for token_list in token_lists]
|
||||
|
||||
@staticmethod
|
||||
def _add_embeddings_to_elements(elements, embeddings) -> List[Element]:
|
||||
assert len(elements) == len(embeddings)
|
||||
elements_w_embedding = []
|
||||
for i, element in enumerate(elements):
|
||||
element.embeddings = embeddings[i]
|
||||
elements_w_embedding.append(element)
|
||||
return elements
|
||||
@@ -0,0 +1,15 @@
|
||||
class PageCountExceededError(ValueError):
|
||||
"""Error raised, when number of pages exceeds pdf_hi_res_max_pages limit."""
|
||||
|
||||
def __init__(self, document_pages: int, pdf_hi_res_max_pages: int):
|
||||
self.document_pages = document_pages
|
||||
self.pdf_hi_res_max_pages = pdf_hi_res_max_pages
|
||||
self.message = (
|
||||
f"Maximum number of PDF file pages exceeded - "
|
||||
f"pages={document_pages}, maximum={pdf_hi_res_max_pages}."
|
||||
)
|
||||
super().__init__(self.message)
|
||||
|
||||
|
||||
class UnprocessableEntityError(Exception):
|
||||
"""Error raised when a file is not valid."""
|
||||
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,150 @@
|
||||
from typing import IO, Optional, Tuple, Union
|
||||
|
||||
from charset_normalizer import detect
|
||||
|
||||
from unstructured.errors import UnprocessableEntityError
|
||||
from unstructured.partition.common.common import convert_to_bytes
|
||||
|
||||
ENCODE_REC_THRESHOLD = 0.8
|
||||
|
||||
# popular encodings from https://en.wikipedia.org/wiki/Popularity_of_text_encodings
|
||||
COMMON_ENCODINGS = [
|
||||
"utf_8",
|
||||
"iso_8859_1",
|
||||
"iso_8859_6",
|
||||
"iso_8859_8",
|
||||
"ascii",
|
||||
"big5",
|
||||
"utf_16",
|
||||
"utf_16_be",
|
||||
"utf_16_le",
|
||||
"utf_32",
|
||||
"utf_32_be",
|
||||
"utf_32_le",
|
||||
"euc_jis_2004",
|
||||
"euc_jisx0213",
|
||||
"euc_jp",
|
||||
"euc_kr",
|
||||
"gb18030",
|
||||
"shift_jis",
|
||||
"shift_jis_2004",
|
||||
"shift_jisx0213",
|
||||
]
|
||||
|
||||
|
||||
def format_encoding_str(encoding: str) -> str:
|
||||
"""Format input encoding string (e.g., `utf-8`, `iso-8859-1`, etc).
|
||||
Parameters
|
||||
----------
|
||||
encoding
|
||||
The encoding string to be formatted (e.g., `UTF-8`, `utf_8`, `ISO-8859-1`, `iso_8859_1`,
|
||||
etc).
|
||||
"""
|
||||
formatted_encoding = encoding.lower().replace("_", "-")
|
||||
|
||||
# Special case for Arabic and Hebrew charsets with directional annotations
|
||||
annotated_encodings = ["iso-8859-6-i", "iso-8859-6-e", "iso-8859-8-i", "iso-8859-8-e"]
|
||||
if formatted_encoding in annotated_encodings:
|
||||
formatted_encoding = formatted_encoding[:-2] # remove the annotation
|
||||
|
||||
return formatted_encoding
|
||||
|
||||
|
||||
def validate_encoding(encoding: str) -> bool:
|
||||
"""Checks if an encoding string is valid. Helps to avoid errors in cases where
|
||||
invalid encodings are extracted from malformed documents."""
|
||||
for common_encoding in COMMON_ENCODINGS:
|
||||
if format_encoding_str(common_encoding) == format_encoding_str(encoding):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def detect_file_encoding(
|
||||
filename: str = "",
|
||||
file: Optional[Union[bytes, IO[bytes]]] = None,
|
||||
) -> Tuple[str, str]:
|
||||
if filename:
|
||||
with open(filename, "rb") as f:
|
||||
byte_data = f.read()
|
||||
elif file:
|
||||
byte_data = convert_to_bytes(file)
|
||||
else:
|
||||
raise FileNotFoundError("No filename nor file were specified")
|
||||
|
||||
result = detect(byte_data)
|
||||
encoding = result["encoding"]
|
||||
confidence = result["confidence"]
|
||||
|
||||
if encoding is None or confidence is None or confidence < ENCODE_REC_THRESHOLD:
|
||||
# Encoding detection failed, fallback to predefined encodings
|
||||
for enc in COMMON_ENCODINGS:
|
||||
try:
|
||||
if filename:
|
||||
with open(filename, encoding=enc) as f:
|
||||
file_text = f.read()
|
||||
else:
|
||||
file_text = byte_data.decode(enc)
|
||||
encoding = enc
|
||||
break
|
||||
except (UnicodeDecodeError, UnicodeError):
|
||||
continue
|
||||
else:
|
||||
# NOTE: Use UnprocessableEntityError instead of UnicodeDecodeError to avoid
|
||||
# logging the entire file content. UnicodeDecodeError automatically stores
|
||||
# the complete input data, which can be problematic for large files.
|
||||
raise UnprocessableEntityError(
|
||||
"Unable to determine file encoding after trying all common encodings. "
|
||||
"File may be corrupted or in an unsupported format."
|
||||
) from None
|
||||
|
||||
else:
|
||||
# NOTE: Catch UnicodeDecodeError to avoid logging the entire file content.
|
||||
# UnicodeDecodeError automatically stores the complete input data in its
|
||||
# 'object' attribute, which can cause issues with large files in logging
|
||||
# and error reporting systems.
|
||||
try:
|
||||
file_text = byte_data.decode(encoding)
|
||||
except (UnicodeDecodeError, UnicodeError):
|
||||
raise UnprocessableEntityError(
|
||||
f"File encoding detection failed: detected '{encoding}' but decode failed. "
|
||||
f"File may be corrupted or in an unsupported format."
|
||||
) from None
|
||||
|
||||
formatted_encoding = format_encoding_str(encoding)
|
||||
|
||||
return formatted_encoding, file_text
|
||||
|
||||
|
||||
def read_txt_file(
|
||||
filename: str = "",
|
||||
file: Optional[Union[bytes, IO[bytes]]] = None,
|
||||
encoding: Optional[str] = None,
|
||||
) -> Tuple[str, str]:
|
||||
"""Extracts document metadata from a plain text document."""
|
||||
if filename:
|
||||
if encoding:
|
||||
formatted_encoding = format_encoding_str(encoding)
|
||||
with open(filename, encoding=formatted_encoding) as f:
|
||||
try:
|
||||
file_text = f.read()
|
||||
except (UnicodeDecodeError, UnicodeError) as error:
|
||||
raise error
|
||||
else:
|
||||
formatted_encoding, file_text = detect_file_encoding(filename)
|
||||
elif file:
|
||||
if encoding:
|
||||
formatted_encoding = format_encoding_str(encoding)
|
||||
try:
|
||||
file_content = file if isinstance(file, bytes) else file.read()
|
||||
if isinstance(file_content, bytes):
|
||||
file_text = file_content.decode(formatted_encoding)
|
||||
else:
|
||||
file_text = file_content
|
||||
except (UnicodeDecodeError, UnicodeError) as error:
|
||||
raise error
|
||||
else:
|
||||
formatted_encoding, file_text = detect_file_encoding(file=file)
|
||||
else:
|
||||
raise FileNotFoundError("No filename was specified")
|
||||
|
||||
return formatted_encoding, file_text
|
||||
@@ -0,0 +1,81 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
from typing import IO
|
||||
|
||||
from unstructured.errors import UnprocessableEntityError
|
||||
from unstructured.partition.common.common import exactly_one
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
|
||||
@requires_dependencies(["pypandoc"])
|
||||
def convert_file_to_text(filename: str, source_format: str, target_format: str) -> str:
|
||||
"""Uses pandoc to convert the source document to a raw text string."""
|
||||
import pypandoc
|
||||
|
||||
try:
|
||||
text: str = pypandoc.convert_file(
|
||||
filename, target_format, format=source_format, sandbox=True
|
||||
)
|
||||
except FileNotFoundError as err:
|
||||
msg = (
|
||||
f"Error converting the file to text. Ensure you have the pandoc package installed on"
|
||||
f" your system. Installation instructions are available at"
|
||||
f" https://pandoc.org/installing.html. The original exception text was:\n{err}"
|
||||
)
|
||||
raise FileNotFoundError(msg)
|
||||
except RuntimeError as err:
|
||||
err_str = str(err)
|
||||
if source_format == "epub" and (
|
||||
"Couldn't extract ePub file" in err_str
|
||||
or "No entry on path" in err_str
|
||||
or re.search(r"exitcode ['\"]?64['\"]?", err_str)
|
||||
):
|
||||
raise UnprocessableEntityError(f"Invalid EPUB file: {err_str}")
|
||||
|
||||
supported_source_formats, _ = pypandoc.get_pandoc_formats()
|
||||
|
||||
if source_format == "rtf" and source_format not in supported_source_formats:
|
||||
additional_info = (
|
||||
"Support for RTF files is not available in the current pandoc installation. "
|
||||
"It was introduced in pandoc 2.14.2.\n"
|
||||
"Reference: https://pandoc.org/releases.html#pandoc-2.14.2-2021-08-21"
|
||||
)
|
||||
else:
|
||||
additional_info = ""
|
||||
|
||||
msg = (
|
||||
f"{err}\n\n{additional_info}\n\n"
|
||||
f"Current version of pandoc: {pypandoc.get_pandoc_version()}\n"
|
||||
"Make sure you have the right version installed in your system. Please follow the"
|
||||
" pandoc installation instructions in README.md to install the right version."
|
||||
)
|
||||
raise RuntimeError(msg)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def convert_file_to_html_text_using_pandoc(
|
||||
source_format: str, filename: str | None = None, file: IO[bytes] | None = None
|
||||
) -> str:
|
||||
"""Converts a document to HTML raw text.
|
||||
|
||||
Enables the doucment to be processed using `partition_html()`.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file)
|
||||
|
||||
if file is not None:
|
||||
with tempfile.TemporaryDirectory() as temp_dir_path:
|
||||
tmp_file_path = os.path.join(temp_dir_path, f"tmp_file.{source_format}")
|
||||
with open(tmp_file_path, "wb") as tmp_file:
|
||||
tmp_file.write(file.read())
|
||||
return convert_file_to_text(
|
||||
filename=tmp_file_path, source_format=source_format, target_format="html"
|
||||
)
|
||||
|
||||
assert filename is not None
|
||||
return convert_file_to_text(
|
||||
filename=filename, source_format=source_format, target_format="html"
|
||||
)
|
||||
@@ -0,0 +1,834 @@
|
||||
"""Automatically detect file-type based on inspection of the file's contents.
|
||||
|
||||
Auto-detection proceeds via a sequence of strategies. The first strategy to confidently determine a
|
||||
file-type returns that value. A strategy that is not applicable, either because it lacks the input
|
||||
required or fails to determine a file-type, returns `None` and execution continues with the next
|
||||
strategy.
|
||||
|
||||
`_FileTypeDetector` is the main object and implements the three strategies.
|
||||
|
||||
The three strategies are:
|
||||
|
||||
- Use MIME-type asserted by caller in the `content_type` argument.
|
||||
- Guess a MIME-type using libmagic, falling back to the `filetype` package when libmagic is
|
||||
unavailable.
|
||||
- Map filename-extension to a `FileType` member.
|
||||
|
||||
A file that fails all three strategies is assigned the value `FileType.UNK`, for "unknown".
|
||||
|
||||
`_FileTypeDetectionContext` encapsulates the various arguments received by `detect_filetype()` and
|
||||
provides values derived from them. This object is immutable and can be passed to delegates of
|
||||
`_FileTypeDetector` to provide whatever context they need on the current detection instance.
|
||||
|
||||
`_FileTypeDetector` delegates to _differentiator_ objects like `_ZipFileDifferentiator` for
|
||||
specialized discrimination and/or confirmation of ambiguous or frequently mis-identified
|
||||
MIME-types. Additional differentiators are planned, one for `application/x-ole-storage`
|
||||
(DOC, PPT, XLS, and MSG file-types) and perhaps others.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
import functools
|
||||
import importlib.util
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import zipfile
|
||||
from typing import IO, Callable, Iterator, Optional
|
||||
|
||||
import filetype as ft
|
||||
from olefile import OleFileIO
|
||||
from oxmsg.storage import Storage
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
from unstructured.documents.elements import Element
|
||||
from unstructured.file_utils.encoding import detect_file_encoding, format_encoding_str
|
||||
from unstructured.file_utils.model import FileType
|
||||
from unstructured.logger import logger
|
||||
from unstructured.nlp.patterns import EMAIL_HEAD_RE, LIST_OF_DICTS_PATTERN
|
||||
from unstructured.partition.common.common import add_element_metadata, exactly_one
|
||||
from unstructured.partition.common.metadata import set_element_hierarchy
|
||||
from unstructured.utils import get_call_args_applying_defaults, lazyproperty
|
||||
|
||||
try:
|
||||
importlib.import_module("magic")
|
||||
LIBMAGIC_AVAILABLE = True
|
||||
except ImportError:
|
||||
LIBMAGIC_AVAILABLE = False # pyright: ignore[reportConstantRedefinition]
|
||||
|
||||
|
||||
def detect_filetype(
|
||||
file_path: str | None = None,
|
||||
file: IO[bytes] | tempfile.SpooledTemporaryFile | None = None,
|
||||
encoding: str | None = None,
|
||||
content_type: str | None = None,
|
||||
metadata_file_path: Optional[str] = None,
|
||||
) -> FileType:
|
||||
"""Determine file-type of specified file using libmagic and/or fallback methods.
|
||||
|
||||
One of `file_path` or `file` must be specified. A `file_path` that does not
|
||||
correspond to a file on the filesystem raises `ValueError`.
|
||||
|
||||
Args:
|
||||
content_type: MIME-type of document-source, when already known. Providing
|
||||
a value for this argument disables auto-detection unless it does not map
|
||||
to a FileType member or is ambiguous, in which case it is ignored.
|
||||
encoding: Only used for textual file-types. When omitted, `utf-8` is
|
||||
assumed. Should generally be omitted except to resolve a problem with
|
||||
textual file-types like HTML.
|
||||
metadata_file_path: Only used when `file` is provided and then only as a
|
||||
source for a filename-extension that may be needed as a secondary
|
||||
content-type indicator. Ignored with the document is specified using
|
||||
`file_path`.
|
||||
|
||||
Returns:
|
||||
A member of the `FileType` enumeration, `FileType.UNK` when the file type
|
||||
could not be determined or is not supported.
|
||||
|
||||
Raises:
|
||||
ValueError: when:
|
||||
- `file_path` is specified but does not correspond to a file on the
|
||||
filesystem.
|
||||
- Neither `file_path` nor `file` were specified.
|
||||
"""
|
||||
file_buffer = file
|
||||
if isinstance(file, tempfile.SpooledTemporaryFile):
|
||||
file_buffer = io.BytesIO(file.read())
|
||||
file.seek(0)
|
||||
|
||||
ctx = _FileTypeDetectionContext.new(
|
||||
file_path=file_path,
|
||||
file=file_buffer,
|
||||
encoding=encoding,
|
||||
content_type=content_type,
|
||||
metadata_file_path=metadata_file_path,
|
||||
)
|
||||
return _FileTypeDetector.file_type(ctx)
|
||||
|
||||
|
||||
def is_json_processable(
|
||||
filename: Optional[str] = None,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
file_text: Optional[str] = None,
|
||||
encoding: Optional[str] = "utf-8",
|
||||
) -> bool:
|
||||
"""True when file looks like a JSON array of objects.
|
||||
|
||||
Uses regex on a file prefix, so not entirely reliable but good enough if you already know the
|
||||
file is JSON.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file, file_text=file_text)
|
||||
|
||||
if file_text is None:
|
||||
file_text = _FileTypeDetectionContext.new(
|
||||
file_path=filename, file=file, encoding=encoding
|
||||
).text_head
|
||||
|
||||
return re.match(LIST_OF_DICTS_PATTERN, file_text) is not None
|
||||
|
||||
|
||||
def is_ndjson_processable(
|
||||
filename: Optional[str] = None,
|
||||
file: Optional[IO[bytes]] = None,
|
||||
file_text: Optional[str] = None,
|
||||
encoding: Optional[str] = "utf-8",
|
||||
) -> bool:
|
||||
"""True when file looks like a JSON array of objects.
|
||||
|
||||
Uses regex on a file prefix, so not entirely reliable but good enough if you already know the
|
||||
file is JSON.
|
||||
"""
|
||||
exactly_one(filename=filename, file=file, file_text=file_text)
|
||||
|
||||
if file_text is None:
|
||||
file_text = _FileTypeDetectionContext.new(
|
||||
file_path=filename, file=file, encoding=encoding
|
||||
).text_head
|
||||
return file_text.lstrip().startswith("{")
|
||||
|
||||
|
||||
class _FileTypeDetector:
|
||||
"""Determines file type from a variety of possible inputs."""
|
||||
|
||||
def __init__(self, ctx: _FileTypeDetectionContext):
|
||||
self._ctx = ctx
|
||||
|
||||
@classmethod
|
||||
def file_type(cls, ctx: _FileTypeDetectionContext) -> FileType:
|
||||
"""Detect file-type of document-source described by `ctx`."""
|
||||
return cls(ctx)._file_type
|
||||
|
||||
@property
|
||||
def _file_type(self) -> FileType:
|
||||
"""FileType member corresponding to this document source."""
|
||||
# -- An explicit content-type most commonly asserted by the client/SDK and is therefore
|
||||
# -- inherently unreliable. On the other hand, binary file-types can be detected with 100%
|
||||
# -- accuracy. So start with binary types and only then consider an asserted content-type,
|
||||
# -- generally as a last resort.
|
||||
|
||||
if (
|
||||
( # strategy 1: most binary types can be detected with 100% accuracy
|
||||
predicted_file_type := self._known_binary_file_type
|
||||
)
|
||||
or ( # strategy 2: use content-type asserted by caller
|
||||
predicted_file_type := self._file_type_from_content_type
|
||||
)
|
||||
or ( # strategy 3: guess MIME-type using libmagic and use that
|
||||
predicted_file_type := self._file_type_from_guessed_mime_type
|
||||
)
|
||||
or ( # strategy 4: use filename-extension, like ".docx" -> FileType.DOCX
|
||||
predicted_file_type := self._file_type_from_file_extension
|
||||
)
|
||||
):
|
||||
result_file_type = predicted_file_type
|
||||
else:
|
||||
# give up and report FileType.UNK
|
||||
result_file_type = FileType.UNK
|
||||
|
||||
if result_file_type == FileType.JSON:
|
||||
# edge case where JSON/NDJSON content without file extension
|
||||
# (magic lib can't distinguish them)
|
||||
result_file_type = self._disambiguate_json_file_type
|
||||
|
||||
return result_file_type
|
||||
|
||||
@property
|
||||
def _known_binary_file_type(self) -> FileType | None:
|
||||
"""Detect file-type for binary types we can positively detect."""
|
||||
if file_type := _OleFileDetector.file_type(self._ctx):
|
||||
return file_type
|
||||
|
||||
self._ctx.rule_out_cfb_content_types()
|
||||
|
||||
if file_type := _ZipFileDetector.file_type(self._ctx):
|
||||
return file_type
|
||||
|
||||
self._ctx.rule_out_zip_content_types()
|
||||
|
||||
return None
|
||||
|
||||
@property
|
||||
def _file_type_from_content_type(self) -> FileType | None:
|
||||
"""Map passed content-type argument to a file-type, subject to certain rules."""
|
||||
|
||||
# -- when no content-type was asserted by caller, this strategy is not applicable --
|
||||
if not self._ctx.content_type:
|
||||
return None
|
||||
|
||||
# -- otherwise we trust the passed `content_type` as long as `FileType` recognizes it --
|
||||
return FileType.from_mime_type(self._ctx.content_type)
|
||||
|
||||
@property
|
||||
def _disambiguate_json_file_type(self) -> FileType:
|
||||
"""Disambiguate JSON/NDJSON file-type based on file contents."""
|
||||
if is_json_processable(file_text=self._ctx.text_head):
|
||||
return FileType.JSON
|
||||
if is_ndjson_processable(file_text=self._ctx.text_head):
|
||||
return FileType.NDJSON
|
||||
raise ValueError("Unable to process JSON file")
|
||||
|
||||
@property
|
||||
def _file_type_from_guessed_mime_type(self) -> FileType | None:
|
||||
"""FileType based on auto-detection of MIME-type by libmagic.
|
||||
|
||||
In some cases refinements are necessary on the magic-derived MIME-types. This process
|
||||
includes applying those rules, most of which are accumulated through practical experience.
|
||||
"""
|
||||
mime_type = self._ctx.mime_type
|
||||
extension = self._ctx.extension
|
||||
|
||||
# -- when libmagic is not installed, the `filetype` package is used instead.
|
||||
# -- `filetype.guess()` returns `None` for file-types it does not support, which
|
||||
# -- unfortunately includes all the textual file-types like CSV, EML, HTML, MD, RST, RTF,
|
||||
# -- TSV, and TXT. When we have no guessed MIME-type, this strategy is not applicable.
|
||||
if mime_type is None:
|
||||
return None
|
||||
|
||||
if mime_type.endswith("xml"):
|
||||
return FileType.HTML if extension in (".html", ".htm") else FileType.XML
|
||||
|
||||
if differentiator := _TextFileDifferentiator.applies(self._ctx):
|
||||
return differentiator.file_type
|
||||
|
||||
# -- All source-code files (e.g. *.py, *.js) are classified as plain text for the moment --
|
||||
if self._ctx.has_code_mime_type:
|
||||
return FileType.TXT
|
||||
|
||||
if mime_type.endswith("empty"):
|
||||
return FileType.EMPTY
|
||||
|
||||
if mime_type.endswith("json") and self._ctx.extension == ".ndjson":
|
||||
return FileType.NDJSON
|
||||
|
||||
# -- if no more-specific rules apply, use the MIME-type -> FileType mapping when present --
|
||||
file_type = FileType.from_mime_type(mime_type)
|
||||
return file_type if file_type != FileType.UNK else None
|
||||
|
||||
@lazyproperty
|
||||
def _file_type_from_file_extension(self) -> FileType | None:
|
||||
"""Determine file-type from filename extension.
|
||||
|
||||
Returns `None` when no filename is available or when the extension does not map to a
|
||||
supported file-type.
|
||||
"""
|
||||
return FileType.from_extension(self._ctx.extension)
|
||||
|
||||
|
||||
class _FileTypeDetectionContext:
|
||||
"""Provides all arguments to auto-file detection and values derived from them.
|
||||
|
||||
NOTE that `._content_type` is mutable via `.rule_out_*_content_types()` methods, so it should
|
||||
not be assumed to be a constant value across those calls.
|
||||
|
||||
This keeps computation of derived values out of the file-detection code but more importantly
|
||||
allows the main filetype-detector to pass the full context to any delegates without coupling
|
||||
itself to which values it might need.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
file_path: str | None = None,
|
||||
*,
|
||||
file: IO[bytes] | None = None,
|
||||
encoding: str | None = None,
|
||||
content_type: str | None = None,
|
||||
metadata_file_path: str | None = None,
|
||||
):
|
||||
self._file_path_arg = file_path
|
||||
self._file_arg = file
|
||||
self._encoding_arg = encoding
|
||||
self._content_type = content_type
|
||||
self._metadata_file_path = metadata_file_path
|
||||
|
||||
@classmethod
|
||||
def new(
|
||||
cls,
|
||||
*,
|
||||
file_path: str | None,
|
||||
file: IO[bytes] | None,
|
||||
encoding: str | None,
|
||||
content_type: str | None = None,
|
||||
metadata_file_path: str | None = None,
|
||||
) -> _FileTypeDetectionContext:
|
||||
self = cls(
|
||||
file_path=file_path,
|
||||
file=file,
|
||||
encoding=encoding,
|
||||
content_type=content_type,
|
||||
metadata_file_path=metadata_file_path,
|
||||
)
|
||||
self._validate()
|
||||
return self
|
||||
|
||||
@property
|
||||
def content_type(self) -> str | None:
|
||||
"""MIME-type asserted by caller; not based on inspection of file by this process.
|
||||
|
||||
Would commonly occur when the file was downloaded via HTTP and a `"Content-Type:` header was
|
||||
present on the response. These are often ambiguous and sometimes just wrong so get some
|
||||
further verification. All lower-case when not `None`.
|
||||
"""
|
||||
# -- Note `._content_type` is mutable via `.invalidate_content_type()` so this cannot be a
|
||||
# -- `@lazyproperty`.
|
||||
return self._content_type.lower() if self._content_type else None
|
||||
|
||||
@lazyproperty
|
||||
def encoding(self) -> str:
|
||||
"""Character-set used to encode text of this file.
|
||||
|
||||
Relevant for textual file-types only, like HTML, TXT, JSON, etc.
|
||||
"""
|
||||
return format_encoding_str(self._encoding_arg or "utf-8")
|
||||
|
||||
@lazyproperty
|
||||
def extension(self) -> str:
|
||||
"""Best filename-extension we can muster, "" when there is no available source."""
|
||||
# -- get from file_path, or file when it has a name (path) --
|
||||
with self.open() as file:
|
||||
if hasattr(file, "name") and file.name:
|
||||
return os.path.splitext(file.name)[1].lower()
|
||||
|
||||
# -- otherwise use metadata file-path when provided --
|
||||
if file_path := self._metadata_file_path:
|
||||
return os.path.splitext(file_path)[1].lower()
|
||||
|
||||
# -- otherwise empty str means no extension, same as a path like "a/b/name-no-ext" --
|
||||
return ""
|
||||
|
||||
@lazyproperty
|
||||
def file_head(self) -> bytes:
|
||||
"""The initial bytes of the file to be recognized, for use with libmagic detection."""
|
||||
with self.open() as file:
|
||||
return file.read(8192)
|
||||
|
||||
@lazyproperty
|
||||
def file_path(self) -> str | None:
|
||||
"""Filesystem path to file to be inspected, when provided on call.
|
||||
|
||||
None when the caller specified the source as a file-like object instead. Useful for user
|
||||
feedback on an error, but users of context should have little use for it otherwise.
|
||||
"""
|
||||
if (file_path := self._file_path_arg) is None:
|
||||
return None
|
||||
|
||||
return os.path.realpath(file_path) if os.path.islink(file_path) else file_path
|
||||
|
||||
@lazyproperty
|
||||
def has_code_mime_type(self) -> bool:
|
||||
"""True when `mime_type` plausibly indicates a programming language source-code file."""
|
||||
mime_type = self.mime_type
|
||||
|
||||
if mime_type is None:
|
||||
return False
|
||||
|
||||
# -- check Go separately to avoid matching other MIME type containing "go" --
|
||||
if mime_type == "text/x-go":
|
||||
return True
|
||||
|
||||
return any(
|
||||
lang in mime_type
|
||||
for lang in [
|
||||
"c#",
|
||||
"c++",
|
||||
"cpp",
|
||||
"csharp",
|
||||
"java",
|
||||
"javascript",
|
||||
"php",
|
||||
"python",
|
||||
"ruby",
|
||||
"swift",
|
||||
"typescript",
|
||||
]
|
||||
)
|
||||
|
||||
@lazyproperty
|
||||
def is_zipfile(self) -> bool:
|
||||
"""True when file is a Zip archive."""
|
||||
with self.open() as file:
|
||||
return zipfile.is_zipfile(file)
|
||||
|
||||
@lazyproperty
|
||||
def mime_type(self) -> str | None:
|
||||
"""The best MIME-type we can get from `magic` (or `filetype` package).
|
||||
|
||||
A `str` return value is always in lower-case.
|
||||
"""
|
||||
file_path = self.file_path
|
||||
|
||||
if LIBMAGIC_AVAILABLE:
|
||||
import magic
|
||||
|
||||
mime_type = (
|
||||
magic.from_file(file_path, mime=True)
|
||||
if file_path
|
||||
else magic.from_buffer(self.file_head, mime=True)
|
||||
)
|
||||
return mime_type.lower() if mime_type else None
|
||||
|
||||
mime_type = ft.guess_mime(file_path) if file_path else ft.guess_mime(self.file_head)
|
||||
|
||||
if mime_type is None:
|
||||
logger.warning(
|
||||
"libmagic is unavailable but assists in filetype detection. Please consider"
|
||||
" installing libmagic for better results."
|
||||
)
|
||||
return None
|
||||
|
||||
return mime_type.lower()
|
||||
|
||||
@contextlib.contextmanager
|
||||
def open(self) -> Iterator[IO[bytes]]:
|
||||
"""Encapsulates complexity of dealing with file-path or file-like-object.
|
||||
|
||||
Provides an `IO[bytes]` object as the "common-denominator" document source.
|
||||
|
||||
Must be used as a context manager using a `with` statement:
|
||||
|
||||
with self._file as file:
|
||||
do things with file
|
||||
|
||||
File is guaranteed to be at read position 0 when called.
|
||||
"""
|
||||
if self.file_path:
|
||||
with open(self.file_path, "rb") as f:
|
||||
yield f
|
||||
else:
|
||||
file = self._file_arg
|
||||
assert file is not None # -- guaranteed by `._validate()` --
|
||||
file.seek(0)
|
||||
yield file
|
||||
|
||||
def rule_out_cfb_content_types(self) -> None:
|
||||
"""Invalidate content-type when a legacy MS-Office file-type is asserted.
|
||||
|
||||
Used before returning `None`; at that point we know the file is not one of these formats
|
||||
so if the asserted `content-type` is a legacy MS-Office type we know it's wrong and should
|
||||
not be used as a fallback later in the detection process.
|
||||
"""
|
||||
if FileType.from_mime_type(self._content_type) in (
|
||||
FileType.DOC,
|
||||
FileType.MSG,
|
||||
FileType.PPT,
|
||||
FileType.XLS,
|
||||
):
|
||||
self._content_type = None
|
||||
|
||||
def rule_out_zip_content_types(self) -> None:
|
||||
"""Invalidate content-type when an MS-Office 2007+ file-type is asserted.
|
||||
|
||||
Used before returning `None`; at that point we know the file is not one of these formats
|
||||
so if the asserted `content-type` is an MS-Office 2007+ type we know it's wrong and should
|
||||
not be used as a fallback later in the detection process.
|
||||
"""
|
||||
if FileType.from_mime_type(self._content_type) in (
|
||||
FileType.DOCX,
|
||||
FileType.EPUB,
|
||||
FileType.ODT,
|
||||
FileType.PPTX,
|
||||
FileType.XLSX,
|
||||
FileType.ZIP,
|
||||
):
|
||||
self._content_type = None
|
||||
|
||||
@lazyproperty
|
||||
def text_head(self) -> str:
|
||||
"""The initial characters of the text file for use with text-format differentiation.
|
||||
|
||||
Raises:
|
||||
UnicodeDecodeError if file cannot be read as text.
|
||||
"""
|
||||
# TODO: only attempts fallback character-set detection for file-path case, not for
|
||||
# file-like object case. Seems like we should do both.
|
||||
|
||||
if file := self._file_arg:
|
||||
file.seek(0)
|
||||
content = file.read(4096)
|
||||
file.seek(0)
|
||||
return (
|
||||
content
|
||||
if isinstance(content, str)
|
||||
else content.decode(encoding=self.encoding, errors="ignore")
|
||||
)
|
||||
|
||||
file_path = self.file_path
|
||||
assert file_path is not None # -- guaranteed by `._validate` --
|
||||
|
||||
try:
|
||||
with open(file_path, encoding=self.encoding) as f:
|
||||
return f.read(4096)
|
||||
except UnicodeDecodeError:
|
||||
encoding, _ = detect_file_encoding(filename=file_path)
|
||||
with open(file_path, encoding=encoding) as f:
|
||||
return f.read(4096)
|
||||
|
||||
def _validate(self) -> None:
|
||||
"""Raise if the context is invalid."""
|
||||
if self.file_path and not os.path.isfile(self.file_path):
|
||||
raise FileNotFoundError(f"no such file {self._file_path_arg}")
|
||||
if not self.file_path and not self._file_arg:
|
||||
raise ValueError("either `file_path` or `file` argument must be provided")
|
||||
|
||||
|
||||
class _OleFileDetector:
|
||||
"""Detect and differentiate a CFB file, aka. "OLE" file.
|
||||
|
||||
Compound File Binary Format (CFB), aka. OLE file, is use by Microsoft for legacy MS Office
|
||||
files (DOC, PPT, XLS) as well as for Outlook MSG files.
|
||||
"""
|
||||
|
||||
def __init__(self, ctx: _FileTypeDetectionContext):
|
||||
self._ctx = ctx
|
||||
|
||||
@classmethod
|
||||
def file_type(cls, ctx: _FileTypeDetectionContext) -> FileType | None:
|
||||
"""Specific file-type when file is a CFB file, `None` otherwise."""
|
||||
return cls(ctx)._file_type
|
||||
|
||||
@property
|
||||
def _file_type(self) -> FileType | None:
|
||||
"""Differentiated file-type for Microsoft Compound File Binary Format (CFBF).
|
||||
|
||||
Returns one of:
|
||||
- `FileType.DOC`
|
||||
- `FileType.PPT`
|
||||
- `FileType.XLS`
|
||||
- `FileType.MSG`
|
||||
- `None` when the file is not one of these.
|
||||
"""
|
||||
# -- all CFB files share common magic number, start with that --
|
||||
if not self._is_ole_file:
|
||||
return None
|
||||
|
||||
# -- check storage contents of the ole file for file-type specific stream names --
|
||||
if (ole_file_type := self._ole_file_type) is not None:
|
||||
return ole_file_type
|
||||
|
||||
return None
|
||||
|
||||
@lazyproperty
|
||||
def _is_ole_file(self) -> bool:
|
||||
"""True when file has CFB magic first 8 bytes."""
|
||||
with self._ctx.open() as file:
|
||||
return file.read(8) == b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"
|
||||
|
||||
@lazyproperty
|
||||
def _ole_file_type(self) -> FileType | None:
|
||||
with self._ctx.open() as f:
|
||||
ole = OleFileIO(f) # pyright: ignore[reportUnknownVariableType]
|
||||
root_storage = Storage.from_ole(ole) # pyright: ignore[reportUnknownMemberType]
|
||||
|
||||
for stream in root_storage.streams:
|
||||
if stream.name == "WordDocument":
|
||||
return FileType.DOC
|
||||
elif stream.name == "PowerPoint Document":
|
||||
return FileType.PPT
|
||||
elif stream.name == "Workbook":
|
||||
return FileType.XLS
|
||||
elif stream.name == "__properties_version1.0":
|
||||
return FileType.MSG
|
||||
|
||||
return None
|
||||
|
||||
|
||||
class _TextFileDifferentiator:
|
||||
"""Refine a textual file-type that may not be as specific as it could be."""
|
||||
|
||||
def __init__(self, ctx: _FileTypeDetectionContext):
|
||||
self._ctx = ctx
|
||||
|
||||
@classmethod
|
||||
def applies(cls, ctx: _FileTypeDetectionContext) -> _TextFileDifferentiator | None:
|
||||
"""Constructs an instance, but only if this differentiator applies in `ctx`."""
|
||||
mime_type = ctx.mime_type
|
||||
return (
|
||||
cls(ctx)
|
||||
if mime_type and (mime_type == "message/rfc822" or mime_type.startswith("text"))
|
||||
else None
|
||||
)
|
||||
|
||||
@lazyproperty
|
||||
def file_type(self) -> FileType:
|
||||
"""Differentiated file-type for textual content.
|
||||
|
||||
Always produces a file-type, worst case that's `FileType.TXT` when nothing more specific
|
||||
applies.
|
||||
"""
|
||||
extension = self._ctx.extension
|
||||
|
||||
if extension in [
|
||||
".csv",
|
||||
".eml",
|
||||
".html",
|
||||
".json",
|
||||
".markdown",
|
||||
".md",
|
||||
".org",
|
||||
".p7s",
|
||||
".rst",
|
||||
".rtf",
|
||||
".tab",
|
||||
".tsv",
|
||||
]:
|
||||
return FileType.from_extension(extension) or FileType.TXT
|
||||
|
||||
# NOTE(crag): for older versions of the OS libmagic package, such as is currently
|
||||
# installed on the Unstructured docker image, .json files resolve to "text/plain"
|
||||
# rather than "application/json". this corrects for that case.
|
||||
if self._is_json:
|
||||
return FileType.JSON
|
||||
|
||||
if self._is_csv:
|
||||
return FileType.CSV
|
||||
|
||||
if self._is_eml:
|
||||
return FileType.EML
|
||||
|
||||
if extension in (".text", ".txt"):
|
||||
return FileType.TXT
|
||||
|
||||
# Safety catch
|
||||
if file_type := FileType.from_mime_type(self._ctx.mime_type):
|
||||
return file_type
|
||||
|
||||
return FileType.TXT
|
||||
|
||||
@lazyproperty
|
||||
def _is_csv(self) -> bool:
|
||||
"""True when file is plausibly in Comma Separated Values (CSV) format."""
|
||||
|
||||
def count_commas(text: str):
|
||||
"""Counts the number of commas in a line, excluding commas in quotes."""
|
||||
pattern = r"(?=(?:[^\"]*\"[^\"]*\")*[^\"]*$),"
|
||||
matches = re.findall(pattern, text)
|
||||
return len(matches)
|
||||
|
||||
lines = self._ctx.text_head.strip().splitlines()
|
||||
if len(lines) < 2:
|
||||
return False
|
||||
# -- check at most the first 10 lines --
|
||||
lines = lines[: len(lines)] if len(lines) < 10 else lines[:10]
|
||||
# -- any lines without at least one comma disqualifies the file --
|
||||
if any("," not in line for line in lines):
|
||||
return False
|
||||
header_count = count_commas(lines[0])
|
||||
return all(count_commas(line) == header_count for line in lines[1:])
|
||||
|
||||
@lazyproperty
|
||||
def _is_eml(self) -> bool:
|
||||
"""Checks if a text/plain file is actually a .eml file.
|
||||
|
||||
Uses a regex pattern to see if the start of the file matches the typical pattern for a .eml
|
||||
file.
|
||||
"""
|
||||
return EMAIL_HEAD_RE.match(self._ctx.text_head) is not None
|
||||
|
||||
@lazyproperty
|
||||
def _is_json(self) -> bool:
|
||||
"""True when file is JSON collection.
|
||||
|
||||
A JSON file that contains only a string, number, or boolean, while valid JSON, will fail
|
||||
this test since it is not partitionable.
|
||||
"""
|
||||
text_head = self._ctx.text_head
|
||||
|
||||
# -- an empty file is not JSON --
|
||||
if not text_head.lstrip():
|
||||
return False
|
||||
|
||||
# -- has to be a list or object, no string, number, or bool --
|
||||
if text_head.lstrip()[0] not in "[{":
|
||||
return False
|
||||
|
||||
try:
|
||||
with self._ctx.open() as file:
|
||||
json.load(file)
|
||||
return True
|
||||
except json.JSONDecodeError:
|
||||
return False
|
||||
|
||||
|
||||
class _ZipFileDetector:
|
||||
"""Detect and differentiate a Zip-archive file."""
|
||||
|
||||
def __init__(self, ctx: _FileTypeDetectionContext):
|
||||
self._ctx = ctx
|
||||
|
||||
@classmethod
|
||||
def file_type(cls, ctx: _FileTypeDetectionContext) -> FileType | None:
|
||||
"""Most specific file-type available when file is a Zip file, `None` otherwise.
|
||||
|
||||
MS-Office 2007+ files are detected with 100% accuracy. Otherwise this returns `None`, even
|
||||
when we can tell it's a Zip file, so later strategies can have a crack at it. In
|
||||
particular, ODT and EPUB files are Zip archives but are not detected here.
|
||||
"""
|
||||
return cls(ctx)._file_type
|
||||
|
||||
@lazyproperty
|
||||
def _file_type(self) -> FileType | None:
|
||||
"""Differentiated file-type for a Zip archive.
|
||||
|
||||
Returns `FileType.DOCX`, `FileType.PPTX`, or `FileType.XLSX` when one of those applies,
|
||||
`None` otherwise.
|
||||
"""
|
||||
if not self._ctx.is_zipfile:
|
||||
return None
|
||||
|
||||
with self._ctx.open() as file:
|
||||
zip = zipfile.ZipFile(file)
|
||||
|
||||
filenames = zip.namelist()
|
||||
|
||||
if any(re.match(r"word/document.*\.xml$", filename) for filename in filenames):
|
||||
return FileType.DOCX
|
||||
|
||||
if any(re.match(r"xl/workbook.*\.xml$", filename) for filename in filenames):
|
||||
return FileType.XLSX
|
||||
|
||||
if any(re.match(r"ppt/presentation.*\.xml$", filename) for filename in filenames):
|
||||
return FileType.PPTX
|
||||
|
||||
# -- ODT and EPUB files place their MIME-type in `mimetype` in the archive root --
|
||||
if "mimetype" in filenames:
|
||||
with zip.open("mimetype") as f:
|
||||
mime_type = f.read().decode("utf-8").strip()
|
||||
return FileType.from_mime_type(mime_type)
|
||||
|
||||
return FileType.ZIP
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
|
||||
def add_metadata(func: Callable[_P, list[Element]]) -> Callable[_P, list[Element]]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> list[Element]:
|
||||
elements = func(*args, **kwargs)
|
||||
call_args = get_call_args_applying_defaults(func, *args, **kwargs)
|
||||
|
||||
if call_args.get("metadata_filename"):
|
||||
call_args["filename"] = call_args.get("metadata_filename")
|
||||
|
||||
metadata_kwargs = {
|
||||
kwarg: call_args.get(kwarg) for kwarg in ("filename", "url", "text_as_html")
|
||||
}
|
||||
# NOTE (yao): do not use cast here as cast(None) still is None
|
||||
if not str(kwargs.get("model_name", "")).startswith("chipper"):
|
||||
# NOTE(alan): Skip hierarchy if using chipper, as it should take care of that
|
||||
elements = set_element_hierarchy(elements)
|
||||
|
||||
for element in elements:
|
||||
# NOTE(robinson) - Attached files have already run through this logic
|
||||
# in their own partitioning function
|
||||
if element.metadata.attached_to_filename is None:
|
||||
add_element_metadata(element, **metadata_kwargs)
|
||||
|
||||
return elements
|
||||
|
||||
return wrapper
|
||||
|
||||
|
||||
def add_filetype(
|
||||
filetype: FileType,
|
||||
) -> Callable[[Callable[_P, list[Element]]], Callable[_P, list[Element]]]:
|
||||
"""Post-process element-metadata for list[Element] from partitioning.
|
||||
|
||||
This decorator adds a post-processing step to a document partitioner.
|
||||
|
||||
- Adds `.metadata.filetype` (source-document MIME-type) metadata value
|
||||
|
||||
This "partial" decorator is present because `partition_image()` does not apply
|
||||
`.metadata.filetype` this way since each image type has its own MIME-type (e.g. `image.jpeg`,
|
||||
`image/png`, etc.).
|
||||
"""
|
||||
|
||||
def decorator(func: Callable[_P, list[Element]]) -> Callable[_P, list[Element]]:
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> list[Element]:
|
||||
elements = func(*args, **kwargs)
|
||||
|
||||
for element in elements:
|
||||
# NOTE(robinson) - Attached files have already run through this logic
|
||||
# in their own partitioning function
|
||||
if element.metadata.attached_to_filename is None:
|
||||
add_element_metadata(element, filetype=filetype.mime_type)
|
||||
|
||||
return elements
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def add_metadata_with_filetype(
|
||||
filetype: FileType,
|
||||
) -> Callable[[Callable[_P, list[Element]]], Callable[_P, list[Element]]]:
|
||||
"""..."""
|
||||
|
||||
def decorator(func: Callable[_P, list[Element]]) -> Callable[_P, list[Element]]:
|
||||
return add_filetype(filetype=filetype)(add_metadata(func))
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,9 @@
|
||||
GOOGLE_DRIVE_EXPORT_TYPES = {
|
||||
"application/vnd.google-apps.document": "application/"
|
||||
"vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
"application/vnd.google-apps.spreadsheet": "application/"
|
||||
"vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
"application/vnd.google-apps.presentation": "application/"
|
||||
"vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
"application/vnd.google-apps.photo": "image/jpeg",
|
||||
}
|
||||
@@ -0,0 +1,534 @@
|
||||
"""Domain-model for file-types."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import enum
|
||||
from typing import TYPE_CHECKING, Callable, Iterable, Type, cast
|
||||
|
||||
from typing_extensions import ParamSpec
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from unstructured.documents.elements import Element
|
||||
else:
|
||||
Element = None
|
||||
|
||||
|
||||
def _create_file_type_enum(
|
||||
cls: Type["FileType"],
|
||||
value: str,
|
||||
partitioner_shortname: str | None,
|
||||
importable_package_dependencies: Iterable[str],
|
||||
extra_name: str | None,
|
||||
extensions: Iterable[str],
|
||||
canonical_mime_type: str,
|
||||
alias_mime_types: Iterable[str],
|
||||
partitioner_full_module_path: str | None = None,
|
||||
) -> "FileType":
|
||||
"""
|
||||
Moving here instead of directly in the FileType.__new__ allows us
|
||||
to dynamically create new enum properties.
|
||||
|
||||
FileType.__new__ does not work with dynamic properties.
|
||||
"""
|
||||
val = object.__new__(cls)
|
||||
val._value_ = value
|
||||
val._partitioner_shortname = partitioner_shortname
|
||||
val._importable_package_dependencies = tuple(importable_package_dependencies)
|
||||
val._extra_name = extra_name
|
||||
val._extensions = tuple(extensions)
|
||||
val._canonical_mime_type = canonical_mime_type
|
||||
val._alias_mime_types = tuple(alias_mime_types)
|
||||
val._partitioner_full_module_path = partitioner_full_module_path
|
||||
return val
|
||||
|
||||
|
||||
class FileType(enum.Enum):
|
||||
"""The collection of file-types recognized by `unstructured`.
|
||||
|
||||
Note not all of these can be partitioned, e.g. WAV and ZIP have no partitioner.
|
||||
"""
|
||||
|
||||
_partitioner_shortname: str | None
|
||||
"""Like "docx", from which partitioner module and function-name can be derived via template."""
|
||||
|
||||
_importable_package_dependencies: tuple[str, ...]
|
||||
"""Packages that must be available for import for this file-type's partitioner to work."""
|
||||
|
||||
_extra_name: str | None
|
||||
"""`pip install` extra that provides package dependencies for this file-type."""
|
||||
|
||||
_extensions: tuple[str, ...]
|
||||
"""Filename-extensions recognized as this file-type. Use for secondary identification only."""
|
||||
|
||||
_canonical_mime_type: str
|
||||
"""The MIME-type used as `.metadata.filetype` for this file-type."""
|
||||
|
||||
_alias_mime_types: tuple[str, ...]
|
||||
"""MIME-types accepted as identifying this file-type."""
|
||||
|
||||
_partitioner_full_module_path: str | None
|
||||
"""Fully-qualified name of module providing partitioner for this file-type."""
|
||||
|
||||
def __new__(
|
||||
cls,
|
||||
value: str,
|
||||
partitioner_shortname: str | None,
|
||||
importable_package_dependencies: Iterable[str],
|
||||
extra_name: str | None,
|
||||
extensions: Iterable[str],
|
||||
canonical_mime_type: str,
|
||||
alias_mime_types: Iterable[str],
|
||||
partitioner_full_module_path: str | None = None,
|
||||
):
|
||||
return _create_file_type_enum(
|
||||
cls,
|
||||
value,
|
||||
partitioner_shortname,
|
||||
importable_package_dependencies,
|
||||
extra_name,
|
||||
extensions,
|
||||
canonical_mime_type,
|
||||
alias_mime_types,
|
||||
partitioner_full_module_path,
|
||||
)
|
||||
|
||||
def __lt__(self, other: FileType) -> bool:
|
||||
"""Makes `FileType` members comparable with relational operators, at least with `<`.
|
||||
|
||||
This makes them sortable, in particular it supports sorting for pandas groupby functions.
|
||||
"""
|
||||
return self.name < other.name
|
||||
|
||||
@classmethod
|
||||
def from_extension(cls, extension: str | None) -> FileType | None:
|
||||
"""Select a FileType member based on an extension.
|
||||
|
||||
`extension` must include the leading period, like `".pdf"`. Extension is suitable as a
|
||||
secondary file-type identification method but is unreliable for primary identification.
|
||||
|
||||
Returns `None` when `extension` is not registered for any supported file-type.
|
||||
"""
|
||||
if extension in (None, "", "."):
|
||||
return None
|
||||
# -- not super efficient but plenty fast enough for once-or-twice-per-file use and avoids
|
||||
# -- limitations on defining a class variable on an Enum.
|
||||
for m in cls.__members__.values():
|
||||
if extension in m._extensions:
|
||||
return m
|
||||
return None
|
||||
|
||||
@classmethod
|
||||
def from_mime_type(cls, mime_type: str | None) -> FileType | None:
|
||||
"""Select a FileType member based on a MIME-type.
|
||||
|
||||
Returns `None` when `mime_type` is `None` or does not map to the canonical MIME-type of a
|
||||
`FileType` member or one of its alias MIME-types.
|
||||
"""
|
||||
if mime_type is None:
|
||||
return None
|
||||
# -- not super efficient but plenty fast enough for once-or-twice-per-file use and avoids
|
||||
# -- limitations on defining a class variable on an Enum.
|
||||
for m in cls.__members__.values():
|
||||
if mime_type == m._canonical_mime_type or mime_type in m._alias_mime_types:
|
||||
return m
|
||||
return None
|
||||
|
||||
@property
|
||||
def extra_name(self) -> str | None:
|
||||
"""The `pip` "extra" that must be installed to provide this file-type's dependencies.
|
||||
|
||||
Like "image" for PNG, as in `pip install "unstructured[image]"`.
|
||||
|
||||
`None` when partitioning this file-type requires only the base `unstructured` install.
|
||||
"""
|
||||
return self._extra_name
|
||||
|
||||
@property
|
||||
def importable_package_dependencies(self) -> tuple[str, ...]:
|
||||
"""Packages that must be importable for this file-type's partitioner to work.
|
||||
|
||||
In general, these are the packages provided by the `pip install` "extra" for this file-type,
|
||||
like `pip install "unstructured[docx]"` loads the `python-docx` package.
|
||||
|
||||
Note that these names are the ones used in an `import` statement, which is not necessarily
|
||||
the same as the _distribution_ package name used by `pip`. For example, the DOCX
|
||||
distribution package name is `"python-docx"` whereas the _importable_ package name is
|
||||
`"docx"`. This latter name as it appears like `import docx` is what is provided by this
|
||||
property.
|
||||
|
||||
The return value is an empty tuple for file-types that do not require optional dependencies.
|
||||
|
||||
Note this property does not complain when accessed on a non-partitionable file-type, it
|
||||
simply returns an empty tuple because file-types that are not partitionable require no
|
||||
optional dependencies.
|
||||
"""
|
||||
return self._importable_package_dependencies
|
||||
|
||||
@property
|
||||
def is_partitionable(self) -> bool:
|
||||
"""True when there is a partitioner for this file-type.
|
||||
|
||||
Note this does not check whether the dependencies for this file-type are installed so
|
||||
attempting to partition a file of this type may still fail. This is meant for
|
||||
distinguishing file-types like WAV, ZIP, EMPTY, and UNK which are legitimate file-types
|
||||
but have no associated partitioner.
|
||||
"""
|
||||
return bool(self._partitioner_shortname) or bool(self._partitioner_full_module_path)
|
||||
|
||||
@property
|
||||
def mime_type(self) -> str:
|
||||
"""The canonical MIME-type for this file-type, suitable for use in metadata.
|
||||
|
||||
This value is used in `.metadata.filetype` for elements partitioned from files of this
|
||||
type. In general it is the "offical", "recommended", or "defacto-standard" MIME-type for
|
||||
files of this type, in that order, as available.
|
||||
"""
|
||||
return self._canonical_mime_type
|
||||
|
||||
@property
|
||||
def partitioner_function_name(self) -> str:
|
||||
"""Name of partitioner function for this file-type. Like "partition_docx".
|
||||
|
||||
Raises when this property is accessed on a file-type that is not partitionable. Use
|
||||
`.is_partitionable` to avoid exceptions when partitionability is unknown.
|
||||
"""
|
||||
# -- Raise when this property is accessed on a FileType member that has no partitioner
|
||||
# -- shortname. This prevents a harder-to-find bug from appearing far away from this call
|
||||
# -- when code would try to `getattr(module, None)` or whatever.
|
||||
if full_module_path := self._partitioner_full_module_path:
|
||||
return full_module_path.split(".")[-1]
|
||||
|
||||
if (shortname := self._partitioner_shortname) is None:
|
||||
raise ValueError(
|
||||
f"`.partitioner_function_name` is undefined because FileType.{self.name} is not"
|
||||
f" partitionable. Use `.is_partitionable` to determine whether a `FileType`"
|
||||
f" is partitionable."
|
||||
)
|
||||
return f"partition_{shortname}"
|
||||
|
||||
@property
|
||||
def partitioner_module_qname(self) -> str:
|
||||
"""Fully-qualified name of module providing partitioner for this file-type.
|
||||
|
||||
e.g. "unstructured.partition.docx" for FileType.DOCX.
|
||||
"""
|
||||
# -- Raise when this property is accessed on a FileType member that has no partitioner
|
||||
# -- shortname. This prevents a harder-to-find bug from appearing far away from this call
|
||||
# -- when code would try to `importlib.import_module(None)` or whatever.
|
||||
if full_module_path := self._partitioner_full_module_path:
|
||||
return ".".join(full_module_path.split(".")[:-1])
|
||||
|
||||
if (shortname := self._partitioner_shortname) is None:
|
||||
raise ValueError(
|
||||
f"`.partitioner_module_qname` is undefined because FileType.{self.name} is not"
|
||||
f" partitionable. Use `.is_partitionable` to determine whether a `FileType`"
|
||||
f" is partitionable."
|
||||
)
|
||||
return f"unstructured.partition.{shortname}"
|
||||
|
||||
@property
|
||||
def partitioner_shortname(self) -> str | None:
|
||||
"""Familiar name of partitioner, like "image" for file-types that use `partition_image()`.
|
||||
|
||||
One use is to determine whether a file-type is one of the five image types, all of which
|
||||
are processed by `partition_image()`.
|
||||
|
||||
`None` for file-types that are not partitionable, although `.is_partitionable` is the
|
||||
preferred way of discovering that.
|
||||
"""
|
||||
return self._partitioner_shortname
|
||||
|
||||
BMP = (
|
||||
"bmp", # -- value for this Enum member, like BMP = "bmp" in a simple enum --
|
||||
"image", # -- partitioner_shortname --
|
||||
["unstructured_inference"], # -- importable_package_dependencies --
|
||||
"image", # -- extra_name - like `pip install "unstructured[image]"` in this case --
|
||||
[".bmp"], # -- extensions - filename extensions that map to this file-type --
|
||||
"image/bmp", # -- canonical_mime_type - MIME-type written to `.metadata.filetype` --
|
||||
cast(list[str], []), # -- alias_mime-types - other MIME-types that map to this file-type --
|
||||
)
|
||||
CSV = (
|
||||
"csv",
|
||||
"csv",
|
||||
["pandas"],
|
||||
"csv",
|
||||
[".csv"],
|
||||
"text/csv",
|
||||
[
|
||||
"application/csv",
|
||||
"application/x-csv",
|
||||
"text/comma-separated-values",
|
||||
"text/x-comma-separated-values",
|
||||
"text/x-csv",
|
||||
],
|
||||
)
|
||||
DOC = ("doc", "doc", ["docx"], "doc", [".doc"], "application/msword", cast(list[str], []))
|
||||
DOCX = (
|
||||
"docx",
|
||||
"docx",
|
||||
["docx"],
|
||||
"docx",
|
||||
[".docx"],
|
||||
"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
|
||||
cast(list[str], []),
|
||||
)
|
||||
EML = (
|
||||
"eml",
|
||||
"email",
|
||||
cast(list[str], []),
|
||||
None,
|
||||
[".eml", ".p7s"],
|
||||
"message/rfc822",
|
||||
cast(list[str], []),
|
||||
)
|
||||
EPUB = (
|
||||
"epub",
|
||||
"epub",
|
||||
["pypandoc"],
|
||||
"epub",
|
||||
[".epub"],
|
||||
"application/epub",
|
||||
["application/epub+zip"],
|
||||
)
|
||||
HEIC = (
|
||||
"heic",
|
||||
"image",
|
||||
["unstructured_inference"],
|
||||
"image",
|
||||
[".heic"],
|
||||
"image/heic",
|
||||
cast(list[str], []),
|
||||
)
|
||||
HTML = (
|
||||
"html",
|
||||
"html",
|
||||
cast(list[str], []),
|
||||
None,
|
||||
[".html", ".htm"],
|
||||
"text/html",
|
||||
cast(list[str], []),
|
||||
)
|
||||
JPG = (
|
||||
"jpg",
|
||||
"image",
|
||||
["unstructured_inference"],
|
||||
"image",
|
||||
[".jpeg", ".jpg"],
|
||||
"image/jpeg",
|
||||
cast(list[str], []),
|
||||
)
|
||||
JSON = (
|
||||
"json",
|
||||
"json",
|
||||
cast(list[str], []),
|
||||
None,
|
||||
[".json"],
|
||||
"application/json",
|
||||
cast(list[str], []),
|
||||
)
|
||||
MD = ("md", "md", ["markdown"], "md", [".md"], "text/markdown", ["text/x-markdown"])
|
||||
MSG = (
|
||||
"msg",
|
||||
"msg",
|
||||
["oxmsg"],
|
||||
"msg",
|
||||
[".msg"],
|
||||
"application/vnd.ms-outlook",
|
||||
cast(list[str], []),
|
||||
)
|
||||
NDJSON = (
|
||||
"ndjson",
|
||||
"ndjson",
|
||||
cast(list[str], []),
|
||||
None,
|
||||
[".ndjson"],
|
||||
"application/x-ndjson",
|
||||
cast(list[str], []),
|
||||
)
|
||||
ODT = (
|
||||
"odt",
|
||||
"odt",
|
||||
["docx", "pypandoc"],
|
||||
"odt",
|
||||
[".odt"],
|
||||
"application/vnd.oasis.opendocument.text",
|
||||
cast(list[str], []),
|
||||
)
|
||||
ORG = ("org", "org", ["pypandoc"], "org", [".org"], "text/org", cast(list[str], []))
|
||||
PDF = (
|
||||
"pdf",
|
||||
"pdf",
|
||||
["pdf2image", "pdfminer", "PIL"],
|
||||
"pdf",
|
||||
[".pdf"],
|
||||
"application/pdf",
|
||||
cast(list[str], []),
|
||||
)
|
||||
PNG = (
|
||||
"png",
|
||||
"image",
|
||||
["unstructured_inference"],
|
||||
"image",
|
||||
[".png"],
|
||||
"image/png",
|
||||
cast(list[str], []),
|
||||
)
|
||||
PPT = (
|
||||
"ppt",
|
||||
"ppt",
|
||||
["pptx"],
|
||||
"ppt",
|
||||
[".ppt"],
|
||||
"application/vnd.ms-powerpoint",
|
||||
cast(list[str], []),
|
||||
)
|
||||
PPTX = (
|
||||
"pptx",
|
||||
"pptx",
|
||||
["pptx"],
|
||||
"pptx",
|
||||
[".pptx"],
|
||||
"application/vnd.openxmlformats-officedocument.presentationml.presentation",
|
||||
cast(list[str], []),
|
||||
)
|
||||
RST = ("rst", "rst", ["pypandoc"], "rst", [".rst"], "text/x-rst", cast(list[str], []))
|
||||
RTF = ("rtf", "rtf", ["pypandoc"], "rtf", [".rtf"], "text/rtf", ["application/rtf"])
|
||||
TIFF = (
|
||||
"tiff",
|
||||
"image",
|
||||
["unstructured_inference"],
|
||||
"image",
|
||||
[".tiff"],
|
||||
"image/tiff",
|
||||
cast(list[str], []),
|
||||
)
|
||||
TSV = ("tsv", "tsv", ["pandas"], "tsv", [".tab", ".tsv"], "text/tsv", cast(list[str], []))
|
||||
TXT = (
|
||||
"txt",
|
||||
"text",
|
||||
cast(list[str], []),
|
||||
None,
|
||||
[
|
||||
".txt",
|
||||
".text",
|
||||
# NOTE(robinson) - for now we are treating code files as plain text
|
||||
".c",
|
||||
".cc",
|
||||
".cpp",
|
||||
".cs",
|
||||
".cxx",
|
||||
".go",
|
||||
".java",
|
||||
".js",
|
||||
".log",
|
||||
".php",
|
||||
".py",
|
||||
".rb",
|
||||
".swift",
|
||||
".ts",
|
||||
".yaml",
|
||||
".yml",
|
||||
],
|
||||
"text/plain",
|
||||
[
|
||||
# NOTE(robinson) - In the future, we may have special processing for YAML files
|
||||
# instead of treating them as plaintext.
|
||||
"text/yaml",
|
||||
"application/x-yaml",
|
||||
"application/yaml",
|
||||
"text/x-yaml",
|
||||
],
|
||||
)
|
||||
WAV = (
|
||||
"wav",
|
||||
None,
|
||||
cast(list[str], []),
|
||||
None,
|
||||
[".wav"],
|
||||
"audio/wav",
|
||||
[
|
||||
"audio/vnd.wav",
|
||||
"audio/vnd.wave",
|
||||
"audio/wave",
|
||||
"audio/x-pn-wav",
|
||||
"audio/x-wav",
|
||||
],
|
||||
)
|
||||
XLS = (
|
||||
"xls",
|
||||
"xlsx",
|
||||
["pandas", "openpyxl"],
|
||||
"xlsx",
|
||||
[".xls"],
|
||||
"application/vnd.ms-excel",
|
||||
cast(list[str], []),
|
||||
)
|
||||
XLSX = (
|
||||
"xlsx",
|
||||
"xlsx",
|
||||
["pandas", "openpyxl"],
|
||||
"xlsx",
|
||||
[".xlsx"],
|
||||
"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
|
||||
cast(list[str], []),
|
||||
)
|
||||
XML = ("xml", "xml", cast(list[str], []), None, [".xml"], "application/xml", ["text/xml"])
|
||||
ZIP = ("zip", None, cast(list[str], []), None, [".zip"], "application/zip", cast(list[str], []))
|
||||
|
||||
UNK = (
|
||||
"unk",
|
||||
None,
|
||||
cast(list[str], []),
|
||||
None,
|
||||
cast(list[str], []),
|
||||
"application/octet-stream",
|
||||
cast(list[str], []),
|
||||
)
|
||||
EMPTY = (
|
||||
"empty",
|
||||
None,
|
||||
cast(list[str], []),
|
||||
None,
|
||||
cast(list[str], []),
|
||||
"inode/x-empty",
|
||||
cast(list[str], []),
|
||||
)
|
||||
|
||||
|
||||
def create_file_type(
|
||||
name: str,
|
||||
*,
|
||||
canonical_mime_type: str,
|
||||
importable_package_dependencies: Iterable[str] | None = None,
|
||||
extra_name: str | None = None,
|
||||
extensions: Iterable[str] | None = None,
|
||||
alias_mime_types: Iterable[str] | None = None,
|
||||
) -> FileType:
|
||||
"""Register a new FileType member."""
|
||||
type_ = _create_file_type_enum(
|
||||
FileType,
|
||||
name,
|
||||
None,
|
||||
importable_package_dependencies or cast(list[str], []),
|
||||
extra_name,
|
||||
extensions or cast(list[str], []),
|
||||
canonical_mime_type,
|
||||
alias_mime_types or cast(list[str], []),
|
||||
None,
|
||||
)
|
||||
type_._name_ = name
|
||||
FileType._member_map_[name] = type_
|
||||
return type_
|
||||
|
||||
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
|
||||
def register_partitioner(
|
||||
file_type: FileType,
|
||||
) -> Callable[[Callable[_P, list[Element]]], Callable[_P, list[Element]]]:
|
||||
def decorator(func: Callable[_P, list[Element]]) -> Callable[_P, list[Element]]:
|
||||
file_type._partitioner_full_module_path = func.__module__ + "." + func.__name__
|
||||
return func
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
Adds support for working with newline-delimited JSON (ndjson) files. This format is useful for
|
||||
streaming json content that would otherwise not be possible using raw JSON files.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import IO, Any
|
||||
|
||||
|
||||
def dumps(obj: list[dict[str, Any]], **kwargs) -> str:
|
||||
"""
|
||||
Converts the list of dictionaries into string representation
|
||||
|
||||
Args:
|
||||
obj (list[dict[str, Any]]): List of dictionaries to convert
|
||||
**kwargs: Additional keyword arguments to pass to json.dumps
|
||||
|
||||
Returns:
|
||||
str: string representation of the list of dictionaries
|
||||
"""
|
||||
return "\n".join(json.dumps(each, **kwargs) for each in obj)
|
||||
|
||||
|
||||
def dump(obj: list[dict[str, Any]], fp: IO, **kwargs) -> None:
|
||||
"""
|
||||
Writes the list of dictionaries to a newline-delimited file
|
||||
|
||||
Args:
|
||||
obj (list[dict[str, Any]]): List of dictionaries to convert
|
||||
fp (IO): File pointer to write the string representation to
|
||||
**kwargs: Additional keyword arguments to pass to json.dumps
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# Indent breaks ndjson formatting
|
||||
kwargs["indent"] = None
|
||||
text = dumps(obj, **kwargs)
|
||||
fp.write(text)
|
||||
|
||||
|
||||
def loads(s: str, **kwargs) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Converts the raw string into a list of dictionaries
|
||||
|
||||
Args:
|
||||
s (str): Raw string to convert
|
||||
**kwargs: Additional keyword arguments to pass to json.loads
|
||||
|
||||
Returns:
|
||||
list[dict[str, Any]]: List of dictionaries parsed from the input string
|
||||
"""
|
||||
return [json.loads(line, **kwargs) for line in s.splitlines()]
|
||||
|
||||
|
||||
def load(fp: IO, **kwargs) -> list[dict[str, Any]]:
|
||||
"""
|
||||
Converts the contents of the file into a list of dictionaries
|
||||
|
||||
Args:
|
||||
fp (IO): File pointer to read the string representation from
|
||||
**kwargs: Additional keyword arguments to pass to json.loads
|
||||
|
||||
Returns:
|
||||
list[dict[str, Any]]: List of dictionaries parsed from the file
|
||||
"""
|
||||
return loads(fp.read(), **kwargs)
|
||||
@@ -0,0 +1,24 @@
|
||||
import logging
|
||||
|
||||
from unstructured.utils import scarf_analytics
|
||||
|
||||
logger = logging.getLogger("unstructured")
|
||||
trace_logger = logging.getLogger("unstructured.trace")
|
||||
|
||||
# Create a custom logging level
|
||||
DETAIL = 15
|
||||
logging.addLevelName(DETAIL, "DETAIL")
|
||||
|
||||
|
||||
# Create a custom log method for the "DETAIL" level
|
||||
def detail(self, message, *args, **kws):
|
||||
if self.isEnabledFor(DETAIL):
|
||||
self._log(DETAIL, message, args, **kws)
|
||||
|
||||
|
||||
# Note(Trevor,Crag): to opt out of scarf analytics, set the environment variable:
|
||||
# SCARF_NO_ANALYTICS=true. See the README for more info.
|
||||
scarf_analytics()
|
||||
|
||||
# Add the custom log method to the logging.Logger class
|
||||
logging.Logger.detail = detail # type: ignore
|
||||
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,108 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
|
||||
from typing_extensions import TypeAlias
|
||||
|
||||
FrequencyDict: TypeAlias = "dict[tuple[str, int | None], int]"
|
||||
"""Like:
|
||||
{
|
||||
("ListItem", 0): 2,
|
||||
("NarrativeText", None): 2,
|
||||
("Title", 0): 5,
|
||||
("UncategorizedText", None): 6,
|
||||
}
|
||||
"""
|
||||
|
||||
|
||||
def get_element_type_frequency(
|
||||
elements: str,
|
||||
) -> FrequencyDict:
|
||||
"""
|
||||
Calculate the frequency of Element Types from a list of elements.
|
||||
|
||||
Args:
|
||||
elements (str): String-formatted json of all elements (as a result of elements_to_json).
|
||||
Returns:
|
||||
Element type and its frequency in dictionary format.
|
||||
"""
|
||||
frequency: dict[tuple[str, int | None], int] = {}
|
||||
if len(elements) == 0:
|
||||
return frequency
|
||||
for element in json.loads(elements):
|
||||
type = element.get("type")
|
||||
category_depth = element["metadata"].get("category_depth")
|
||||
key = (type, category_depth)
|
||||
if key not in frequency:
|
||||
frequency[key] = 1
|
||||
else:
|
||||
frequency[key] += 1
|
||||
return frequency
|
||||
|
||||
|
||||
def calculate_element_type_percent_match(
|
||||
output: FrequencyDict,
|
||||
source: FrequencyDict,
|
||||
category_depth_weight: float = 0.5,
|
||||
) -> float:
|
||||
"""Calculate the percent match between two frequency dictionary.
|
||||
|
||||
Intended to use with `get_element_type_frequency` function. The function counts the absolute
|
||||
exact match (type and depth), and counts the weighted match (correct type but different depth),
|
||||
then normalized with source's total elements.
|
||||
"""
|
||||
if len(output) == 0 or len(source) == 0:
|
||||
return 0.0
|
||||
|
||||
output_copy = output.copy()
|
||||
source_copy = source.copy()
|
||||
total_source_element_count = 0
|
||||
total_match_element_count = 0
|
||||
|
||||
unmatched_depth_output: dict[str, int] = {}
|
||||
unmatched_depth_source: dict[str, int] = {}
|
||||
|
||||
# loop through the output list to find match with source
|
||||
for k, _ in output_copy.items():
|
||||
if k in source_copy:
|
||||
match_count = min(output_copy[k], source_copy[k])
|
||||
total_match_element_count += match_count
|
||||
total_source_element_count += match_count
|
||||
|
||||
# update the dictionary by removing already matched values
|
||||
output_copy[k] -= match_count
|
||||
source_copy[k] -= match_count
|
||||
|
||||
# add unmatched leftovers from output_copy to a new dictionary
|
||||
element_type = k[0]
|
||||
if element_type not in unmatched_depth_output:
|
||||
unmatched_depth_output[element_type] = output_copy[k]
|
||||
else:
|
||||
unmatched_depth_output[element_type] += output_copy[k]
|
||||
|
||||
# add unmatched leftovers from source_copy to a new dictionary
|
||||
unmatched_depth_source = _convert_to_frequency_without_depth(source_copy)
|
||||
|
||||
# loop through the source list to match any existing partial match left
|
||||
for k, _ in unmatched_depth_source.items():
|
||||
total_source_element_count += unmatched_depth_source[k]
|
||||
if k in unmatched_depth_output:
|
||||
match_count = min(unmatched_depth_output[k], unmatched_depth_source[k])
|
||||
total_match_element_count += match_count * category_depth_weight
|
||||
|
||||
return min(max(total_match_element_count / total_source_element_count, 0.0), 1.0)
|
||||
|
||||
|
||||
def _convert_to_frequency_without_depth(d: FrequencyDict) -> dict[str, int]:
|
||||
"""
|
||||
Takes in element frequency with depth of format (type, depth): value
|
||||
and converts to dictionary without depth of format type: value
|
||||
"""
|
||||
res: dict[str, int] = {}
|
||||
for k, v in d.items():
|
||||
element_type = k[0]
|
||||
if element_type not in res:
|
||||
res[element_type] = v
|
||||
else:
|
||||
res[element_type] += v
|
||||
return res
|
||||
@@ -0,0 +1,897 @@
|
||||
#! /usr/bin/env python3
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import concurrent.futures
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from tqdm import tqdm
|
||||
|
||||
from unstructured.metrics.element_type import (
|
||||
calculate_element_type_percent_match,
|
||||
get_element_type_frequency,
|
||||
)
|
||||
from unstructured.metrics.object_detection import (
|
||||
ObjectDetectionEvalProcessor,
|
||||
)
|
||||
from unstructured.metrics.table.table_eval import TableEvalProcessor
|
||||
from unstructured.metrics.text_extraction import calculate_accuracy, calculate_percent_missing_text
|
||||
from unstructured.metrics.utils import (
|
||||
_count,
|
||||
_display,
|
||||
_format_grouping_output,
|
||||
_mean,
|
||||
_prepare_output_cct,
|
||||
_pstdev,
|
||||
_read_text_file,
|
||||
_rename_aggregated_columns,
|
||||
_stdev,
|
||||
_write_to_file,
|
||||
)
|
||||
|
||||
logger = logging.getLogger("unstructured.eval")
|
||||
handler = logging.StreamHandler()
|
||||
handler.name = "eval_log_handler"
|
||||
formatter = logging.Formatter("%(asctime)s %(processName)-10s %(levelname)-8s %(message)s")
|
||||
handler.setFormatter(formatter)
|
||||
|
||||
# Only want to add the handler once
|
||||
if "eval_log_handler" not in [h.name for h in logger.handlers]:
|
||||
logger.addHandler(handler)
|
||||
|
||||
logger.setLevel(logging.DEBUG)
|
||||
|
||||
AGG_HEADERS = ["metric", "average", "sample_sd", "population_sd", "count"]
|
||||
AGG_HEADERS_MAPPING = {
|
||||
"index": "metric",
|
||||
"_mean": "average",
|
||||
"_stdev": "sample_sd",
|
||||
"_pstdev": "population_sd",
|
||||
"_count": "count",
|
||||
}
|
||||
OUTPUT_TYPE_OPTIONS = ["json", "txt"]
|
||||
|
||||
|
||||
@dataclass
|
||||
class BaseMetricsCalculator(ABC):
|
||||
"""Foundation class for specialized metrics calculators.
|
||||
|
||||
It provides a common interface for calculating metrics based on outputs and ground truths.
|
||||
Those can be provided as either directories or lists of files.
|
||||
"""
|
||||
|
||||
documents_dir: str | Path
|
||||
ground_truths_dir: str | Path
|
||||
|
||||
def __post_init__(self):
|
||||
"""Discover all files in the provided directories."""
|
||||
self.documents_dir = Path(self.documents_dir).resolve()
|
||||
self.ground_truths_dir = Path(self.ground_truths_dir).resolve()
|
||||
|
||||
# -- auto-discover all files in the directories --
|
||||
self._document_paths = [
|
||||
path.relative_to(self.documents_dir)
|
||||
for path in self.documents_dir.glob("*")
|
||||
if path.is_file()
|
||||
]
|
||||
self._ground_truth_paths = [
|
||||
path.relative_to(self.ground_truths_dir)
|
||||
for path in self.ground_truths_dir.glob("*")
|
||||
if path.is_file()
|
||||
]
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def default_tsv_name(self):
|
||||
"""Default name for the per-document metrics TSV file."""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def default_agg_tsv_name(self):
|
||||
"""Default name for the aggregated metrics TSV file."""
|
||||
|
||||
@abstractmethod
|
||||
def _generate_dataframes(self, rows: list) -> tuple[pd.DataFrame, pd.DataFrame]:
|
||||
"""Generates pandas DataFrames from the list of rows.
|
||||
|
||||
The first DF (index 0) is a dataframe containing metrics per file.
|
||||
The second DF (index 1) is a dataframe containing the aggregated
|
||||
metrics.
|
||||
"""
|
||||
|
||||
def on_files(
|
||||
self,
|
||||
document_paths: Optional[list[str | Path]] = None,
|
||||
ground_truth_paths: Optional[list[str | Path]] = None,
|
||||
) -> BaseMetricsCalculator:
|
||||
"""Overrides the default list of files to process."""
|
||||
if document_paths:
|
||||
self._document_paths = [Path(p) for p in document_paths]
|
||||
|
||||
if ground_truth_paths:
|
||||
self._ground_truth_paths = [Path(p) for p in ground_truth_paths]
|
||||
|
||||
return self
|
||||
|
||||
def calculate(
|
||||
self,
|
||||
executor: Optional[concurrent.futures.Executor] = None,
|
||||
export_dir: Optional[str | Path] = None,
|
||||
visualize_progress: bool = True,
|
||||
display_agg_df: bool = True,
|
||||
) -> pd.DataFrame:
|
||||
"""Calculates metrics for each document using the provided executor.
|
||||
|
||||
* Optionally, the results can be exported and displayed.
|
||||
* It loops through the list of structured output from all of `documents_dir` or
|
||||
selected files from `document_paths`, and compares them with gold-standard
|
||||
of the same file name under `ground_truths_dir` or selected files from `ground_truth_paths`.
|
||||
|
||||
Args:
|
||||
executor: concurrent.futures.Executor instance
|
||||
export_dir: directory to export the results
|
||||
visualize_progress: whether to display progress bar
|
||||
display_agg_df: whether to display the aggregated results
|
||||
|
||||
Returns:
|
||||
Metrics for each document as a pandas DataFrame
|
||||
"""
|
||||
if executor is None:
|
||||
executor = self._default_executor()
|
||||
rows = self._process_all_documents(executor, visualize_progress)
|
||||
df, agg_df = self._generate_dataframes(rows)
|
||||
|
||||
if export_dir is not None:
|
||||
_write_to_file(export_dir, self.default_tsv_name, df)
|
||||
_write_to_file(export_dir, self.default_agg_tsv_name, agg_df)
|
||||
|
||||
if display_agg_df is True:
|
||||
_display(agg_df)
|
||||
return df
|
||||
|
||||
@classmethod
|
||||
def _default_executor(cls):
|
||||
max_processors = int(os.environ.get("MAX_PROCESSES", os.cpu_count()))
|
||||
logger.info(f"Configuring a pool of {max_processors} processors for parallel processing.")
|
||||
return cls._get_executor_class()(max_workers=max_processors)
|
||||
|
||||
@classmethod
|
||||
def _get_executor_class(
|
||||
cls,
|
||||
) -> type[concurrent.futures.ThreadPoolExecutor] | type[concurrent.futures.ProcessPoolExecutor]:
|
||||
return concurrent.futures.ProcessPoolExecutor
|
||||
|
||||
def _process_all_documents(
|
||||
self, executor: concurrent.futures.Executor, visualize_progress: bool
|
||||
) -> list:
|
||||
"""Triggers processing of all documents using the provided executor.
|
||||
|
||||
Failures are omitted from the returned result.
|
||||
"""
|
||||
with executor:
|
||||
return [
|
||||
row
|
||||
for row in tqdm(
|
||||
executor.map(self._try_process_document, self._document_paths),
|
||||
total=len(self._document_paths),
|
||||
leave=False,
|
||||
disable=not visualize_progress,
|
||||
)
|
||||
if row is not None
|
||||
]
|
||||
|
||||
def _try_process_document(self, doc: Path) -> Optional[list]:
|
||||
"""Safe wrapper around the document processing method."""
|
||||
logger.info(f"Processing {doc}")
|
||||
try:
|
||||
return self._process_document(doc)
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to process document {doc}: {e}")
|
||||
return None
|
||||
|
||||
@abstractmethod
|
||||
def _process_document(self, doc: Path) -> Optional[list]:
|
||||
"""Should return all metadata and metrics for a single document."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TableStructureMetricsCalculator(BaseMetricsCalculator):
|
||||
"""Calculates the following metrics for tables:
|
||||
- tables found accuracy
|
||||
- table-level accuracy
|
||||
- element in column index accuracy
|
||||
- element in row index accuracy
|
||||
- element's column content accuracy
|
||||
- element's row content accuracy
|
||||
It also calculates the aggregated accuracy.
|
||||
"""
|
||||
|
||||
cutoff: Optional[float] = None
|
||||
weighted_average: bool = True
|
||||
include_false_positives: bool = True
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
|
||||
@property
|
||||
def supported_metric_names(self):
|
||||
return [
|
||||
"total_tables",
|
||||
"table_level_acc",
|
||||
"table_detection_recall",
|
||||
"table_detection_precision",
|
||||
"table_detection_f1",
|
||||
"composite_structure_acc",
|
||||
"element_col_level_index_acc",
|
||||
"element_row_level_index_acc",
|
||||
"element_col_level_content_acc",
|
||||
"element_row_level_content_acc",
|
||||
]
|
||||
|
||||
@property
|
||||
def default_tsv_name(self):
|
||||
return "all-docs-table-structure-accuracy.tsv"
|
||||
|
||||
@property
|
||||
def default_agg_tsv_name(self):
|
||||
return "aggregate-table-structure-accuracy.tsv"
|
||||
|
||||
def _process_document(self, doc: Path) -> Optional[list]:
|
||||
doc_path = Path(doc)
|
||||
out_filename = doc_path.stem
|
||||
doctype = Path(out_filename).suffix
|
||||
src_gt_filename = out_filename + ".json"
|
||||
connector = doc_path.parts[-2] if len(doc_path.parts) > 1 else None
|
||||
|
||||
if src_gt_filename in self._ground_truth_paths: # type: ignore
|
||||
return None
|
||||
|
||||
prediction_file = self.documents_dir / doc
|
||||
if not prediction_file.exists():
|
||||
logger.warning(f"Prediction file {prediction_file} does not exist, skipping")
|
||||
return None
|
||||
|
||||
ground_truth_file = self.ground_truths_dir / src_gt_filename
|
||||
if not ground_truth_file.exists():
|
||||
logger.warning(f"Ground truth file {ground_truth_file} does not exist, skipping")
|
||||
return None
|
||||
|
||||
processor_from_text_as_html = TableEvalProcessor.from_json_files(
|
||||
prediction_file=prediction_file,
|
||||
ground_truth_file=ground_truth_file,
|
||||
cutoff=self.cutoff,
|
||||
source_type="html",
|
||||
)
|
||||
report_from_html = processor_from_text_as_html.process_file()
|
||||
return [
|
||||
out_filename,
|
||||
doctype,
|
||||
connector,
|
||||
report_from_html.total_predicted_tables,
|
||||
] + [getattr(report_from_html, metric) for metric in self.supported_metric_names]
|
||||
|
||||
def _generate_dataframes(self, rows):
|
||||
headers = [
|
||||
"filename",
|
||||
"doctype",
|
||||
"connector",
|
||||
"total_predicted_tables",
|
||||
] + self.supported_metric_names
|
||||
|
||||
df = pd.DataFrame(rows, columns=headers)
|
||||
df["_table_weights"] = df["total_tables"]
|
||||
|
||||
if self.include_false_positives:
|
||||
# we give false positive tables a 1 table worth of weight in computing table level acc
|
||||
df["_table_weights"][df.total_tables.eq(0) & df.total_predicted_tables.gt(0)] = 1
|
||||
|
||||
# filter down to only those with actual and/or predicted tables
|
||||
has_tables_df = df[df["_table_weights"] > 0]
|
||||
|
||||
if not self.weighted_average:
|
||||
# for all non zero elements assign them value 1
|
||||
df["_table_weights"] = df["_table_weights"].apply(
|
||||
lambda table_weight: 1 if table_weight != 0 else 0
|
||||
)
|
||||
|
||||
if has_tables_df.empty:
|
||||
agg_df = pd.DataFrame(
|
||||
[[metric, None, None, None, 0] for metric in self.supported_metric_names]
|
||||
).reset_index()
|
||||
else:
|
||||
element_metrics_results = {}
|
||||
for metric in self.supported_metric_names:
|
||||
metric_df = has_tables_df[has_tables_df[metric].notnull()]
|
||||
agg_metric = metric_df[metric].agg([_stdev, _pstdev, _count]).transpose()
|
||||
if metric.startswith("total_tables"):
|
||||
agg_metric["_mean"] = metric_df[metric].mean()
|
||||
elif metric.startswith("table_level_acc"):
|
||||
agg_metric["_mean"] = np.round(
|
||||
np.average(metric_df[metric], weights=metric_df["_table_weights"]),
|
||||
3,
|
||||
)
|
||||
else:
|
||||
# false positive tables do not contribute to table structure and content
|
||||
# extraction metrics
|
||||
agg_metric["_mean"] = np.round(
|
||||
np.average(metric_df[metric], weights=metric_df["total_tables"]),
|
||||
3,
|
||||
)
|
||||
if agg_metric.empty:
|
||||
element_metrics_results[metric] = pd.Series(
|
||||
data=[None, None, None, 0], index=["_mean", "_stdev", "_pstdev", "_count"]
|
||||
)
|
||||
else:
|
||||
element_metrics_results[metric] = agg_metric
|
||||
agg_df = pd.DataFrame(element_metrics_results).transpose().reset_index()
|
||||
agg_df = agg_df.rename(columns=AGG_HEADERS_MAPPING)
|
||||
return df, agg_df
|
||||
|
||||
|
||||
@dataclass
|
||||
class TextExtractionMetricsCalculator(BaseMetricsCalculator):
|
||||
"""Calculates text accuracy and percent missing between document and ground truth texts.
|
||||
|
||||
It also calculates the aggregated accuracy and percent missing.
|
||||
"""
|
||||
|
||||
group_by: Optional[str] = None
|
||||
weights: tuple[int, int, int] = (1, 1, 1)
|
||||
document_type: str = "json"
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self._validate_inputs()
|
||||
|
||||
@property
|
||||
def default_tsv_name(self) -> str:
|
||||
return "all-docs-cct.tsv"
|
||||
|
||||
@property
|
||||
def default_agg_tsv_name(self) -> str:
|
||||
return "aggregate-scores-cct.tsv"
|
||||
|
||||
def calculate(
|
||||
self,
|
||||
executor: Optional[concurrent.futures.Executor] = None,
|
||||
export_dir: Optional[str | Path] = None,
|
||||
visualize_progress: bool = True,
|
||||
display_agg_df: bool = True,
|
||||
) -> pd.DataFrame:
|
||||
"""See the parent class for the method's docstring."""
|
||||
df = super().calculate(
|
||||
executor=executor,
|
||||
export_dir=export_dir,
|
||||
visualize_progress=visualize_progress,
|
||||
display_agg_df=display_agg_df,
|
||||
)
|
||||
|
||||
if export_dir is not None and self.group_by:
|
||||
get_mean_grouping(self.group_by, df, export_dir, "text_extraction")
|
||||
return df
|
||||
|
||||
def _validate_inputs(self):
|
||||
if not self._document_paths:
|
||||
logger.info("No output files to calculate to edit distances for, exiting")
|
||||
sys.exit(0)
|
||||
if self.document_type not in OUTPUT_TYPE_OPTIONS:
|
||||
raise ValueError(
|
||||
"Specified file type under `documents_dir` or `output_list` should be one of "
|
||||
f"`json` or `txt`. The given file type is {self.document_type}, exiting."
|
||||
)
|
||||
for path in self._document_paths:
|
||||
try:
|
||||
path.suffixes[-1]
|
||||
except IndexError:
|
||||
logger.error(f"File {path} does not have a suffix, skipping")
|
||||
continue
|
||||
if path.suffixes[-1] != f".{self.document_type}":
|
||||
logger.warning(
|
||||
"The directory contains file type inconsistent with the given input. "
|
||||
"Please note that some files will be skipped."
|
||||
)
|
||||
if not all(path.suffixes[-1] == f".{self.document_type}" for path in self._document_paths):
|
||||
logger.warning(
|
||||
"The directory contains file type inconsistent with the given input. "
|
||||
"Please note that some files will be skipped."
|
||||
)
|
||||
|
||||
def _process_document(self, doc: Path) -> Optional[list]:
|
||||
filename = doc.stem
|
||||
doctype = doc.suffixes[-2]
|
||||
connector = doc.parts[0] if len(doc.parts) > 1 else None
|
||||
|
||||
output_cct, source_cct = self._get_ccts(doc)
|
||||
# NOTE(amadeusz): Levenshtein distance calculation takes too long
|
||||
# skip it if file sizes differ wildly
|
||||
if 0.5 < len(output_cct.encode()) / len(source_cct.encode()) < 2.0:
|
||||
accuracy = round(calculate_accuracy(output_cct, source_cct, self.weights), 3)
|
||||
else:
|
||||
# 0.01 to distinguish it was set manually
|
||||
accuracy = 0.01
|
||||
percent_missing = round(calculate_percent_missing_text(output_cct, source_cct), 3)
|
||||
return [filename, doctype, connector, accuracy, percent_missing]
|
||||
|
||||
def _get_ccts(self, doc: Path) -> tuple[str, str]:
|
||||
output_cct = _prepare_output_cct(
|
||||
docpath=self.documents_dir / doc, output_type=self.document_type
|
||||
)
|
||||
source_cct = _read_text_file(self.ground_truths_dir / doc.with_suffix(".txt"))
|
||||
|
||||
return output_cct, source_cct
|
||||
|
||||
def _generate_dataframes(self, rows):
|
||||
headers = ["filename", "doctype", "connector", "cct-accuracy", "cct-%missing"]
|
||||
df = pd.DataFrame(rows, columns=headers)
|
||||
|
||||
acc = df[["cct-accuracy"]].agg([_mean, _stdev, _pstdev, _count]).transpose()
|
||||
miss = df[["cct-%missing"]].agg([_mean, _stdev, _pstdev, _count]).transpose()
|
||||
if acc.shape[1] == 0 and miss.shape[1] == 0:
|
||||
agg_df = pd.DataFrame(columns=AGG_HEADERS)
|
||||
else:
|
||||
agg_df = pd.concat((acc, miss)).reset_index()
|
||||
agg_df.columns = AGG_HEADERS
|
||||
|
||||
return df, agg_df
|
||||
|
||||
|
||||
@dataclass
|
||||
class ElementTypeMetricsCalculator(BaseMetricsCalculator):
|
||||
"""
|
||||
Calculates element type frequency accuracy, percent missing and
|
||||
aggregated accuracy between document and ground truth.
|
||||
"""
|
||||
|
||||
group_by: Optional[str] = None
|
||||
|
||||
def calculate(
|
||||
self,
|
||||
executor: Optional[concurrent.futures.Executor] = None,
|
||||
export_dir: Optional[str | Path] = None,
|
||||
visualize_progress: bool = True,
|
||||
display_agg_df: bool = False,
|
||||
) -> pd.DataFrame:
|
||||
"""See the parent class for the method's docstring."""
|
||||
df = super().calculate(
|
||||
executor=executor,
|
||||
export_dir=export_dir,
|
||||
visualize_progress=visualize_progress,
|
||||
display_agg_df=display_agg_df,
|
||||
)
|
||||
|
||||
if export_dir is not None and self.group_by:
|
||||
get_mean_grouping(self.group_by, df, export_dir, "element_type")
|
||||
return df
|
||||
|
||||
@property
|
||||
def default_tsv_name(self) -> str:
|
||||
return "all-docs-element-type-frequency.tsv"
|
||||
|
||||
@property
|
||||
def default_agg_tsv_name(self) -> str:
|
||||
return "aggregate-scores-element-type.tsv"
|
||||
|
||||
def _process_document(self, doc: Path) -> Optional[list]:
|
||||
filename = doc.stem
|
||||
doctype = doc.suffixes[-2]
|
||||
connector = doc.parts[0] if len(doc.parts) > 1 else None
|
||||
|
||||
output = get_element_type_frequency(_read_text_file(self.documents_dir / doc))
|
||||
source = get_element_type_frequency(
|
||||
_read_text_file(self.ground_truths_dir / doc.with_suffix(".json"))
|
||||
)
|
||||
accuracy = round(calculate_element_type_percent_match(output, source), 3)
|
||||
return [filename, doctype, connector, accuracy]
|
||||
|
||||
def _generate_dataframes(self, rows):
|
||||
headers = ["filename", "doctype", "connector", "element-type-accuracy"]
|
||||
df = pd.DataFrame(rows, columns=headers)
|
||||
if df.empty:
|
||||
agg_df = pd.DataFrame(["element-type-accuracy", None, None, None, 0]).transpose()
|
||||
else:
|
||||
agg_df = df.agg({"element-type-accuracy": [_mean, _stdev, _pstdev, _count]}).transpose()
|
||||
agg_df = agg_df.reset_index()
|
||||
|
||||
agg_df.columns = AGG_HEADERS
|
||||
|
||||
return df, agg_df
|
||||
|
||||
|
||||
def get_mean_grouping(
|
||||
group_by: str,
|
||||
data_input: Union[pd.DataFrame, str],
|
||||
export_dir: str,
|
||||
eval_name: str,
|
||||
agg_name: Optional[str] = None,
|
||||
export_filename: Optional[str] = None,
|
||||
) -> None:
|
||||
"""Aggregates accuracy and missing metrics by column name 'doctype' or 'connector',
|
||||
or 'all' for all rows. Export to TSV.
|
||||
If `all`, passing export_name is recommended.
|
||||
|
||||
Args:
|
||||
group_by (str): Grouping category ('doctype' or 'connector' or 'all').
|
||||
data_input (Union[pd.DataFrame, str]): DataFrame or path to a CSV/TSV file.
|
||||
export_dir (str): Directory for the exported TSV file.
|
||||
eval_name (str): Evaluated metric ('text_extraction' or 'element_type').
|
||||
agg_name (str, optional): String to use with export filename. Default is `cct` for
|
||||
group_by `text_extraction` and `element-type` for `element_type`
|
||||
export_name (str, optional): Export filename.
|
||||
"""
|
||||
if group_by not in ("doctype", "connector") and group_by != "all":
|
||||
raise ValueError("Invalid grouping category. Returning a non-group evaluation.")
|
||||
|
||||
if eval_name == "text_extraction":
|
||||
agg_fields = ["cct-accuracy", "cct-%missing"]
|
||||
agg_name = "cct"
|
||||
elif eval_name == "element_type":
|
||||
agg_fields = ["element-type-accuracy"]
|
||||
agg_name = "element-type"
|
||||
elif eval_name == "object_detection":
|
||||
agg_fields = ["f1_score", "m_ap"]
|
||||
agg_name = "object-detection"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown metric for eval {eval_name}. "
|
||||
f"Expected `text_extraction` or `element_type` or `table_extraction`."
|
||||
)
|
||||
|
||||
if isinstance(data_input, str):
|
||||
if not os.path.exists(data_input):
|
||||
raise FileNotFoundError(f"File {data_input} not found.")
|
||||
if data_input.endswith(".csv"):
|
||||
df = pd.read_csv(data_input, header=None)
|
||||
elif data_input.endswith(".tsv"):
|
||||
df = pd.read_csv(data_input, sep="\t")
|
||||
elif data_input.endswith(".txt"):
|
||||
df = pd.read_csv(data_input, sep="\t", header=None)
|
||||
else:
|
||||
raise ValueError("Please provide a .csv or .tsv file.")
|
||||
else:
|
||||
df = data_input
|
||||
|
||||
if df.empty:
|
||||
raise SystemExit("Data is empty. Exiting.")
|
||||
elif group_by != "all" and (group_by not in df.columns or df[group_by].isnull().all()):
|
||||
raise SystemExit(
|
||||
f"Data cannot be aggregated by `{group_by}`."
|
||||
f" Check if it's empty or the column is missing/empty."
|
||||
)
|
||||
|
||||
grouped_df = []
|
||||
if group_by and group_by != "all":
|
||||
for field in agg_fields:
|
||||
grouped_df.append(
|
||||
_rename_aggregated_columns(
|
||||
df.groupby(group_by).agg({field: [_mean, _stdev, _pstdev, _count]})
|
||||
)
|
||||
)
|
||||
if group_by == "all":
|
||||
df["grouping_key"] = 0
|
||||
for field in agg_fields:
|
||||
grouped_df.append(
|
||||
_rename_aggregated_columns(
|
||||
df.groupby("grouping_key").agg({field: [_mean, _stdev, _pstdev, _count]})
|
||||
)
|
||||
)
|
||||
grouped_df = _format_grouping_output(*grouped_df)
|
||||
if "grouping_key" in grouped_df.columns.get_level_values(0):
|
||||
grouped_df = grouped_df.drop("grouping_key", axis=1, level=0)
|
||||
|
||||
if export_filename:
|
||||
if not export_filename.endswith(".tsv"):
|
||||
export_filename = export_filename + ".tsv"
|
||||
_write_to_file(export_dir, export_filename, grouped_df)
|
||||
else:
|
||||
_write_to_file(export_dir, f"all-{group_by}-agg-{agg_name}.tsv", grouped_df)
|
||||
|
||||
|
||||
def filter_metrics(
|
||||
data_input: Union[str, pd.DataFrame],
|
||||
filter_list: Union[str, List[str]],
|
||||
filter_by: str = "filename",
|
||||
export_filename: Optional[str] = None,
|
||||
export_dir: str = "metrics",
|
||||
return_type: str = "file",
|
||||
) -> Optional[pd.DataFrame]:
|
||||
"""Reads the data_input file and filter only selected row available in filter_list.
|
||||
|
||||
Args:
|
||||
data_input (str, dataframe): the source data, path to file or dataframe
|
||||
filter_list (str, list): the filter, path to file or list of string
|
||||
filter_by (str): data_input's column to filter the filter_list to
|
||||
export_filename (str, optional): export filename. required when return_type is "file"
|
||||
export_dir (str, optional): export directory. default to <current directory>/metrics
|
||||
return_type (str): "file" or "dataframe"
|
||||
"""
|
||||
if isinstance(data_input, str):
|
||||
if not os.path.exists(data_input):
|
||||
raise FileNotFoundError(f"File {data_input} not found.")
|
||||
if data_input.endswith(".csv"):
|
||||
df = pd.read_csv(data_input, header=None)
|
||||
elif data_input.endswith(".tsv"):
|
||||
df = pd.read_csv(data_input, sep="\t")
|
||||
elif data_input.endswith(".txt"):
|
||||
df = pd.read_csv(data_input, sep="\t", header=None)
|
||||
else:
|
||||
raise ValueError("Please provide a .csv or .tsv file.")
|
||||
else:
|
||||
df = data_input
|
||||
|
||||
if isinstance(filter_list, str):
|
||||
if not os.path.exists(filter_list):
|
||||
raise FileNotFoundError(f"File {filter_list} not found.")
|
||||
if filter_list.endswith(".csv"):
|
||||
filter_df = pd.read_csv(filter_list, header=None)
|
||||
elif filter_list.endswith(".tsv"):
|
||||
filter_df = pd.read_csv(filter_list, sep="\t")
|
||||
elif filter_list.endswith(".txt"):
|
||||
filter_df = pd.read_csv(filter_list, sep="\t", header=None)
|
||||
else:
|
||||
raise ValueError("Please provide a .csv or .tsv file.")
|
||||
filter_list = filter_df.iloc[:, 0].astype(str).values.tolist()
|
||||
elif not isinstance(filter_list, list):
|
||||
raise ValueError("Please provide a List of strings or path to file.")
|
||||
|
||||
if filter_by not in df.columns:
|
||||
raise ValueError("`filter_by` key does not exists in the data provided.")
|
||||
|
||||
res = df[df[filter_by].isin(filter_list)]
|
||||
|
||||
if res.empty:
|
||||
raise SystemExit("No common file names between data_input and filter_list. Exiting.")
|
||||
|
||||
if return_type == "dataframe":
|
||||
return res
|
||||
elif return_type == "file" and export_filename:
|
||||
_write_to_file(export_dir, export_filename, res)
|
||||
elif return_type == "file" and not export_filename:
|
||||
raise ValueError("Please provide `export_filename`.")
|
||||
else:
|
||||
raise ValueError("Return type must be either `dataframe` or `file`.")
|
||||
|
||||
|
||||
@dataclass
|
||||
class ObjectDetectionMetricsCalculatorBase(BaseMetricsCalculator, ABC):
|
||||
"""
|
||||
Calculates object detection metrics for each document:
|
||||
- f1 score
|
||||
- precision
|
||||
- recall
|
||||
- average precision (mAP)
|
||||
It also calculates aggregated metrics.
|
||||
"""
|
||||
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self._document_paths = [
|
||||
path.relative_to(self.documents_dir)
|
||||
for path in self.documents_dir.rglob("analysis/*/layout_dump/object_detection.json")
|
||||
if path.is_file()
|
||||
]
|
||||
|
||||
@property
|
||||
def supported_metric_names(self):
|
||||
return ["f1_score", "precision", "recall", "m_ap"]
|
||||
|
||||
@property
|
||||
def default_tsv_name(self):
|
||||
return "all-docs-object-detection-metrics.tsv"
|
||||
|
||||
@property
|
||||
def default_agg_tsv_name(self):
|
||||
return "aggregate-object-detection-metrics.tsv"
|
||||
|
||||
def _find_file_in_ground_truth(self, file_stem: str) -> Optional[Path]:
|
||||
"""Find the file corresponding to OD model dump file among the set of ground truth files
|
||||
|
||||
The files in ground truth paths keep the original extension and have .json suffix added,
|
||||
e.g.:
|
||||
some_document.pdf.json
|
||||
poster.jpg.json
|
||||
|
||||
To compare to `file_stem` we need to take the prefix part of the file, thus double-stem
|
||||
is applied.
|
||||
"""
|
||||
for path in self._ground_truth_paths:
|
||||
if Path(path.stem).stem == file_stem:
|
||||
return path
|
||||
return None
|
||||
|
||||
def _get_paths(self, doc: Path) -> tuple(str, Path, Path):
|
||||
"""Resolves ground doctype, prediction file path and ground truth path.
|
||||
|
||||
As OD dump directory structure differes from other simple outputs, it needs
|
||||
a specific processing to match the output OD dump file with corresponding
|
||||
OD GT file.
|
||||
|
||||
The outputs are placed in a dicrectory structure:
|
||||
|
||||
analysis
|
||||
|- document_name
|
||||
|- layout_dump
|
||||
|- object_detection.json
|
||||
|- bboxes # not used in this evaluation
|
||||
|
||||
and the GT file is pleced in od_gt directory for given dataset
|
||||
|
||||
dataset_name
|
||||
|- od_gt
|
||||
|- document_name.pdf.json
|
||||
|
||||
Args:
|
||||
doc (Path): path to the OD dump file
|
||||
|
||||
Returns:
|
||||
tuple: doctype, prediction file path, ground truth path
|
||||
"""
|
||||
od_dump_path = Path(doc)
|
||||
file_stem = od_dump_path.parts[-3] # we take the `document_name` - so the filename stem
|
||||
|
||||
src_gt_filename = self._find_file_in_ground_truth(file_stem)
|
||||
|
||||
if src_gt_filename not in self._ground_truth_paths:
|
||||
raise ValueError(f"Ground truth file {src_gt_filename} not found in list of GT files")
|
||||
|
||||
doctype = Path(src_gt_filename.stem).suffix[1:]
|
||||
|
||||
prediction_file = self.documents_dir / doc
|
||||
if not prediction_file.exists():
|
||||
logger.warning(f"Prediction file {prediction_file} does not exist, skipping")
|
||||
raise ValueError(f"Prediction file {prediction_file} does not exist")
|
||||
|
||||
ground_truth_file = self.ground_truths_dir / src_gt_filename
|
||||
if not ground_truth_file.exists():
|
||||
logger.warning(f"Ground truth file {ground_truth_file} does not exist, skipping")
|
||||
raise ValueError(f"Ground truth file {ground_truth_file} does not exist")
|
||||
|
||||
return doctype, prediction_file, ground_truth_file
|
||||
|
||||
def _generate_dataframes(self, rows) -> tuple[pd.DataFrame, pd.DataFrame]:
|
||||
headers = ["filename", "doctype", "connector"] + self.supported_metric_names
|
||||
df = pd.DataFrame(rows, columns=headers)
|
||||
|
||||
if df.empty:
|
||||
agg_df = pd.DataFrame(columns=AGG_HEADERS)
|
||||
else:
|
||||
element_metrics_results = {}
|
||||
for metric in self.supported_metric_names:
|
||||
metric_df = df[df[metric].notnull()]
|
||||
agg_metric = metric_df[metric].agg([_mean, _stdev, _pstdev, _count]).transpose()
|
||||
if agg_metric.empty:
|
||||
element_metrics_results[metric] = pd.Series(
|
||||
data=[None, None, None, 0], index=["_mean", "_stdev", "_pstdev", "_count"]
|
||||
)
|
||||
else:
|
||||
element_metrics_results[metric] = agg_metric
|
||||
agg_df = pd.DataFrame(element_metrics_results).transpose().reset_index()
|
||||
agg_df.columns = AGG_HEADERS
|
||||
|
||||
return df, agg_df
|
||||
|
||||
|
||||
class ObjectDetectionPerClassMetricsCalculator(ObjectDetectionMetricsCalculatorBase):
|
||||
def __post_init__(self):
|
||||
super().__post_init__()
|
||||
self.per_class_metric_names: list[str] | None = None
|
||||
self._set_supported_metrics()
|
||||
|
||||
@property
|
||||
def supported_metric_names(self):
|
||||
if self.per_class_metric_names:
|
||||
return self.per_class_metric_names
|
||||
else:
|
||||
raise ValueError("per_class_metrics not initialized - cannot get class names")
|
||||
|
||||
@property
|
||||
def default_tsv_name(self):
|
||||
return "all-docs-object-detection-metrics-per-class.tsv"
|
||||
|
||||
@property
|
||||
def default_agg_tsv_name(self):
|
||||
return "aggregate-object-detection-metrics-per-class.tsv"
|
||||
|
||||
def _process_document(self, doc: Path) -> Optional[list]:
|
||||
"""Calculate both class-aggregated and per-class metrics for a single document.
|
||||
|
||||
Args:
|
||||
doc (Path): path to the OD dump file
|
||||
|
||||
Returns:
|
||||
tuple: a tuple of aggregated and per-class metrics for a single document
|
||||
"""
|
||||
try:
|
||||
doctype, prediction_file, ground_truth_file = self._get_paths(doc)
|
||||
except ValueError as e:
|
||||
logger.error(f"Failed to process document {doc}: {e}")
|
||||
return None
|
||||
|
||||
processor = ObjectDetectionEvalProcessor.from_json_files(
|
||||
prediction_file_path=prediction_file,
|
||||
ground_truth_file_path=ground_truth_file,
|
||||
)
|
||||
_, per_class_metrics = processor.get_metrics()
|
||||
|
||||
per_class_metrics_row = [
|
||||
ground_truth_file.stem,
|
||||
doctype,
|
||||
None, # connector
|
||||
]
|
||||
|
||||
for combined_metric_name in self.supported_metric_names:
|
||||
metric = "_".join(combined_metric_name.split("_")[:-1])
|
||||
class_name = combined_metric_name.split("_")[-1]
|
||||
class_metrics = getattr(per_class_metrics, metric)
|
||||
per_class_metrics_row.append(class_metrics[class_name])
|
||||
return per_class_metrics_row
|
||||
|
||||
def _set_supported_metrics(self):
|
||||
"""Sets the supported metrics based on the classes found in the ground truth files.
|
||||
The difference between per class and aggregated calculator is that the list of classes
|
||||
(so the metrics) bases on the contents of the GT / prediction files.
|
||||
"""
|
||||
metrics = ["f1_score", "precision", "recall", "m_ap"]
|
||||
classes = set()
|
||||
for gt_file in self._ground_truth_paths:
|
||||
gt_file_path = self.ground_truths_dir / gt_file
|
||||
with open(gt_file_path) as f:
|
||||
gt = json.load(f)
|
||||
gt_classes = gt["object_detection_classes"]
|
||||
classes.update(gt_classes)
|
||||
per_class_metric_names = []
|
||||
for metric in metrics:
|
||||
for class_name in classes:
|
||||
per_class_metric_names.append(f"{metric}_{class_name}")
|
||||
self.per_class_metric_names = sorted(per_class_metric_names)
|
||||
|
||||
|
||||
class ObjectDetectionAggregatedMetricsCalculator(ObjectDetectionMetricsCalculatorBase):
|
||||
"""Calculates object detection metrics for each document and aggregates by all classes"""
|
||||
|
||||
@property
|
||||
def supported_metric_names(self):
|
||||
return ["f1_score", "precision", "recall", "m_ap"]
|
||||
|
||||
@property
|
||||
def default_tsv_name(self):
|
||||
return "all-docs-object-detection-metrics.tsv"
|
||||
|
||||
@property
|
||||
def default_agg_tsv_name(self):
|
||||
return "aggregate-object-detection-metrics.tsv"
|
||||
|
||||
def _process_document(self, doc: Path) -> Optional[list]:
|
||||
"""Calculate both class-aggregated and per-class metrics for a single document.
|
||||
|
||||
Args:
|
||||
doc (Path): path to the OD dump file
|
||||
|
||||
Returns:
|
||||
list: a list of aggregated metrics for a single document
|
||||
"""
|
||||
try:
|
||||
doctype, prediction_file, ground_truth_file = self._get_paths(doc)
|
||||
except ValueError as e:
|
||||
logger.error(f"Failed to process document {doc}: {e}")
|
||||
return None
|
||||
|
||||
processor = ObjectDetectionEvalProcessor.from_json_files(
|
||||
prediction_file_path=prediction_file,
|
||||
ground_truth_file_path=ground_truth_file,
|
||||
)
|
||||
metrics, _ = processor.get_metrics()
|
||||
|
||||
return [
|
||||
ground_truth_file.stem,
|
||||
doctype,
|
||||
None, # connector
|
||||
] + [getattr(metrics, metric) for metric in self.supported_metric_names]
|
||||
@@ -0,0 +1,719 @@
|
||||
"""
|
||||
Implements object detection metrics: average precision, precision, recall, and f1 score.
|
||||
"""
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
IOU_THRESHOLDS = torch.tensor(
|
||||
[0.5000, 0.5500, 0.6000, 0.6500, 0.7000, 0.7500, 0.8000, 0.8500, 0.9000, 0.9500]
|
||||
)
|
||||
SCORE_THRESHOLD = 0.1
|
||||
RECALL_THRESHOLDS = torch.arange(0, 1.01, 0.01)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ObjectDetectionAggregatedEvaluation:
|
||||
"""Class representing a gathered class-aggregated object detection metrics"""
|
||||
|
||||
f1_score: float
|
||||
precision: float
|
||||
recall: float
|
||||
m_ap: float
|
||||
|
||||
|
||||
@dataclass
|
||||
class ObjectDetectionPerClassEvaluation:
|
||||
"""Class representing a gathered object detection metrics per-class"""
|
||||
|
||||
f1_score: dict[str, float]
|
||||
precision: dict[str, float]
|
||||
recall: dict[str, float]
|
||||
m_ap: dict[str, float]
|
||||
|
||||
@classmethod
|
||||
def from_tensors(cls, ap, precision, recall, f1, class_labels):
|
||||
f1_score = {class_labels[i]: f1[i] for i in range(len(class_labels))}
|
||||
precision = {class_labels[i]: precision[i] for i in range(len(class_labels))}
|
||||
recall = {class_labels[i]: recall[i] for i in range(len(class_labels))}
|
||||
m_ap = {class_labels[i]: ap[i] for i in range(len(class_labels))}
|
||||
|
||||
return cls(f1_score, precision, recall, m_ap)
|
||||
|
||||
|
||||
class ObjectDetectionEvalProcessor:
|
||||
iou_thresholds = IOU_THRESHOLDS
|
||||
score_threshold = SCORE_THRESHOLD
|
||||
recall_thresholds = RECALL_THRESHOLDS
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
document_preds: list[torch.Tensor],
|
||||
document_targets: list[torch.Tensor],
|
||||
pages_height: list[int],
|
||||
pages_width: list[int],
|
||||
class_labels: list[str],
|
||||
device: str = "cpu",
|
||||
):
|
||||
"""
|
||||
Initializes the ObjectDetection prediction and ground truth.
|
||||
|
||||
Args:
|
||||
document_preds (list): list (of length pages of document) of
|
||||
Tensors of shape (num_predictions, 6)
|
||||
format: (x1, y1, x2, y2, confidence,class_label)
|
||||
where x1,y1,x2,y2 are according to image size
|
||||
document_targets (list): list (of length pages of document) of
|
||||
Tensors of shape (num_targets, 6)
|
||||
format: (label, x1, y1, x2, y2)
|
||||
where x,y,w,h are according to image size
|
||||
pages_height (list): list of height of each page in the document
|
||||
pages_width (list): list of width of each page in the document
|
||||
class_labels (list): list of class labels
|
||||
"""
|
||||
self.device = device
|
||||
self.document_preds = [pred.to(device) for pred in document_preds]
|
||||
self.document_targets = [target.to(device) for target in document_targets]
|
||||
self.pages_height = pages_height
|
||||
self.pages_width = pages_width
|
||||
self.class_labels = class_labels
|
||||
|
||||
@classmethod
|
||||
def from_json_files(
|
||||
cls,
|
||||
prediction_file_path: Path,
|
||||
ground_truth_file_path: Path,
|
||||
) -> "ObjectDetectionEvalProcessor":
|
||||
"""
|
||||
Initializes the ObjectDetection prediction and ground truth,
|
||||
and converts the data to the required format.
|
||||
|
||||
Args:
|
||||
prediction_file_path (Path): path to json file with predictions dump from OD model
|
||||
ground_truth_file_path (Path): path to json file with OD ground truth data
|
||||
"""
|
||||
# TODO: Test after https://unstructured-ai.atlassian.net/browse/ML-92
|
||||
# is done.
|
||||
with open(prediction_file_path) as f:
|
||||
predictions_data = json.load(f)
|
||||
with open(ground_truth_file_path) as f:
|
||||
ground_truth_data = json.load(f)
|
||||
|
||||
assert sorted(predictions_data["object_detection_classes"]) == sorted(
|
||||
ground_truth_data["object_detection_classes"]
|
||||
), "Classes in predictions and ground truth do not match."
|
||||
assert len(predictions_data["pages"]) == len(
|
||||
ground_truth_data["pages"]
|
||||
), "Pages number in predictions and ground truth do not match."
|
||||
for pred_page, gt_page in zip(
|
||||
sorted(predictions_data["pages"], key=lambda p: p["number"]),
|
||||
sorted(ground_truth_data["pages"], key=lambda p: p["number"]),
|
||||
):
|
||||
assert pred_page["number"] == gt_page["number"], (
|
||||
f"Page numbers in predictions {prediction_file_path.name} "
|
||||
f"({pred_page['number']}) and ground truth {ground_truth_file_path.name} "
|
||||
f"({gt_page['number']}) do not match."
|
||||
)
|
||||
page_num = pred_page["number"]
|
||||
|
||||
# TODO: translate the bboxes instead of raising error
|
||||
assert pred_page["size"] == gt_page["size"], (
|
||||
f"Page sizes in predictions {prediction_file_path.name} "
|
||||
f"({pred_page['size'][0]} x {pred_page['size'][1]}) "
|
||||
f"and ground truth {ground_truth_file_path.name} ({gt_page['size'][0]} x "
|
||||
f"{gt_page['size'][1]}) do not match for page {page_num}."
|
||||
)
|
||||
|
||||
class_labels = predictions_data["object_detection_classes"]
|
||||
document_preds = cls._process_data(predictions_data, class_labels, prediction=True)
|
||||
document_targets = cls._process_data(ground_truth_data, class_labels)
|
||||
pages_height, pages_width = cls._parse_page_dimensions(predictions_data)
|
||||
|
||||
return cls(document_preds, document_targets, pages_height, pages_width, class_labels)
|
||||
|
||||
def get_metrics(
|
||||
self,
|
||||
) -> tuple[ObjectDetectionAggregatedEvaluation, ObjectDetectionPerClassEvaluation]:
|
||||
"""Get per document OD metrics.
|
||||
|
||||
Returns:
|
||||
tuple: Tuple of ObjectDetectionAggregatedEvaluation and
|
||||
ObjectDetectionPerClassEvaluation
|
||||
"""
|
||||
document_matchings = []
|
||||
for preds, targets, height, width in zip(
|
||||
self.document_preds, self.document_targets, self.pages_height, self.pages_width
|
||||
):
|
||||
# iterate over each page
|
||||
page_matching_tensors = self._compute_page_detection_matching(
|
||||
preds=preds,
|
||||
targets=targets,
|
||||
height=height,
|
||||
width=width,
|
||||
)
|
||||
document_matchings.append(page_matching_tensors)
|
||||
|
||||
# compute metrics for all detections and targets
|
||||
mean_ap, mean_precision, mean_recall, mean_f1 = (
|
||||
-1.0,
|
||||
-1.0,
|
||||
-1.0,
|
||||
-1.0,
|
||||
)
|
||||
|
||||
num_cls = len(self.class_labels)
|
||||
mean_ap_per_class = np.full(num_cls, np.nan)
|
||||
mean_precision_per_class = np.full(num_cls, np.nan)
|
||||
mean_recall_per_class = np.full(num_cls, np.nan)
|
||||
mean_f1_per_class = np.full(num_cls, np.nan)
|
||||
|
||||
if len(document_matchings):
|
||||
matching_info_tensors = [torch.cat(x, 0) for x in list(zip(*document_matchings))]
|
||||
|
||||
# shape (n_class, nb_iou_thresh)
|
||||
(
|
||||
ap_per_present_classes,
|
||||
precision_per_present_classes,
|
||||
recall_per_present_classes,
|
||||
f1_per_present_classes,
|
||||
present_classes,
|
||||
) = self._compute_detection_metrics(
|
||||
*matching_info_tensors,
|
||||
)
|
||||
|
||||
# Precision, recall and f1 are computed for IoU threshold range, averaged over classes
|
||||
# results before version 3.0.4 (Dec 11 2022) were computed only for smallest value
|
||||
# (i.e IoU 0.5 if metric is @0.5:0.95)
|
||||
mean_precision, mean_recall, mean_f1 = (
|
||||
precision_per_present_classes.mean(),
|
||||
recall_per_present_classes.mean(),
|
||||
f1_per_present_classes.mean(),
|
||||
)
|
||||
|
||||
# MaP is averaged over IoU thresholds and over classes
|
||||
mean_ap = ap_per_present_classes.mean()
|
||||
|
||||
# Fill array of per-class AP scores with values for classes that were present in the
|
||||
# dataset
|
||||
ap_per_class = ap_per_present_classes.mean(1)
|
||||
precision_per_class = precision_per_present_classes.mean(1)
|
||||
recall_per_class = recall_per_present_classes.mean(1)
|
||||
f1_per_class = f1_per_present_classes.mean(1)
|
||||
for i, class_index in enumerate(present_classes):
|
||||
mean_ap_per_class[class_index] = float(ap_per_class[i])
|
||||
|
||||
mean_precision_per_class[class_index] = float(precision_per_class[i])
|
||||
mean_recall_per_class[class_index] = float(recall_per_class[i])
|
||||
mean_f1_per_class[class_index] = float(f1_per_class[i])
|
||||
|
||||
od_per_class_evaluation = ObjectDetectionPerClassEvaluation.from_tensors(
|
||||
ap=mean_ap_per_class,
|
||||
precision=mean_precision_per_class,
|
||||
recall=mean_recall_per_class,
|
||||
f1=mean_f1_per_class,
|
||||
class_labels=self.class_labels,
|
||||
)
|
||||
|
||||
od_evaluation = ObjectDetectionAggregatedEvaluation(
|
||||
f1_score=float(mean_f1),
|
||||
precision=float(mean_precision),
|
||||
recall=float(mean_recall),
|
||||
m_ap=float(mean_ap),
|
||||
)
|
||||
|
||||
return od_evaluation, od_per_class_evaluation
|
||||
|
||||
@staticmethod
|
||||
def _parse_page_dimensions(data: dict) -> tuple[list, list]:
|
||||
"""
|
||||
Process the page dimensions from the json file to the required format.
|
||||
"""
|
||||
pages_height = []
|
||||
pages_width = []
|
||||
for page in data["pages"]:
|
||||
pages_height.append(page["size"]["height"])
|
||||
pages_width.append(page["size"]["width"])
|
||||
return pages_height, pages_width
|
||||
|
||||
@staticmethod
|
||||
def _process_data(data: dict, class_labels, prediction: bool = False) -> list[dict]:
|
||||
"""
|
||||
Process the elements from the json file to the required format.
|
||||
"""
|
||||
pages_list = []
|
||||
for page in data["pages"]:
|
||||
page_elements = []
|
||||
for element in page["elements"]:
|
||||
# Extract coordinates, confidence, and class label from each prediction
|
||||
class_label = element["type"]
|
||||
class_idx = class_labels.index(class_label)
|
||||
x1, y1, x2, y2 = element["bbox"]
|
||||
if prediction:
|
||||
confidence = element["prob"]
|
||||
page_elements.append([x1, y1, x2, y2, confidence, class_idx])
|
||||
else:
|
||||
page_elements.append([class_idx, x1, y1, x2, y2])
|
||||
page_tensor = torch.tensor(page_elements)
|
||||
pages_list.append(page_tensor)
|
||||
|
||||
return pages_list
|
||||
|
||||
@staticmethod
|
||||
def _get_top_k_idx_per_cls(
|
||||
preds_scores: torch.Tensor, preds_cls: torch.Tensor, top_k: int
|
||||
) -> torch.Tensor:
|
||||
# From: https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/utils/detection_utils.py # noqa E501
|
||||
"""
|
||||
Get the indexes of all the top k predictions for every class
|
||||
|
||||
Args:
|
||||
preds_scores: The confidence scores, vector of shape (n_pred)
|
||||
preds_cls: The predicted class, vector of shape (n_pred)
|
||||
top_k: Number of predictions to keep per class, ordered by confidence score
|
||||
|
||||
Returns:
|
||||
top_k_idx: Indexes of the top k predictions. length <= (k * n_unique_class)
|
||||
"""
|
||||
n_unique_cls = torch.max(preds_cls)
|
||||
mask = preds_cls.view(-1, 1) == torch.arange(
|
||||
n_unique_cls + 1, device=preds_scores.device
|
||||
).view(1, -1)
|
||||
preds_scores_per_cls = preds_scores.view(-1, 1) * mask
|
||||
|
||||
sorted_scores_per_cls, sorting_idx = preds_scores_per_cls.sort(0, descending=True)
|
||||
idx_with_satisfying_scores = sorted_scores_per_cls[:top_k, :].nonzero(as_tuple=False)
|
||||
top_k_idx = sorting_idx[idx_with_satisfying_scores.split(1, dim=1)]
|
||||
return top_k_idx.view(-1)
|
||||
|
||||
@staticmethod
|
||||
def _change_bbox_bounds_for_image_size(
|
||||
boxes: np.ndarray, img_shape: tuple[int, int]
|
||||
) -> np.ndarray:
|
||||
# From: https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/utils/detection_utils.py # noqa E501
|
||||
"""
|
||||
Clips bboxes to image boundaries.
|
||||
|
||||
Args:
|
||||
bboxes: Input bounding boxes in XYXY format of [..., 4] shape
|
||||
img_shape: Image shape (height, width).
|
||||
Returns:
|
||||
clipped_boxes: Clipped bboxes in XYXY format of [..., 4] shape
|
||||
"""
|
||||
boxes[..., [0, 2]] = boxes[..., [0, 2]].clip(min=0, max=img_shape[1])
|
||||
boxes[..., [1, 3]] = boxes[..., [1, 3]].clip(min=0, max=img_shape[0])
|
||||
return boxes
|
||||
|
||||
@staticmethod
|
||||
def _box_iou(box1: torch.Tensor, box2: torch.Tensor) -> torch.Tensor:
|
||||
# From: https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/utils/detection_utils.py # noqa E501
|
||||
"""
|
||||
Return intersection-over-union (Jaccard index) of boxes.
|
||||
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
|
||||
|
||||
Args:
|
||||
box1: Tensor of shape [N, 4]
|
||||
box2: Tensor of shape [M, 4]
|
||||
|
||||
Returns:
|
||||
iou: Tensor of shape [N, M]: the NxM matrix containing the pairwise IoU values
|
||||
for every element in boxes1 and boxes2
|
||||
"""
|
||||
|
||||
def box_area(box):
|
||||
# box = 4xn
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
|
||||
area1 = box_area(box1.T)
|
||||
area2 = box_area(box2.T)
|
||||
|
||||
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
|
||||
inter = (
|
||||
(torch.min(box1[:, None, 2:], box2[:, 2:]) - torch.max(box1[:, None, :2], box2[:, :2]))
|
||||
.clamp(0)
|
||||
.prod(2)
|
||||
)
|
||||
return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter)
|
||||
|
||||
def _compute_targets(
|
||||
self,
|
||||
preds_box_xyxy: torch.Tensor,
|
||||
preds_cls: torch.Tensor,
|
||||
targets_box_xyxy: torch.Tensor,
|
||||
targets_cls: torch.Tensor,
|
||||
preds_matched: torch.Tensor,
|
||||
targets_matched: torch.Tensor,
|
||||
preds_idx_to_use: torch.Tensor,
|
||||
iou_thresholds: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
# From: https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/utils/detection_utils.py # noqa E501
|
||||
"""
|
||||
Computes the matching targets based on IoU for regular scenarios.
|
||||
|
||||
Args:
|
||||
preds_box_xyxy: (torch.Tensor) Predicted bounding boxes in XYXY format.
|
||||
preds_cls: (torch.Tensor) Predicted classes.
|
||||
targets_box_xyxy: (torch.Tensor) Target bounding boxes in XYXY format.
|
||||
targets_cls: (torch.Tensor) Target classes.
|
||||
preds_matched: (torch.Tensor) Tensor indicating which predictions are matched.
|
||||
targets_matched: (torch.Tensor) Tensor indicating which targets are matched.
|
||||
preds_idx_to_use: (torch.Tensor) Indices of predictions to use.
|
||||
|
||||
Returns:
|
||||
targets: Computed matching targets.
|
||||
"""
|
||||
# shape = (n_preds x n_targets)
|
||||
iou = self._box_iou(preds_box_xyxy[preds_idx_to_use], targets_box_xyxy)
|
||||
|
||||
# Fill IoU values at index (i, j) with 0 when the prediction (i) and target(j)
|
||||
# are of different class
|
||||
# Filling with 0 is equivalent to ignore these values
|
||||
# since with want IoU > iou_threshold > 0
|
||||
cls_mismatch = preds_cls[preds_idx_to_use].view(-1, 1) != targets_cls.view(1, -1)
|
||||
iou[cls_mismatch] = 0
|
||||
|
||||
# The matching priority is first detection confidence and then IoU value.
|
||||
# The detection is already sorted by confidence in NMS,
|
||||
# so here for each prediction we order the targets by iou.
|
||||
sorted_iou, target_sorted = iou.sort(descending=True, stable=True)
|
||||
|
||||
# Only iterate over IoU values higher than min threshold to speed up the process
|
||||
for pred_selected_i, target_sorted_i in (sorted_iou > iou_thresholds[0]).nonzero(
|
||||
as_tuple=False
|
||||
):
|
||||
# pred_selected_i and target_sorted_i are relative to filters/sorting,
|
||||
# so we extract their absolute indexes
|
||||
pred_i = preds_idx_to_use[pred_selected_i]
|
||||
target_i = target_sorted[pred_selected_i, target_sorted_i]
|
||||
|
||||
# Vector[j], True when IoU(pred_i, target_i) is above the (j)th threshold
|
||||
is_iou_above_threshold = sorted_iou[pred_selected_i, target_sorted_i] > iou_thresholds
|
||||
|
||||
# Vector[j], True when both pred_i and target_i are not matched yet
|
||||
# for the (j)th threshold
|
||||
are_candidates_free = torch.logical_and(
|
||||
~preds_matched[pred_i, :], ~targets_matched[target_i, :]
|
||||
)
|
||||
|
||||
# Vector[j], True when (pred_i, target_i) can be matched for the (j)th threshold
|
||||
are_candidates_good = torch.logical_and(is_iou_above_threshold, are_candidates_free)
|
||||
|
||||
# For every threshold (j) where target_i and pred_i can be matched together
|
||||
# ( are_candidates_good[j]==True )
|
||||
# fill the matching placeholders with True
|
||||
targets_matched[target_i, are_candidates_good] = True
|
||||
preds_matched[pred_i, are_candidates_good] = True
|
||||
|
||||
# When all the targets are matched with a prediction for every IoU Threshold, stop.
|
||||
if targets_matched.all():
|
||||
break
|
||||
|
||||
return preds_matched
|
||||
|
||||
def _compute_page_detection_matching(
|
||||
self,
|
||||
preds: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
height: int,
|
||||
width: int,
|
||||
top_k: int = 100,
|
||||
return_on_cpu: bool = True,
|
||||
) -> tuple:
|
||||
# Adapted from: https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/utils/detection_utils.py # noqa E501
|
||||
"""
|
||||
Match predictions (NMS output) and the targets (ground truth) with respect to metric
|
||||
and confidence score for a given image.
|
||||
|
||||
Args:
|
||||
preds: Tensor of shape (num_img_predictions, 6)
|
||||
format: (x1, y1, x2, y2, confidence, class_label)
|
||||
where x1,y1,x2,y2 are according to image size
|
||||
targets: targets for this image of shape (num_img_targets, 5)
|
||||
format: (label, x1, y1, x2, y2)
|
||||
where x1,y1,x2,y2 are according to image size
|
||||
height: dimensions of the image
|
||||
width: dimensions of the image
|
||||
top_k: Number of predictions to keep per class, ordered by confidence score
|
||||
return_on_cpu: If True, the output will be returned on "CPU", otherwise it will be
|
||||
returned on "device"
|
||||
|
||||
Returns:
|
||||
preds_matched: Tensor of shape (num_img_predictions, n_thresholds)
|
||||
True when prediction (i) is matched with a target with respect to
|
||||
the (j)th threshold
|
||||
preds_to_ignore: Tensor of shape (num_img_predictions, n_thresholds)
|
||||
True when prediction (i) is matched with a crowd target with
|
||||
respect to the (j)th threshold
|
||||
preds_scores: Tensor of shape (num_img_predictions),
|
||||
confidence score for every prediction
|
||||
preds_cls: Tensor of shape (num_img_predictions),
|
||||
predicted class for every prediction
|
||||
targets_cls: Tensor of shape (num_img_targets),
|
||||
ground truth class for every target
|
||||
"""
|
||||
thresholds = self.iou_thresholds.to(device=self.device)
|
||||
num_thresholds = len(thresholds)
|
||||
|
||||
if preds is None or len(preds) == 0:
|
||||
preds_matched = torch.zeros((0, num_thresholds), dtype=torch.bool, device=self.device)
|
||||
preds_to_ignore = torch.zeros((0, num_thresholds), dtype=torch.bool, device=self.device)
|
||||
preds_scores = torch.tensor([], dtype=torch.float32, device=self.device)
|
||||
preds_cls = torch.tensor([], dtype=torch.float32, device=self.device)
|
||||
targets_cls = targets[:, 0].to(device=self.device)
|
||||
return preds_matched, preds_to_ignore, preds_scores, preds_cls, targets_cls
|
||||
|
||||
preds_matched = torch.zeros(
|
||||
len(preds), num_thresholds, dtype=torch.bool, device=self.device
|
||||
)
|
||||
targets_matched = torch.zeros(
|
||||
len(targets), num_thresholds, dtype=torch.bool, device=self.device
|
||||
)
|
||||
preds_to_ignore = torch.zeros(
|
||||
len(preds), num_thresholds, dtype=torch.bool, device=self.device
|
||||
)
|
||||
|
||||
preds_cls, preds_box, preds_scores = preds[:, -1], preds[:, 0:4], preds[:, 4]
|
||||
targets_cls, targets_box = targets[:, 0], targets[:, 1:5]
|
||||
|
||||
# Ignore all but the predictions that were top_k for their class
|
||||
preds_idx_to_use = self._get_top_k_idx_per_cls(preds_scores, preds_cls, top_k)
|
||||
preds_to_ignore[:, :] = True
|
||||
preds_to_ignore[preds_idx_to_use] = False
|
||||
|
||||
if len(targets) > 0: # or len(crowd_targets) > 0:
|
||||
self._change_bbox_bounds_for_image_size(preds, (height, width))
|
||||
|
||||
preds_matched = self._compute_targets(
|
||||
preds_box,
|
||||
preds_cls,
|
||||
targets_box,
|
||||
targets_cls,
|
||||
preds_matched,
|
||||
targets_matched,
|
||||
preds_idx_to_use,
|
||||
thresholds,
|
||||
)
|
||||
|
||||
return preds_matched, preds_to_ignore, preds_scores, preds_cls, targets_cls
|
||||
|
||||
def _compute_detection_metrics(
|
||||
self,
|
||||
preds_matched: torch.Tensor,
|
||||
preds_to_ignore: torch.Tensor,
|
||||
preds_scores: torch.Tensor,
|
||||
preds_cls: torch.Tensor,
|
||||
targets_cls: torch.Tensor,
|
||||
) -> tuple:
|
||||
# Adapted from: https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/utils/detection_utils.py # noqa E501
|
||||
"""
|
||||
Compute the list of precision, recall, MaP and f1 for every class.
|
||||
|
||||
Args:
|
||||
preds_matched: Tensor of shape (num_predictions, n_iou_thresholds)
|
||||
True when prediction (i) is matched with a target with respect
|
||||
to the (j)th IoU threshold
|
||||
preds_to_ignore Tensor of shape (num_predictions, n_iou_thresholds)
|
||||
True when prediction (i) is matched with a crowd target with
|
||||
respect to the (j)th IoU threshold
|
||||
preds_scores: Tensor of shape (num_predictions),
|
||||
confidence score for every prediction
|
||||
preds_cls: Tensor of shape (num_predictions),
|
||||
predicted class for every prediction
|
||||
targets_cls: Tensor of shape (num_targets),
|
||||
ground truth class for every target box to be detected
|
||||
|
||||
Returns:
|
||||
ap, precision, recall, f1: Tensors of shape (n_class, nb_iou_thrs)
|
||||
unique_classes: Vector with all unique target classes
|
||||
"""
|
||||
|
||||
preds_matched, preds_to_ignore = preds_matched.to(self.device), preds_to_ignore.to(
|
||||
self.device
|
||||
)
|
||||
preds_scores, preds_cls, targets_cls = (
|
||||
preds_scores.to(self.device),
|
||||
preds_cls.to(self.device),
|
||||
targets_cls.to(self.device),
|
||||
)
|
||||
|
||||
recall_thresholds = self.recall_thresholds.to(self.device)
|
||||
score_threshold = self.score_threshold
|
||||
|
||||
unique_classes = torch.unique(targets_cls).long()
|
||||
|
||||
n_class, nb_iou_thrs = len(unique_classes), preds_matched.shape[-1]
|
||||
|
||||
ap = torch.zeros((n_class, nb_iou_thrs), device=self.device)
|
||||
precision = torch.zeros((n_class, nb_iou_thrs), device=self.device)
|
||||
recall = torch.zeros((n_class, nb_iou_thrs), device=self.device)
|
||||
|
||||
for cls_i, class_value in enumerate(unique_classes):
|
||||
cls_preds_idx, cls_targets_idx = (preds_cls == class_value), (
|
||||
targets_cls == class_value
|
||||
)
|
||||
(
|
||||
cls_ap,
|
||||
cls_precision,
|
||||
cls_recall,
|
||||
) = self._compute_detection_metrics_per_cls(
|
||||
preds_matched=preds_matched[cls_preds_idx],
|
||||
preds_to_ignore=preds_to_ignore[cls_preds_idx],
|
||||
preds_scores=preds_scores[cls_preds_idx],
|
||||
n_targets=cls_targets_idx.sum(),
|
||||
recall_thresholds=recall_thresholds,
|
||||
score_threshold=score_threshold,
|
||||
)
|
||||
ap[cls_i, :] = cls_ap
|
||||
precision[cls_i, :] = cls_precision
|
||||
recall[cls_i, :] = cls_recall
|
||||
|
||||
f1 = 2 * precision * recall / (precision + recall + 1e-16)
|
||||
return ap, precision, recall, f1, unique_classes
|
||||
|
||||
def _compute_detection_metrics_per_cls(
|
||||
self,
|
||||
preds_matched: torch.Tensor,
|
||||
preds_to_ignore: torch.Tensor,
|
||||
preds_scores: torch.Tensor,
|
||||
n_targets: int,
|
||||
recall_thresholds: torch.Tensor,
|
||||
score_threshold: float,
|
||||
):
|
||||
# Adapted from: https://github.com/Deci-AI/super-gradients/blob/master/src/super_gradients/training/utils/detection_utils.py # noqa E501
|
||||
"""
|
||||
Compute the list of precision, recall and MaP of a given class for every recall threshold.
|
||||
|
||||
Args:
|
||||
preds_matched: Tensor of shape (num_predictions, n_thresholds)
|
||||
True when prediction (i) is matched with a target
|
||||
with respect to the(j)th threshold
|
||||
preds_to_ignore Tensor of shape (num_predictions, n_thresholds)
|
||||
True when prediction (i) is matched with a crowd target
|
||||
with respect to the (j)th threshold
|
||||
preds_scores: Tensor of shape (num_predictions),
|
||||
confidence score for every prediction
|
||||
n_targets: Number of target boxes of this class
|
||||
recall_thresholds: Tensor of shape (max_n_rec_thresh)
|
||||
list of recall thresholds used to compute MaP
|
||||
score_threshold: Minimum confidence score to consider a prediction
|
||||
for the computation of precision and recall (not MaP)
|
||||
|
||||
Returns:
|
||||
ap, precision, recall: Tensors of shape (nb_thrs)
|
||||
"""
|
||||
|
||||
nb_iou_thrs = preds_matched.shape[-1]
|
||||
|
||||
tps = preds_matched
|
||||
fps = torch.logical_and(
|
||||
torch.logical_not(preds_matched), torch.logical_not(preds_to_ignore)
|
||||
)
|
||||
|
||||
if len(tps) == 0:
|
||||
return (
|
||||
torch.zeros(nb_iou_thrs, device=self.device),
|
||||
torch.zeros(nb_iou_thrs, device=self.device),
|
||||
torch.zeros(nb_iou_thrs, device=self.device),
|
||||
)
|
||||
|
||||
# Sort by decreasing score
|
||||
dtype = (
|
||||
torch.uint8
|
||||
if preds_scores.is_cuda and preds_scores.dtype is torch.bool
|
||||
else preds_scores.dtype
|
||||
)
|
||||
sort_ind = torch.argsort(preds_scores.to(dtype), descending=True)
|
||||
tps = tps[sort_ind, :]
|
||||
fps = fps[sort_ind, :]
|
||||
preds_scores = preds_scores[sort_ind].contiguous()
|
||||
|
||||
# Rolling sum over the predictions
|
||||
rolling_tps = torch.cumsum(tps, axis=0, dtype=torch.float)
|
||||
rolling_fps = torch.cumsum(fps, axis=0, dtype=torch.float)
|
||||
|
||||
rolling_recalls = rolling_tps / n_targets
|
||||
rolling_precisions = rolling_tps / (
|
||||
rolling_tps + rolling_fps + torch.finfo(torch.float64).eps
|
||||
)
|
||||
|
||||
# Reversed cummax to only have decreasing values
|
||||
rolling_precisions = rolling_precisions.flip(0).cummax(0).values.flip(0)
|
||||
|
||||
# ==================
|
||||
# RECALL & PRECISION
|
||||
|
||||
# We want the rolling precision/recall at index i so that:
|
||||
# preds_scores[i-1] >= score_threshold > preds_scores[i]
|
||||
# Note: torch.searchsorted works on increasing sequence and preds_scores is decreasing,
|
||||
# so we work with "-"
|
||||
# Note2: right=True due to negation
|
||||
lowest_score_above_threshold = torch.searchsorted(
|
||||
-preds_scores, -score_threshold, right=True
|
||||
)
|
||||
|
||||
if (
|
||||
lowest_score_above_threshold == 0
|
||||
): # Here score_threshold > preds_scores[0], so no pred is above the threshold
|
||||
recall = torch.zeros(nb_iou_thrs, device=self.device)
|
||||
precision = torch.zeros(
|
||||
nb_iou_thrs, device=self.device
|
||||
) # the precision is not really defined when no pred but we need to give it a value
|
||||
else:
|
||||
recall = rolling_recalls[lowest_score_above_threshold - 1]
|
||||
precision = rolling_precisions[lowest_score_above_threshold - 1]
|
||||
|
||||
# ==================
|
||||
# AVERAGE PRECISION
|
||||
|
||||
# shape = (nb_iou_thrs, n_recall_thresholds)
|
||||
recall_thresholds = recall_thresholds.view(1, -1).repeat(nb_iou_thrs, 1)
|
||||
|
||||
# We want the index i so that:
|
||||
# rolling_recalls[i-1] < recall_thresholds[k] <= rolling_recalls[i]
|
||||
# Note: when recall_thresholds[k] > max(rolling_recalls), i = len(rolling_recalls)
|
||||
# Note2: we work with transpose (.T) to apply torch.searchsorted on first dim
|
||||
# instead of the last one
|
||||
recall_threshold_idx = torch.searchsorted(
|
||||
rolling_recalls.T.contiguous(), recall_thresholds, right=False
|
||||
).T
|
||||
|
||||
# When recall_thresholds[k] > max(rolling_recalls),
|
||||
# rolling_precisions[i] is not defined, and we want precision = 0
|
||||
rolling_precisions = torch.cat(
|
||||
(rolling_precisions, torch.zeros(1, nb_iou_thrs, device=self.device)), dim=0
|
||||
)
|
||||
|
||||
# shape = (n_recall_thresholds, nb_iou_thrs)
|
||||
sampled_precision_points = torch.gather(
|
||||
input=rolling_precisions, index=recall_threshold_idx, dim=0
|
||||
)
|
||||
|
||||
# Average over the recall_thresholds
|
||||
ap = sampled_precision_points.mean(0)
|
||||
|
||||
return ap, precision, recall
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
from dataclasses import asdict
|
||||
|
||||
# Example usage
|
||||
prediction_file_paths = [Path("pths/to/predictions.json"), Path("pths/to/predictions2.json")]
|
||||
ground_truth_file_paths = [
|
||||
Path("pths/to/ground_truth.json"),
|
||||
Path("pths/to/ground_truth2.json"),
|
||||
]
|
||||
|
||||
for prediction_file_path, ground_truth_file_path in zip(
|
||||
prediction_file_paths, ground_truth_file_paths
|
||||
):
|
||||
eval_processor = ObjectDetectionEvalProcessor.from_json_files(
|
||||
prediction_file_path, ground_truth_file_path
|
||||
)
|
||||
|
||||
metrics, per_class_metrics = eval_processor.get_metrics()
|
||||
print(f"Metrics for {ground_truth_file_path.name}:\n{asdict(metrics)}")
|
||||
print(f"Per class Metrics for {ground_truth_file_path.name}:\n{asdict(per_class_metrics)}")
|
||||
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,180 @@
|
||||
import difflib
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from unstructured_inference.models.eval import compare_contents_as_df
|
||||
|
||||
|
||||
class TableAlignment:
|
||||
def __init__(self, cutoff: float = 0.8):
|
||||
self.cutoff = cutoff
|
||||
|
||||
@staticmethod
|
||||
def get_content_in_tables(table_data: List[List[Dict[str, Any]]]) -> List[str]:
|
||||
# Replace below docstring with google-style docstring
|
||||
"""Extracts and concatenates the content of cells from each table in a list of tables.
|
||||
|
||||
Args:
|
||||
table_data: A list of tables, each table being a list of cell data dictionaries.
|
||||
|
||||
Returns:
|
||||
List of strings where each string represents the concatenated content of one table.
|
||||
"""
|
||||
return [" ".join([d["content"] for d in td if "content" in d]) for td in table_data]
|
||||
|
||||
@staticmethod
|
||||
def get_table_level_alignment(
|
||||
predicted_table_data: List[List[Dict[str, Any]]],
|
||||
ground_truth_table_data: List[List[Dict[str, Any]]],
|
||||
) -> List[int]:
|
||||
"""Compares predicted table data with ground truth data to find the best
|
||||
matching table index for each predicted table.
|
||||
|
||||
Args:
|
||||
predicted_table_data: A list of predicted tables.
|
||||
ground_truth_table_data: A list of ground truth tables.
|
||||
|
||||
Returns:
|
||||
A list of indices indicating the best match in the ground truth for
|
||||
each predicted table.
|
||||
|
||||
"""
|
||||
ground_truth_texts = TableAlignment.get_content_in_tables(ground_truth_table_data)
|
||||
matched_indices = []
|
||||
for td in predicted_table_data:
|
||||
reference = TableAlignment.get_content_in_tables([td])[0]
|
||||
matches = difflib.get_close_matches(reference, ground_truth_texts, cutoff=0.1, n=1)
|
||||
matched_indices.append(ground_truth_texts.index(matches[0]) if matches else -1)
|
||||
return matched_indices
|
||||
|
||||
@staticmethod
|
||||
def _zip_to_dataframe(table_data: List[Dict[str, Any]]) -> pd.DataFrame:
|
||||
df = pd.DataFrame(table_data, columns=["row_index", "col_index", "content"])
|
||||
df = df.set_index("row_index")
|
||||
df["col_index"] = df["col_index"].astype(str)
|
||||
return df
|
||||
|
||||
@staticmethod
|
||||
def get_element_level_alignment(
|
||||
predicted_table_data: List[List[Dict[str, Any]]],
|
||||
ground_truth_table_data: List[List[Dict[str, Any]]],
|
||||
matched_indices: List[int],
|
||||
cutoff: float = 0.8,
|
||||
) -> Dict[str, float]:
|
||||
"""Aligns elements of the predicted tables with the ground truth tables at the cell level.
|
||||
|
||||
Args:
|
||||
predicted_table_data: A list of predicted tables.
|
||||
ground_truth_table_data: A list of ground truth tables.
|
||||
matched_indices: Indices of the best matching ground truth table for each predicted table.
|
||||
cutoff: The cutoff value for the close matches.
|
||||
|
||||
Returns:
|
||||
A dictionary with column and row alignment accuracies.
|
||||
|
||||
"""
|
||||
content_diff_cols = []
|
||||
content_diff_rows = []
|
||||
col_index_acc = []
|
||||
row_index_acc = []
|
||||
|
||||
for idx, td in zip(matched_indices, predicted_table_data):
|
||||
if idx == -1:
|
||||
content_diff_cols.append(0)
|
||||
content_diff_rows.append(0)
|
||||
col_index_acc.append(0)
|
||||
row_index_acc.append(0)
|
||||
continue
|
||||
ground_truth_td = ground_truth_table_data[idx]
|
||||
|
||||
# Get row and col content accuracy
|
||||
predict_table_df = TableAlignment._zip_to_dataframe(td)
|
||||
ground_truth_table_df = TableAlignment._zip_to_dataframe(ground_truth_td)
|
||||
|
||||
table_content_diff = compare_contents_as_df(
|
||||
ground_truth_table_df.fillna(""),
|
||||
predict_table_df.fillna(""),
|
||||
)
|
||||
content_diff_cols.append(table_content_diff["by_col_token_ratio"])
|
||||
content_diff_rows.append(table_content_diff["by_row_token_ratio"])
|
||||
|
||||
aligned_element_col_count = 0
|
||||
aligned_element_row_count = 0
|
||||
total_element_count = 0
|
||||
# Get row and col index accuracy
|
||||
ground_truth_td_contents_list = [gtd["content"].lower() for gtd in ground_truth_td]
|
||||
used_indices = set()
|
||||
indices_tuple_pairs = []
|
||||
for td_ele in td:
|
||||
content = td_ele["content"].lower()
|
||||
row_index = td_ele["row_index"]
|
||||
col_idx = td_ele["col_index"]
|
||||
|
||||
matches = difflib.get_close_matches(
|
||||
content,
|
||||
ground_truth_td_contents_list,
|
||||
cutoff=cutoff,
|
||||
n=1,
|
||||
)
|
||||
# BUG FIX: the previous matched_idx will only output the first matched index if
|
||||
# the match has duplicates in the
|
||||
# ground_truth_td_contents_list, the current fix will output its correspondence idx
|
||||
# once matching is exhausted, it will go back search again the same fashion
|
||||
matching_indices = []
|
||||
if matches != []:
|
||||
b_indices = [
|
||||
i
|
||||
for i, b_string in enumerate(ground_truth_td_contents_list)
|
||||
if b_string == matches[0] and i not in used_indices
|
||||
]
|
||||
if not b_indices:
|
||||
# If all indices are used, reset used_indices and use the first index
|
||||
used_indices.clear()
|
||||
b_indices = [
|
||||
i
|
||||
for i, b_string in enumerate(ground_truth_td_contents_list)
|
||||
if b_string == matches[0] and i not in used_indices
|
||||
]
|
||||
matching_index = b_indices[0]
|
||||
matching_indices.append(matching_index)
|
||||
used_indices.add(matching_index)
|
||||
else:
|
||||
matching_indices = [-1]
|
||||
matched_idx = matching_indices[0]
|
||||
if matched_idx >= 0:
|
||||
gt_row_index = ground_truth_td[matched_idx]["row_index"]
|
||||
gt_col_index = ground_truth_td[matched_idx]["col_index"]
|
||||
indices_tuple_pairs.append(((row_index, col_idx), (gt_row_index, gt_col_index)))
|
||||
|
||||
for indices_tuple_pair in indices_tuple_pairs:
|
||||
if indices_tuple_pair[0][0] == indices_tuple_pair[1][0]:
|
||||
aligned_element_row_count += 1
|
||||
if indices_tuple_pair[0][1] == indices_tuple_pair[1][1]:
|
||||
aligned_element_col_count += 1
|
||||
total_element_count += 1
|
||||
|
||||
table_col_index_acc = 0
|
||||
table_row_index_acc = 0
|
||||
if total_element_count > 0:
|
||||
table_col_index_acc = round(aligned_element_col_count / total_element_count, 2)
|
||||
table_row_index_acc = round(aligned_element_row_count / total_element_count, 2)
|
||||
|
||||
col_index_acc.append(table_col_index_acc)
|
||||
row_index_acc.append(table_row_index_acc)
|
||||
|
||||
not_found_gt_table_indexes = [
|
||||
id for id in range(len(ground_truth_table_data)) if id not in matched_indices
|
||||
]
|
||||
for _ in not_found_gt_table_indexes:
|
||||
content_diff_cols.append(0)
|
||||
content_diff_rows.append(0)
|
||||
col_index_acc.append(0)
|
||||
row_index_acc.append(0)
|
||||
|
||||
return {
|
||||
"col_index_acc": round(np.mean(col_index_acc), 2),
|
||||
"row_index_acc": round(np.mean(row_index_acc), 2),
|
||||
"col_content_acc": round(np.mean(content_diff_cols) / 100.0, 2),
|
||||
"row_content_acc": round(np.mean(content_diff_rows) / 100.0, 2),
|
||||
}
|
||||
@@ -0,0 +1,342 @@
|
||||
"""
|
||||
The purpose of this script is to create a comprehensive metric for table evaluation
|
||||
1. Verify table identification.
|
||||
a. Concatenate all text in the table and ground truth.
|
||||
b. Calculate the difference to find the closest matches.
|
||||
c. If contents are too different, mark as a failure.
|
||||
|
||||
2. For each identified table:
|
||||
a. Align elements at the level of individual elements.
|
||||
b. Match elements by text.
|
||||
c. Determine indexes for both predicted and actual data.
|
||||
d. Compare index tuples at column and row levels to assess content shifts.
|
||||
e. Compare the token orders by flattened along column and row levels
|
||||
f. Note: Imperfect HTML is acceptable unless it impedes parsing,
|
||||
in which case the table is considered failed.
|
||||
|
||||
Example
|
||||
python table_eval.py \
|
||||
--prediction_file "model_output.pdf.json" \
|
||||
--ground_truth_file "ground_truth.pdf.json"
|
||||
"""
|
||||
|
||||
import difflib
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import click
|
||||
import numpy as np
|
||||
|
||||
from unstructured.metrics.table.table_alignment import TableAlignment
|
||||
from unstructured.metrics.table.table_extraction import (
|
||||
extract_and_convert_tables_from_ground_truth,
|
||||
extract_and_convert_tables_from_prediction,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class TableEvaluation:
|
||||
"""Class representing a gathered table metrics."""
|
||||
|
||||
total_tables: int
|
||||
total_predicted_tables: int
|
||||
table_level_acc: float
|
||||
table_detection_recall: float
|
||||
table_detection_precision: float
|
||||
table_detection_f1: float
|
||||
element_col_level_index_acc: float
|
||||
element_row_level_index_acc: float
|
||||
element_col_level_content_acc: float
|
||||
element_row_level_content_acc: float
|
||||
|
||||
@property
|
||||
def composite_structure_acc(self) -> float:
|
||||
return (
|
||||
self.element_col_level_index_acc
|
||||
+ self.element_row_level_index_acc
|
||||
+ (self.element_col_level_content_acc + self.element_row_level_content_acc) / 2
|
||||
) / 3
|
||||
|
||||
|
||||
def table_level_acc(predicted_table_data, ground_truth_table_data, matched_indices):
|
||||
"""computes for each predicted table its accurary compared to ground truth.
|
||||
|
||||
The accuracy is defined as the SequenceMatcher.ratio() between those two strings. If a
|
||||
prediction does not have a matched ground truth its accuracy is 0
|
||||
"""
|
||||
score = np.zeros((len(matched_indices),))
|
||||
ground_truth_text = TableAlignment.get_content_in_tables(ground_truth_table_data)
|
||||
for idx, predicted in enumerate(predicted_table_data):
|
||||
matched_idx = matched_indices[idx]
|
||||
if matched_idx == -1:
|
||||
# false positive; default score 0
|
||||
continue
|
||||
score[idx] = difflib.SequenceMatcher(
|
||||
None,
|
||||
TableAlignment.get_content_in_tables([predicted])[0],
|
||||
ground_truth_text[matched_idx],
|
||||
).ratio()
|
||||
return score
|
||||
|
||||
|
||||
def _count_predicted_tables(matched_indices: List[int]) -> int:
|
||||
"""Counts the number of predicted tables that have a corresponding match in the ground truth.
|
||||
|
||||
Args:
|
||||
matched_indices: List of indices indicating matches between predicted
|
||||
and ground truth tables.
|
||||
|
||||
Returns:
|
||||
The count of matched predicted tables.
|
||||
|
||||
"""
|
||||
return sum(1 for idx in matched_indices if idx >= 0)
|
||||
|
||||
|
||||
def calculate_table_detection_metrics(
|
||||
matched_indices: list[int], ground_truth_tables_number: int
|
||||
) -> tuple[float, float, float]:
|
||||
"""
|
||||
Calculate the table detection metrics: recall, precision, and f1 score.
|
||||
Args:
|
||||
matched_indices:
|
||||
List of indices indicating matches between predicted and ground truth tables
|
||||
For example: matched_indices[i] = j means that the
|
||||
i-th predicted table is matched with the j-th ground truth table.
|
||||
ground_truth_tables_number: the number of ground truth tables.
|
||||
|
||||
Returns:
|
||||
Tuple of recall, precision, and f1 scores
|
||||
"""
|
||||
predicted_tables_number = len(matched_indices)
|
||||
|
||||
matched_set = set(matched_indices)
|
||||
if -1 in matched_set:
|
||||
matched_set.remove(-1)
|
||||
|
||||
true_positive = len(matched_set)
|
||||
false_positive = predicted_tables_number - true_positive
|
||||
positive = ground_truth_tables_number
|
||||
|
||||
recall = true_positive / positive if positive > 0 else 0
|
||||
precision = (
|
||||
true_positive / (true_positive + false_positive)
|
||||
if true_positive + false_positive > 0
|
||||
else 0
|
||||
)
|
||||
f1 = 2 * precision * recall / (precision + recall) if precision + recall > 0 else 0
|
||||
|
||||
return recall, precision, f1
|
||||
|
||||
|
||||
class TableEvalProcessor:
|
||||
def __init__(
|
||||
self,
|
||||
prediction: List[Dict[str, Any]],
|
||||
ground_truth: List[Dict[str, Any]],
|
||||
cutoff: float = 0.8,
|
||||
source_type: str = "html",
|
||||
):
|
||||
"""
|
||||
Initializes the TableEvalProcessor prediction and ground truth.
|
||||
|
||||
Args:
|
||||
ground_truth: Ground truth table data. The tables text should be in the deckerd format.
|
||||
prediction: Predicted table data.
|
||||
cutoff: The cutoff value for the element level alignment. Default is 0.8.
|
||||
|
||||
Examples:
|
||||
ground_truth: [
|
||||
{
|
||||
"type": "Table",
|
||||
"text": [
|
||||
{
|
||||
"id": "f4c35dae-105b-46f5-a77a-7fbc199d6aca",
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 1,
|
||||
"content": "Cell text"
|
||||
},
|
||||
...
|
||||
}
|
||||
]
|
||||
prediction: [
|
||||
{
|
||||
"element_id": <id_string>,
|
||||
...
|
||||
"metadata": {
|
||||
...
|
||||
"text_as_html": "<table><thead><tr><th rowspan=\"2\">June....
|
||||
</tr></td></table>",
|
||||
"table_as_cells":
|
||||
[
|
||||
{
|
||||
"x": 0,
|
||||
"y": 0,
|
||||
"w": 1,
|
||||
"h": 2,
|
||||
"content": "June"
|
||||
},
|
||||
...
|
||||
]
|
||||
}
|
||||
},
|
||||
]
|
||||
|
||||
"""
|
||||
self.prediction = prediction
|
||||
self.ground_truth = ground_truth
|
||||
self.cutoff = cutoff
|
||||
self.source_type = source_type
|
||||
|
||||
@classmethod
|
||||
def from_json_files(
|
||||
cls,
|
||||
prediction_file: Path,
|
||||
ground_truth_file: Path,
|
||||
cutoff: Optional[float] = None,
|
||||
source_type: str = "html",
|
||||
) -> "TableEvalProcessor":
|
||||
"""Factory classmethod to initialize the object with path to json files instead of dicts
|
||||
|
||||
Args:
|
||||
prediction_file: Path to the json file containing the predicted table data.
|
||||
ground_truth_file: Path to the json file containing the ground truth table data.
|
||||
source_type: 'cells' or 'html'. 'cells' refers to reading 'table_as_cells' field while
|
||||
'html' is extracted from 'text_as_html'
|
||||
cutoff: The cutoff value for the element level alignment.
|
||||
If not set, class default value is used (=0.8).
|
||||
|
||||
Returns:
|
||||
TableEvalProcessor: An instance of the class initialized with the provided data.
|
||||
"""
|
||||
with open(prediction_file) as f:
|
||||
prediction = json.load(f)
|
||||
with open(ground_truth_file) as f:
|
||||
ground_truth = json.load(f)
|
||||
if cutoff is not None:
|
||||
return cls(
|
||||
prediction=prediction,
|
||||
ground_truth=ground_truth,
|
||||
cutoff=cutoff,
|
||||
source_type=source_type,
|
||||
)
|
||||
else:
|
||||
return cls(prediction=prediction, ground_truth=ground_truth, source_type=source_type)
|
||||
|
||||
def process_file(self) -> TableEvaluation:
|
||||
"""Processes the files and computes table-level and element-level accuracy.
|
||||
|
||||
Returns:
|
||||
TableEvaluation: A dataclass object containing the computed metrics.
|
||||
"""
|
||||
ground_truth_table_data = extract_and_convert_tables_from_ground_truth(
|
||||
self.ground_truth,
|
||||
)
|
||||
|
||||
predicted_table_data = extract_and_convert_tables_from_prediction(
|
||||
file_elements=self.prediction, source_type=self.source_type
|
||||
)
|
||||
is_table_in_gt = bool(ground_truth_table_data)
|
||||
is_table_predicted = bool(predicted_table_data)
|
||||
if not is_table_in_gt:
|
||||
# There is no table data in ground truth, you either got perfect score or 0
|
||||
score = 0 if is_table_predicted else np.nan
|
||||
table_acc = 1 if not is_table_predicted else 0
|
||||
return TableEvaluation(
|
||||
total_tables=0,
|
||||
total_predicted_tables=len(predicted_table_data),
|
||||
table_level_acc=table_acc,
|
||||
table_detection_recall=score,
|
||||
table_detection_precision=score,
|
||||
table_detection_f1=score,
|
||||
element_col_level_index_acc=score,
|
||||
element_row_level_index_acc=score,
|
||||
element_col_level_content_acc=score,
|
||||
element_row_level_content_acc=score,
|
||||
)
|
||||
if is_table_in_gt and not is_table_predicted:
|
||||
return TableEvaluation(
|
||||
total_tables=len(ground_truth_table_data),
|
||||
total_predicted_tables=0,
|
||||
table_level_acc=0,
|
||||
table_detection_recall=0,
|
||||
table_detection_precision=0,
|
||||
table_detection_f1=0,
|
||||
element_col_level_index_acc=0,
|
||||
element_row_level_index_acc=0,
|
||||
element_col_level_content_acc=0,
|
||||
element_row_level_content_acc=0,
|
||||
)
|
||||
else:
|
||||
# We have both ground truth tables and predicted tables
|
||||
matched_indices = TableAlignment.get_table_level_alignment(
|
||||
predicted_table_data,
|
||||
ground_truth_table_data,
|
||||
)
|
||||
predicted_table_acc = np.mean(
|
||||
table_level_acc(predicted_table_data, ground_truth_table_data, matched_indices)
|
||||
)
|
||||
|
||||
metrics = TableAlignment.get_element_level_alignment(
|
||||
predicted_table_data,
|
||||
ground_truth_table_data,
|
||||
matched_indices,
|
||||
cutoff=self.cutoff,
|
||||
)
|
||||
|
||||
(
|
||||
table_detection_recall,
|
||||
table_detection_precision,
|
||||
table_detection_f1,
|
||||
) = calculate_table_detection_metrics(
|
||||
matched_indices=matched_indices,
|
||||
ground_truth_tables_number=len(ground_truth_table_data),
|
||||
)
|
||||
|
||||
evaluation = TableEvaluation(
|
||||
total_tables=len(ground_truth_table_data),
|
||||
total_predicted_tables=len(predicted_table_data),
|
||||
table_level_acc=predicted_table_acc,
|
||||
table_detection_recall=table_detection_recall,
|
||||
table_detection_precision=table_detection_precision,
|
||||
table_detection_f1=table_detection_f1,
|
||||
element_col_level_index_acc=metrics.get("col_index_acc", 0),
|
||||
element_row_level_index_acc=metrics.get("row_index_acc", 0),
|
||||
element_col_level_content_acc=metrics.get("col_content_acc", 0),
|
||||
element_row_level_content_acc=metrics.get("row_content_acc", 0),
|
||||
)
|
||||
return evaluation
|
||||
|
||||
|
||||
@click.command()
|
||||
@click.option(
|
||||
"--prediction_file", help="Path to the model prediction JSON file", type=click.Path(exists=True)
|
||||
)
|
||||
@click.option(
|
||||
"--ground_truth_file", help="Path to the ground truth JSON file", type=click.Path(exists=True)
|
||||
)
|
||||
@click.option(
|
||||
"--cutoff",
|
||||
type=float,
|
||||
show_default=True,
|
||||
default=0.8,
|
||||
help="The cutoff value for the element level alignment. \
|
||||
If not set, a default value is used",
|
||||
)
|
||||
def run(prediction_file: str, ground_truth_file: str, cutoff: Optional[float]):
|
||||
"""Runs the table evaluation process and prints the computed metrics."""
|
||||
processor = TableEvalProcessor.from_json_files(
|
||||
Path(prediction_file),
|
||||
Path(ground_truth_file),
|
||||
cutoff=cutoff,
|
||||
)
|
||||
report = processor.process_file()
|
||||
print(report)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
run()
|
||||
@@ -0,0 +1,288 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from bs4 import BeautifulSoup
|
||||
from unstructured_inference.models.tables import cells_to_html
|
||||
|
||||
EMPTY_CELL = {
|
||||
"row_index": "",
|
||||
"col_index": "",
|
||||
"content": "",
|
||||
}
|
||||
|
||||
|
||||
def _move_cells_for_spanned_cells(cells: List[Dict[str, Any]]):
|
||||
"""Move cells to the right if spanned cells have an influence on the rendering.
|
||||
|
||||
Args:
|
||||
cells: List of cells in the table in Deckerd format.
|
||||
|
||||
Returns:
|
||||
List of cells in the table in Deckerd format with cells moved to the right if spanned.
|
||||
"""
|
||||
sorted_cells = sorted(cells, key=lambda x: (x["y"], x["x"]))
|
||||
cells_occupied_by_spanned = set()
|
||||
for cell in sorted_cells:
|
||||
if cell["w"] > 1 or cell["h"] > 1:
|
||||
for i in range(cell["y"], cell["y"] + cell["h"]):
|
||||
for j in range(cell["x"], cell["x"] + cell["w"]):
|
||||
if (i, j) != (cell["y"], cell["x"]):
|
||||
cells_occupied_by_spanned.add((i, j))
|
||||
while (cell["y"], cell["x"]) in cells_occupied_by_spanned:
|
||||
cell_y, cell_x = cell["y"], cell["x"]
|
||||
cells_to_the_right = [c for c in sorted_cells if c["y"] == cell_y and c["x"] >= cell_x]
|
||||
for cell_to_move in cells_to_the_right:
|
||||
cell_to_move["x"] += 1
|
||||
cells_occupied_by_spanned.remove((cell_y, cell_x))
|
||||
return sorted_cells
|
||||
|
||||
|
||||
def html_table_to_deckerd(content: str) -> List[Dict[str, Any]]:
|
||||
"""Convert html format to Deckerd table structure.
|
||||
|
||||
Args:
|
||||
content: The html content with a table to extract.
|
||||
|
||||
Returns:
|
||||
A list of dictionaries where each dictionary represents a cell in the table.
|
||||
"""
|
||||
|
||||
soup = BeautifulSoup(content, "html.parser")
|
||||
table = soup.find("table")
|
||||
rows = table.find_all(["tr"])
|
||||
table_data = []
|
||||
|
||||
for i, row in enumerate(rows):
|
||||
cells = row.find_all(["th", "td"])
|
||||
for j, cell_data in enumerate(cells):
|
||||
cell = {
|
||||
"y": i,
|
||||
"x": j,
|
||||
"w": int(cell_data.attrs.get("colspan", 1)),
|
||||
"h": int(cell_data.attrs.get("rowspan", 1)),
|
||||
"content": cell_data.text,
|
||||
}
|
||||
table_data.append(cell)
|
||||
return _move_cells_for_spanned_cells(table_data)
|
||||
|
||||
|
||||
def deckerd_table_to_html(cells: List[Dict[str, Any]]) -> str:
|
||||
"""Convert Deckerd table structure to html format.
|
||||
|
||||
Args:
|
||||
cells: List of dictionaries where each dictionary represents a cell in the table.
|
||||
|
||||
Returns:
|
||||
A string with the html content of the table.
|
||||
"""
|
||||
transformer_cells = []
|
||||
# determine which cells are in header. Consider row 0 as header
|
||||
# but spans may make it larger
|
||||
first_row_cells = [cell for cell in cells if cell["y"] == 0]
|
||||
header_length = max(cell["w"] for cell in first_row_cells)
|
||||
header_rows = set(range(header_length))
|
||||
for cell in cells:
|
||||
cell_data = {
|
||||
"row_nums": list(range(cell["y"], cell["y"] + cell["h"])),
|
||||
"column_nums": list(range(cell["x"], cell["x"] + cell["w"])),
|
||||
"w": cell["w"],
|
||||
"h": cell["h"],
|
||||
"cell text": cell["content"],
|
||||
"column header": cell["y"] in header_rows,
|
||||
}
|
||||
transformer_cells.append(cell_data)
|
||||
# reuse the existing function to convert to HTML
|
||||
table = cells_to_html(transformer_cells)
|
||||
return table
|
||||
|
||||
|
||||
def _convert_table_from_html(content: str) -> List[Dict[str, Any]]:
|
||||
"""Convert html format to table structure. As a middle step it converts
|
||||
html to the Deckerd format as it's more convenient to work with.
|
||||
|
||||
Args:
|
||||
content: The html content with a table to extract.
|
||||
|
||||
Returns:
|
||||
A list of dictionaries where each dictionary represents a cell in the table.
|
||||
"""
|
||||
deckerd_cells = html_table_to_deckerd(content)
|
||||
return _convert_table_from_deckerd(deckerd_cells)
|
||||
|
||||
|
||||
def _convert_table_from_deckerd(content: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
|
||||
"""Convert deckerd format to table structure.
|
||||
|
||||
Args:
|
||||
content: The deckerd formatted content with a table to extract.
|
||||
|
||||
Returns:
|
||||
A list of dictionaries where each dictionary represents a cell in the table.
|
||||
"""
|
||||
table_data = []
|
||||
for table in content:
|
||||
try:
|
||||
cell_data = {
|
||||
"row_index": table["y"],
|
||||
"col_index": table["x"],
|
||||
"content": table["content"],
|
||||
}
|
||||
except KeyError:
|
||||
cell_data = EMPTY_CELL
|
||||
except TypeError:
|
||||
cell_data = EMPTY_CELL
|
||||
table_data.append(cell_data)
|
||||
return table_data
|
||||
|
||||
|
||||
def _sort_table_cells(table_data: List[List[Dict[str, Any]]]) -> List[List[Dict[str, Any]]]:
|
||||
return sorted(table_data, key=lambda cell: (cell["row_index"], cell["col_index"]))
|
||||
|
||||
|
||||
def extract_and_convert_tables_from_ground_truth(
|
||||
file_elements: List[Dict[str, Any]],
|
||||
) -> List[List[Dict[str, Any]]]:
|
||||
"""Extracts and converts tables data to a structured format based on the specified table type.
|
||||
|
||||
Args:
|
||||
file_elements: List of elements from the ground truth file.
|
||||
|
||||
Returns:
|
||||
A list of tables with each table represented as a list of cell data dictionaries.
|
||||
|
||||
"""
|
||||
ground_truth_table_data = []
|
||||
for element in file_elements:
|
||||
if "type" in element and element["type"] == "Table" and "text" in element:
|
||||
try:
|
||||
converted_data = _convert_table_from_deckerd(
|
||||
element["text"],
|
||||
)
|
||||
ground_truth_table_data.append(_sort_table_cells(converted_data))
|
||||
except Exception as e:
|
||||
print(f"Error converting ground truth data: {e}")
|
||||
ground_truth_table_data.append({})
|
||||
|
||||
return ground_truth_table_data
|
||||
|
||||
|
||||
def extract_and_convert_tables_from_prediction(
|
||||
file_elements: List[Dict[str, Any]], source_type: str = "html"
|
||||
) -> List[List[Dict[str, Any]]]:
|
||||
"""Extracts and converts table data to a structured format
|
||||
|
||||
Args:
|
||||
file_elements: List of elements from the file.
|
||||
source_type: 'cells' or 'html'. 'cells' refers to reading 'table_as_cells' field while
|
||||
'html' is extracted from 'text_as_html'
|
||||
|
||||
Returns:
|
||||
A list of tables with each table represented as a list of cell data dictionaries.
|
||||
|
||||
"""
|
||||
source_type_to_extraction_strategies = {
|
||||
"html": extract_cells_from_text_as_html,
|
||||
"cells": extract_cells_from_table_as_cells,
|
||||
}
|
||||
if source_type not in source_type_to_extraction_strategies:
|
||||
raise ValueError(
|
||||
f'source_type {source_type} is not valid. Allowed source_types are "html" and "cells"'
|
||||
)
|
||||
|
||||
extract_cells_fn = source_type_to_extraction_strategies[source_type]
|
||||
fallback_extract_cells_fn = (
|
||||
extract_cells_from_table_as_cells
|
||||
if source_type == "cells"
|
||||
else extract_cells_from_text_as_html
|
||||
)
|
||||
|
||||
predicted_table_data = []
|
||||
for element in file_elements:
|
||||
if element.get("type") == "Table":
|
||||
extracted_cells = extract_cells_fn(element)
|
||||
if not extracted_cells:
|
||||
extracted_cells = fallback_extract_cells_fn(element)
|
||||
if extracted_cells:
|
||||
sorted_cells = _sort_table_cells(extracted_cells)
|
||||
predicted_table_data.append(sorted_cells)
|
||||
|
||||
return predicted_table_data
|
||||
|
||||
|
||||
def extract_cells_from_text_as_html(element: Dict[str, Any]) -> List[Dict[str, Any]] | None:
|
||||
"""Extracts and parse cells from "text_as_html" field in Element structure
|
||||
|
||||
Args:
|
||||
element: Example element:
|
||||
{
|
||||
"type": "Table",
|
||||
"metadata": {
|
||||
"text_as_html": "<table>
|
||||
<thead>
|
||||
<tr>
|
||||
<th>Month A.</th>
|
||||
</tr>
|
||||
</thead>
|
||||
</tbody>
|
||||
<tr>
|
||||
<td>22</td><
|
||||
</tr>
|
||||
</tbody>
|
||||
</table>"
|
||||
}
|
||||
}
|
||||
|
||||
Returns:
|
||||
List of extracted cells in a format:
|
||||
[
|
||||
{
|
||||
"row_index": 0,
|
||||
"col_index": 0,
|
||||
"content": "Month A.",
|
||||
},
|
||||
...,
|
||||
]
|
||||
"""
|
||||
val = element["metadata"].get("text_as_html")
|
||||
if not val or "<table>" not in val:
|
||||
return None
|
||||
|
||||
predicted_cells = None
|
||||
try:
|
||||
predicted_cells = _convert_table_from_html(val)
|
||||
except Exception as e:
|
||||
print(f"Error converting Unstructured table data: {e}")
|
||||
|
||||
return predicted_cells
|
||||
|
||||
|
||||
def extract_cells_from_table_as_cells(element: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
"""Extracts and parse cells from "table_as_cells" field in Element structure
|
||||
|
||||
Args:
|
||||
element: Example element:
|
||||
{
|
||||
"type": "Table",
|
||||
"metadata": {
|
||||
"table_as_cells": [{"x": 0, "y": 0, "w": 1, "h": 1, "content": "Month A."},
|
||||
{"x": 0, "y": 1, "w": 1, "h": 1, "content": "22"}]
|
||||
}
|
||||
}
|
||||
|
||||
Returns:
|
||||
List of extracted cells in a format:
|
||||
[
|
||||
{
|
||||
"row_index": 0,
|
||||
"col_index": 0,
|
||||
"content": "Month A.",
|
||||
},
|
||||
...,
|
||||
]
|
||||
"""
|
||||
predicted_cells = element["metadata"].get("table_as_cells")
|
||||
converted_cells = None
|
||||
if predicted_cells:
|
||||
converted_cells = _convert_table_from_deckerd(predicted_cells)
|
||||
return converted_cells
|
||||
@@ -0,0 +1,49 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Union
|
||||
|
||||
|
||||
@dataclass
|
||||
class SimpleTableCell:
|
||||
x: int
|
||||
y: int
|
||||
w: int
|
||||
h: int
|
||||
content: str = ""
|
||||
|
||||
def to_dict(self):
|
||||
return {
|
||||
"x": self.x,
|
||||
"y": self.y,
|
||||
"w": self.w,
|
||||
"h": self.h,
|
||||
"content": self.content,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def from_table_transformer_cell(cls, tatr_table_cell: dict[str, Union[list[int], str]]):
|
||||
"""
|
||||
Args:
|
||||
tatr_table_cell (dict):
|
||||
Cell in a format returned by Table Transformer model, for example:
|
||||
{
|
||||
"row_nums": [1,2,3],
|
||||
"column_nums": [2],
|
||||
"cell text": "Text inside cell"
|
||||
}
|
||||
"""
|
||||
|
||||
row_nums = tatr_table_cell.get("row_nums", [])
|
||||
column_nums = tatr_table_cell.get("column_nums", [])
|
||||
|
||||
if not row_nums:
|
||||
raise ValueError(f'Cell {tatr_table_cell} has missing values under "row_nums" key')
|
||||
if not column_nums:
|
||||
raise ValueError(f'Cell {tatr_table_cell} has missing values under "column_nums" key')
|
||||
|
||||
return cls(
|
||||
x=min(column_nums),
|
||||
y=min(row_nums),
|
||||
w=len(column_nums),
|
||||
h=len(row_nums),
|
||||
content=tatr_table_cell.get("cell text", ""),
|
||||
)
|
||||
@@ -0,0 +1,48 @@
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
from PIL import Image
|
||||
|
||||
from unstructured.partition.pdf import convert_pdf_to_images
|
||||
from unstructured.partition.pdf_image.ocr import get_table_tokens
|
||||
from unstructured.partition.utils.ocr_models.ocr_interface import OCRAgent
|
||||
from unstructured.utils import requires_dependencies
|
||||
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def image_or_pdf_to_dataframe(filename: str) -> pd.DataFrame:
|
||||
"""helper to JUST run table transformer on the input image/pdf file. It assumes the input is
|
||||
JUST a table. This is intended to facilitate metric tracking on table structure detection ALONE
|
||||
without mixing metric of element detection model"""
|
||||
from unstructured_inference.models.tables import load_agent, tables_agent
|
||||
|
||||
load_agent()
|
||||
|
||||
if filename.endswith(".pdf"):
|
||||
image = list(convert_pdf_to_images(filename))[0].convert("RGB")
|
||||
else:
|
||||
image = Image.open(filename).convert("RGB")
|
||||
|
||||
ocr_agent = OCRAgent.get_agent(language="eng")
|
||||
|
||||
return tables_agent.run_prediction(
|
||||
image, ocr_tokens=get_table_tokens(image, ocr_agent), result_format="dataframe"
|
||||
)
|
||||
|
||||
|
||||
@requires_dependencies("unstructured_inference")
|
||||
def eval_table_transformer_for_file(
|
||||
filename: str,
|
||||
true_table_filename: str,
|
||||
eval_func: str = "token_ratio",
|
||||
) -> float:
|
||||
"""evaluate the predicted table structure vs. actual table structure by column and row as a
|
||||
number between 0 and 1"""
|
||||
from unstructured_inference.models.eval import compare_contents_as_df
|
||||
|
||||
pred_table = image_or_pdf_to_dataframe(filename).fillna("").replace(np.nan, "")
|
||||
actual_table = pd.read_csv(true_table_filename).astype(str).fillna("").replace(np.nan, "")
|
||||
|
||||
results = np.array(
|
||||
list(compare_contents_as_df(actual_table, pred_table, eval_func=eval_func).values()),
|
||||
)
|
||||
return results.mean() / 100.0
|
||||
@@ -0,0 +1,251 @@
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
from rapidfuzz.distance import Levenshtein
|
||||
|
||||
from unstructured.cleaners.core import clean_bullets, remove_sentence_punctuation
|
||||
|
||||
|
||||
def calculate_accuracy(
|
||||
output: Optional[str],
|
||||
source: Optional[str],
|
||||
weights: Tuple[int, int, int] = (2, 1, 1),
|
||||
) -> float:
|
||||
"""
|
||||
Calculates accuracy by calling calculate_edit_distance function using `return_as=score`.
|
||||
The function will return complement of the edit distance instead.
|
||||
"""
|
||||
return calculate_edit_distance(output, source, weights, return_as="score")
|
||||
|
||||
|
||||
def calculate_edit_distance(
|
||||
output: Optional[str],
|
||||
source: Optional[str],
|
||||
weights: Tuple[int, int, int] = (2, 1, 1),
|
||||
return_as: str = "distance",
|
||||
standardize_whitespaces: bool = True,
|
||||
) -> float:
|
||||
"""
|
||||
Calculates edit distance using Levenshtein distance between two strings.
|
||||
|
||||
Args:
|
||||
output (str): The target string to be compared.
|
||||
source (str): The reference string against which 'output' is compared.
|
||||
weights (Tuple[int, int, int], optional): A tuple containing weights
|
||||
for insertion, deletion, and substitution operations in the edit
|
||||
distance calculation. Default is (2, 1, 1).
|
||||
return_as (str, optional): The type of result to return, one of
|
||||
["score", "distance"].
|
||||
Default is "distance".
|
||||
|
||||
Returns:
|
||||
float: The calculated edit distance or similarity score between
|
||||
the 'output' and 'source' strings.
|
||||
|
||||
Raises:
|
||||
ValueError: If 'return_as' is not one of the valid return types
|
||||
["score", "distance"].
|
||||
|
||||
Note:
|
||||
This function calculates the edit distance (or similarity score) between
|
||||
two strings using the Levenshtein distance algorithm. The 'weights' parameter
|
||||
allows customizing the cost of insertion, deletion, and substitution
|
||||
operations. The 'return_as' parameter determines the type of result to return:
|
||||
- "score": Returns the similarity score, where 1.0 indicates a perfect match.
|
||||
- "distance": Returns the raw edit distance value.
|
||||
|
||||
"""
|
||||
return_types = ["score", "distance"]
|
||||
if return_as not in return_types:
|
||||
raise ValueError("Invalid return value type. Expected one of: %s" % return_types)
|
||||
output = standardize_quotes(prepare_str(output, standardize_whitespaces))
|
||||
source = standardize_quotes(prepare_str(source, standardize_whitespaces))
|
||||
distance = Levenshtein.distance(output, source, weights=weights) # type: ignore
|
||||
# lower bounded the char length for source string at 1.0 because to avoid division by zero
|
||||
# in the case where source string is empty, the distance should be at 100%
|
||||
source_char_len = max(len(source), 1.0) # type: ignore
|
||||
bounded_percentage_distance = min(max(distance / source_char_len, 0.0), 1.0)
|
||||
if return_as == "score":
|
||||
return 1 - bounded_percentage_distance
|
||||
elif return_as == "distance":
|
||||
return distance
|
||||
return 0.0
|
||||
|
||||
|
||||
def bag_of_words(text: str) -> Dict[str, int]:
|
||||
"""
|
||||
Outputs the bag of words (BOW) found in the input text and their frequencies.
|
||||
|
||||
Takes "clean, concatenated text" (CCT) from a document as input.
|
||||
|
||||
Removes sentence punctuation, but not punctuation within a word (ex. apostrophes).
|
||||
"""
|
||||
bow: Dict[str, int] = {}
|
||||
incorrect_word: str = ""
|
||||
words = clean_bullets(remove_sentence_punctuation(text.lower(), ["-", "'"])).split()
|
||||
|
||||
i = 0
|
||||
while i < len(words):
|
||||
if len(words[i]) > 1:
|
||||
if words[i] in bow:
|
||||
bow[words[i]] += 1
|
||||
else:
|
||||
bow[words[i]] = 1
|
||||
i += 1
|
||||
else:
|
||||
j = i
|
||||
incorrect_word = ""
|
||||
|
||||
while j < len(words) and len(words[j]) == 1:
|
||||
incorrect_word += words[j]
|
||||
j += 1
|
||||
|
||||
if len(incorrect_word) == 1 and words[i].isalnum():
|
||||
if incorrect_word in bow:
|
||||
bow[incorrect_word] += 1
|
||||
else:
|
||||
bow[incorrect_word] = 1
|
||||
i = j
|
||||
return bow
|
||||
|
||||
|
||||
def calculate_percent_missing_text(
|
||||
output: Optional[str],
|
||||
source: Optional[str],
|
||||
) -> float:
|
||||
"""
|
||||
Creates the bag of words (BOW) found in each input text and their frequencies, then compares the
|
||||
output BOW against the source BOW to calculate the % of text from the source text missing from
|
||||
the output text.
|
||||
|
||||
Takes "clean, concatenated text" (CCT) from a document output and the ground truth source text
|
||||
as inputs.
|
||||
|
||||
If the output text contains all words from the source text and then some extra, result will be
|
||||
0% missing text - this calculation does not penalize duplication.
|
||||
|
||||
A spaced-out word (ex. h e l l o) is considered missing; individual characters of a word
|
||||
will not be counted as separate words.
|
||||
|
||||
Returns the percentage of missing text represented as a decimal between 0 and 1.
|
||||
"""
|
||||
output = prepare_str(output)
|
||||
source = prepare_str(source)
|
||||
output_bow = bag_of_words(output)
|
||||
source_bow = bag_of_words(source)
|
||||
|
||||
# get total words in source bow while counting missing words
|
||||
total_source_word_count = 0
|
||||
total_missing_word_count = 0
|
||||
|
||||
for source_word, source_count in source_bow.items():
|
||||
total_source_word_count += source_count
|
||||
if source_word not in output_bow:
|
||||
# entire count is missing
|
||||
total_missing_word_count += source_count
|
||||
else:
|
||||
output_count = output_bow[source_word]
|
||||
total_missing_word_count += max(source_count - output_count, 0)
|
||||
|
||||
# calculate percent missing text
|
||||
if total_source_word_count == 0:
|
||||
return 0 # nothing missing because nothing in source document
|
||||
|
||||
fraction_missing = round(total_missing_word_count / total_source_word_count, 3)
|
||||
return min(fraction_missing, 1) # limit to 100%
|
||||
|
||||
|
||||
def prepare_str(string: Optional[str], standardize_whitespaces: bool = False) -> str:
|
||||
if not string:
|
||||
return ""
|
||||
if standardize_whitespaces:
|
||||
return " ".join(string.split())
|
||||
return str(string) # type: ignore
|
||||
|
||||
|
||||
def standardize_quotes(text: str) -> str:
|
||||
"""
|
||||
Converts all unicode quotes to standard ASCII quotes with comprehensive coverage.
|
||||
|
||||
Args:
|
||||
text (str): The input text to be standardized.
|
||||
|
||||
Returns:
|
||||
str: The text with standardized quotes.
|
||||
"""
|
||||
# Double Quotes Dictionary
|
||||
double_quotes = {
|
||||
'"': "U+0022", # noqa 601 # Standard typewriter/programmer's quote
|
||||
'"': "U+201C", # noqa 601 # Left double quotation mark
|
||||
'"': "U+201D", # noqa 601 # Right double quotation mark
|
||||
"„": "U+201E", # Double low-9 quotation mark
|
||||
"‟": "U+201F", # Double high-reversed-9 quotation mark
|
||||
"«": "U+00AB", # Left-pointing double angle quotation mark
|
||||
"»": "U+00BB", # Right-pointing double angle quotation mark
|
||||
"❝": "U+275D", # Heavy double turned comma quotation mark ornament
|
||||
"❞": "U+275E", # Heavy double comma quotation mark ornament
|
||||
"⹂": "U+2E42", # Double low-reversed-9 quotation mark
|
||||
"🙶": "U+1F676", # SANS-SERIF HEAVY DOUBLE TURNED COMMA QUOTATION MARK ORNAMENT
|
||||
"🙷": "U+1F677", # SANS-SERIF HEAVY DOUBLE COMMA QUOTATION MARK ORNAMENT
|
||||
"🙸": "U+1F678", # SANS-SERIF HEAVY LOW DOUBLE COMMA QUOTATION MARK ORNAMENT
|
||||
"⠦": "U+2826", # Braille double closing quotation mark
|
||||
"⠴": "U+2834", # Braille double opening quotation mark
|
||||
"〝": "U+301D", # REVERSED DOUBLE PRIME QUOTATION MARK
|
||||
"〞": "U+301E", # DOUBLE PRIME QUOTATION MARK
|
||||
"〟": "U+301F", # LOW DOUBLE PRIME QUOTATION MARK
|
||||
"＂": "U+FF02", # FULLWIDTH QUOTATION MARK
|
||||
",,": "U+275E", # LOW HEAVY DOUBLE COMMA ORNAMENT
|
||||
}
|
||||
|
||||
# Single Quotes Dictionary
|
||||
single_quotes = {
|
||||
"'": "U+0027", # noqa 601 # Standard typewriter/programmer's quote
|
||||
"'": "U+2018", # noqa 601 # Left single quotation mark
|
||||
"'": "U+2019", # noqa 601 # Right single quotation mark # noqa: W605
|
||||
"‚": "U+201A", # Single low-9 quotation mark
|
||||
"‛": "U+201B", # Single high-reversed-9 quotation mark
|
||||
"‹": "U+2039", # Single left-pointing angle quotation mark
|
||||
"›": "U+203A", # Single right-pointing angle quotation mark
|
||||
"❛": "U+275B", # Heavy single turned comma quotation mark ornament
|
||||
"❜": "U+275C", # Heavy single comma quotation mark ornament
|
||||
"「": "U+300C", # Left corner bracket
|
||||
"」": "U+300D", # Right corner bracket
|
||||
"『": "U+300E", # Left white corner bracket
|
||||
"』": "U+300F", # Right white corner bracket
|
||||
"﹁": "U+FE41", # PRESENTATION FORM FOR VERTICAL LEFT CORNER BRACKET
|
||||
"﹂": "U+FE42", # PRESENTATION FORM FOR VERTICAL RIGHT CORNER BRACKET
|
||||
"﹃": "U+FE43", # PRESENTATION FORM FOR VERTICAL LEFT WHITE CORNER BRACKET
|
||||
"﹄": "U+FE44", # PRESENTATION FORM FOR VERTICAL RIGHT WHITE CORNER BRACKET
|
||||
"＇": "U+FF07", # FULLWIDTH APOSTROPHE
|
||||
"｢": "U+FF62", # HALFWIDTH LEFT CORNER BRACKET
|
||||
"｣": "U+FF63", # HALFWIDTH RIGHT CORNER BRACKET
|
||||
}
|
||||
|
||||
double_quote_standard = '"'
|
||||
single_quote_standard = "'"
|
||||
|
||||
# Apply double quote replacements
|
||||
for unicode_val in double_quotes.values():
|
||||
unicode_char = unicode_to_char(unicode_val)
|
||||
if unicode_char in text:
|
||||
text = text.replace(unicode_char, double_quote_standard)
|
||||
|
||||
# Apply single quote replacements
|
||||
for unicode_val in single_quotes.values():
|
||||
unicode_char = unicode_to_char(unicode_val)
|
||||
if unicode_char in text:
|
||||
text = text.replace(unicode_char, single_quote_standard)
|
||||
|
||||
return text
|
||||
|
||||
|
||||
def unicode_to_char(unicode_val: str) -> str:
|
||||
"""
|
||||
Converts a Unicode value to a character.
|
||||
|
||||
Args:
|
||||
unicode_val (str): The Unicode value to convert.
|
||||
|
||||
Returns:
|
||||
str: The character corresponding to the Unicode value.
|
||||
"""
|
||||
return chr(int(unicode_val.replace("U+", ""), 16))
|
||||
@@ -0,0 +1,246 @@
|
||||
import logging
|
||||
import os
|
||||
import re
|
||||
import statistics
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import click
|
||||
import pandas as pd
|
||||
|
||||
from unstructured.staging.base import elements_from_json, elements_to_text
|
||||
|
||||
logger = logging.getLogger("unstructured.eval")
|
||||
|
||||
|
||||
def _prepare_output_cct(docpath: str, output_type: str) -> str:
|
||||
"""
|
||||
Convert given input document (path) into cct-ready. The function only support conversion
|
||||
from `json` or `txt` file.
|
||||
"""
|
||||
try:
|
||||
if output_type == "json":
|
||||
output_cct = elements_to_text(elements_from_json(docpath))
|
||||
elif output_type == "txt":
|
||||
output_cct = _read_text_file(docpath)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"File type not supported. Expects one of `json` or `txt`, \
|
||||
but received {output_type} instead."
|
||||
)
|
||||
except ValueError as e:
|
||||
logger.error(f"Could not read the file {docpath}")
|
||||
raise e
|
||||
return output_cct
|
||||
|
||||
|
||||
def _listdir_recursive(dir: str) -> List[str]:
|
||||
"""
|
||||
Recursively lists all files in the given directory and its subdirectories.
|
||||
Returns a list of all files found, with each file's path relative to the
|
||||
initial directory.
|
||||
"""
|
||||
listdir = []
|
||||
for dirpath, _, filenames in os.walk(dir):
|
||||
for filename in filenames:
|
||||
# Remove the starting directory from the path to show the relative path
|
||||
relative_path = os.path.relpath(dirpath, dir)
|
||||
if relative_path == ".":
|
||||
listdir.append(filename)
|
||||
else:
|
||||
listdir.append(os.path.join(relative_path, filename))
|
||||
return listdir
|
||||
|
||||
|
||||
def _rename_aggregated_columns(df):
|
||||
"""
|
||||
Renames aggregated columns in a DataFrame based on a predefined mapping.
|
||||
|
||||
Parameters:
|
||||
df (pandas.DataFrame): The DataFrame with aggregated columns to rename.
|
||||
|
||||
Returns:
|
||||
pandas.DataFrame: A new DataFrame with renamed aggregated columns.
|
||||
"""
|
||||
rename_map = {"_mean": "mean", "_stdev": "stdev", "_pstdev": "pstdev", "_count": "count"}
|
||||
return df.rename(columns=rename_map)
|
||||
|
||||
|
||||
def _format_grouping_output(*df):
|
||||
"""
|
||||
Concatenates multiple pandas DataFrame objects along the columns (side-by-side)
|
||||
and resets the index.
|
||||
"""
|
||||
return pd.concat(df, axis=1).reset_index()
|
||||
|
||||
|
||||
def _display(df):
|
||||
"""
|
||||
Displays the evaluation metrics in a formatted text table.
|
||||
"""
|
||||
if len(df) == 0:
|
||||
return
|
||||
headers = df.columns.tolist()
|
||||
col_widths = [
|
||||
max(len(header), max(len(str(item)) for item in df[header])) for header in headers
|
||||
]
|
||||
click.echo(" ".join(header.ljust(col_widths[i]) for i, header in enumerate(headers)))
|
||||
click.echo("-" * sum(col_widths) + "-" * (len(headers) - 1))
|
||||
for _, row in df.iterrows():
|
||||
formatted_row = []
|
||||
for item in row:
|
||||
if isinstance(item, float):
|
||||
formatted_row.append(f"{item:.3f}")
|
||||
else:
|
||||
formatted_row.append(str(item))
|
||||
click.echo(
|
||||
" ".join(formatted_row[i].ljust(col_widths[i]) for i in range(len(formatted_row))),
|
||||
)
|
||||
|
||||
|
||||
def _write_to_file(
|
||||
directory: str, filename: str, df: pd.DataFrame, mode: str = "w", overwrite: bool = True
|
||||
):
|
||||
"""
|
||||
Save the metrics report to tsv file. The function allows an option 1) to choose `mode`
|
||||
as `w` (write) or `a` (append) and 2) to `overwrite` the file if filename existed or not.
|
||||
"""
|
||||
if mode not in ["w", "a"]:
|
||||
raise ValueError("Mode not supported. Mode must be one of [w, a].")
|
||||
if directory:
|
||||
Path(directory).mkdir(exist_ok=True)
|
||||
if "count" in df.columns:
|
||||
df["count"] = df["count"].astype(int)
|
||||
if "filename" in df.columns and "connector" in df.columns:
|
||||
df.sort_values(by=["connector", "filename"], inplace=True)
|
||||
if not overwrite:
|
||||
filename = _get_non_duplicated_filename(directory, filename)
|
||||
df.to_csv(
|
||||
os.path.join(directory, filename), sep="\t", mode=mode, index=False, header=(mode == "w")
|
||||
)
|
||||
|
||||
|
||||
def _sorting_key(filename):
|
||||
"""
|
||||
A function that defines the sorting method for duplicated file names. For example,
|
||||
with filename.ext filename (1).ext filename (2).ext filename (10).ext - this function
|
||||
extracts the integer in the bracket and sort those numbers ascendingly.
|
||||
"""
|
||||
# Regular expression to find the number in the filename
|
||||
numbers = re.findall(r"(\d+)", filename)
|
||||
if numbers:
|
||||
# If there's a number, return it as an integer for sorting
|
||||
return int(numbers[-1])
|
||||
else:
|
||||
# If no number, return 0 so these files come first
|
||||
return 0
|
||||
|
||||
|
||||
def _uniquity_file(file_list, target_filename) -> str:
|
||||
"""
|
||||
Checks the duplicity of the file name from the list and run the numerical check
|
||||
of the minimum number needed as extension to not overwrite the exising file.
|
||||
Returns a string of file name in the format of `filename (<min number>).ext`.
|
||||
"""
|
||||
original_filename, extension = target_filename.rsplit(".", 1)
|
||||
pattern = rf"^{re.escape(original_filename)}(?: \((\d+)\))?\.{re.escape(extension)}$"
|
||||
duplicated_files = sorted([f for f in file_list if re.match(pattern, f)], key=_sorting_key)
|
||||
|
||||
numbers = []
|
||||
for file in duplicated_files:
|
||||
match = re.search(r"\((\d+)\)", file)
|
||||
if match:
|
||||
numbers.append(int(match.group(1)))
|
||||
|
||||
numbers.sort()
|
||||
|
||||
counter = 1
|
||||
for number in numbers:
|
||||
if number == counter:
|
||||
counter += 1
|
||||
else:
|
||||
break
|
||||
|
||||
return original_filename + " (" + str(counter) + ")." + extension
|
||||
|
||||
|
||||
def _get_non_duplicated_filename(dir, filename) -> str:
|
||||
"""
|
||||
Helper function to calls the `_uniquity_file` function. Takes in directory and file name
|
||||
to check on.
|
||||
"""
|
||||
filename = _uniquity_file(os.listdir(dir), filename)
|
||||
return filename
|
||||
|
||||
|
||||
def _mean(scores: Union[pd.Series, List[float]], rounding: Optional[int] = 3) -> Union[float, None]:
|
||||
"""
|
||||
Find mean from the list. Returns None if no element in the list.
|
||||
|
||||
Args:
|
||||
rounding (int): optional argument that allows user to define decimal points. Default at 3.
|
||||
"""
|
||||
if len(scores) == 0:
|
||||
return None
|
||||
mean = statistics.mean(scores)
|
||||
if not rounding:
|
||||
return mean
|
||||
return round(mean, rounding)
|
||||
|
||||
|
||||
def _stdev(scores: List[Optional[float]], rounding: Optional[int] = 3) -> Union[float, None]:
|
||||
"""
|
||||
Find standard deviation from the list.
|
||||
Returns None if only 0 or 1 element in the list.
|
||||
|
||||
Args:
|
||||
rounding (int): optional argument that allows user to define decimal points. Default at 3.
|
||||
"""
|
||||
# Filter out None values
|
||||
scores = [score for score in scores if score is not None]
|
||||
# Proceed only if there are more than one value
|
||||
if len(scores) <= 1:
|
||||
return None
|
||||
if not rounding:
|
||||
return statistics.stdev(scores)
|
||||
return round(statistics.stdev(scores), rounding)
|
||||
|
||||
|
||||
def _pstdev(scores: List[Optional[float]], rounding: Optional[int] = 3) -> Union[float, None]:
|
||||
"""
|
||||
Find population standard deviation from the list.
|
||||
Returns None if only 0 or 1 element in the list.
|
||||
|
||||
Args:
|
||||
rounding (int): optional argument that allows user to define decimal points. Default at 3.
|
||||
"""
|
||||
scores = [score for score in scores if score is not None]
|
||||
if len(scores) <= 1:
|
||||
return None
|
||||
if not rounding:
|
||||
return statistics.pstdev(scores)
|
||||
return round(statistics.pstdev(scores), rounding)
|
||||
|
||||
|
||||
def _count(scores: List[Optional[float]]) -> float:
|
||||
"""
|
||||
Returns the row count of the list.
|
||||
"""
|
||||
return len(scores)
|
||||
|
||||
|
||||
def _read_text_file(path):
|
||||
"""
|
||||
Reads the contents of a text file and returns it as a string.
|
||||
"""
|
||||
# Check if the file exists
|
||||
if not os.path.exists(path):
|
||||
raise FileNotFoundError(f"The file at {path} does not exist.")
|
||||
|
||||
try:
|
||||
with open(path, errors="ignore") as f:
|
||||
text = f.read()
|
||||
return text
|
||||
except OSError as e:
|
||||
# Handle other I/O related errors
|
||||
raise IOError(f"An error occurred when reading the file at {path}: {e}")
|
||||
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