indico_toolkit 6.14.2__tar.gz → 7.2.0__tar.gz

{indico_toolkit-6.14.2 → indico_toolkit-7.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: indico-toolkit
-Version: 6.14.2
+Version: 7.2.0
 Summary: Classes, functions, and abstractions for Indico IPA
 License: MIT
 Author: Indico Solutions Engineering

{indico_toolkit-6.14.2 → indico_toolkit-7.2.0}/indico_toolkit/__init__.py

@@ -21,4 +21,4 @@ __all__ = (
     "ToolkitStaggeredLoopError",
     "ToolkitStatusError",
 )
-__version__ = "6.14.2"
+__version__ = "7.2.0"

indico_toolkit-7.2.0/indico_toolkit/etloutput/__init__.py

@@ -0,0 +1,157 @@
+from typing import TYPE_CHECKING, TypeAlias, TypeVar
+
+from ..results import NULL_BOX, NULL_SPAN, Box, Span
+from ..results.utils import get, has, json_loaded, str_decoded
+from .cell import Cell, CellType
+from .errors import EtlOutputError, TableCellNotFoundError, TokenNotFoundError
+from .etloutput import EtlOutput
+from .range import Range
+from .table import Table
+from .token import Token
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+
+__all__ = (
+    "Box",
+    "Cell",
+    "CellType",
+    "EtlOutput",
+    "EtlOutputError",
+    "load",
+    "load_async",
+    "NULL_BOX",
+    "NULL_SPAN",
+    "Range",
+    "Span",
+    "Table",
+    "TableCellNotFoundError",
+    "Token",
+    "TokenNotFoundError",
+)
+
+Loadable: TypeAlias = "dict[str, object] | list[object] | str | bytes"
+Readable = TypeVar("Readable")
+URI: TypeAlias = str
+
+
+def load(
+    etl_output: "Loadable | Readable",
+    *,
+    reader: "Callable[[Readable | URI], Loadable]",
+    text: bool = True,
+    tokens: bool = True,
+    tables: bool = True,
+) -> EtlOutput:
+    """
+    Load `etl_output` as an `EtlOutput` dataclass.
+
+    `etl_output` can be a dict, JSON string/bytes, or something that can be read by
+    `reader` to produce a loadable type.
+
+    `reader` must be able to load Indico storage URIs as lists or JSON strings/bytes.
+
+    Use `text`, `tokens`, and `tables` to specify what not to load.
+
+    ```
+    result = results.load(submission.result_file, reader=read_uri)
+    etl_outputs = {
+        document: etloutput.load(document.etl_output_uri, reader=read_uri)
+        for document in result.documents
+        if not document.failed
+    }
+    ```
+    """
+    if not isinstance(etl_output, dict):
+        if (isinstance(etl_output, str) and etl_output.startswith("{")) or (
+            isinstance(etl_output, bytes) and etl_output.startswith(b"{")
+        ):
+            etl_output = json_loaded(etl_output)
+        else:
+            etl_output = json_loaded(reader(etl_output))  # type: ignore[arg-type]
+
+    pages = get(etl_output, list, "pages")
+
+    if text and has(pages, str, 0, "text"):
+        text_pages = map(
+            lambda page: str_decoded(reader(get(page, str, "text"))), pages  # type: ignore[arg-type]
+        )
+    else:
+        text_pages = ()  # type: ignore[assignment]
+
+    if tokens and has(pages, str, 0, "tokens"):
+        token_dict_pages = map(
+            lambda page: json_loaded(reader(get(page, str, "tokens"))), pages
+        )
+    else:
+        token_dict_pages = ()  # type: ignore[assignment]
+
+    if tables and has(pages, str, 0, "tables"):
+        table_dict_pages = map(
+            lambda page: json_loaded(reader(get(page, str, "tables"))), pages
+        )
+    else:
+        table_dict_pages = ()  # type: ignore[assignment]
+
+    return EtlOutput.from_pages(text_pages, token_dict_pages, table_dict_pages)
+
+
+async def load_async(
+    etl_output: "Loadable | Readable",
+    *,
+    reader: "Callable[[Readable | URI], Awaitable[Loadable]]",
+    text: bool = True,
+    tokens: bool = True,
+    tables: bool = True,
+) -> EtlOutput:
+    """
+    Load `etl_output` as an `EtlOutput` dataclass.
+
+    `etl_output` can be a dict, JSON string/bytes, or something that can be read by
+    `reader` to produce a loadable type.
+
+    `reader` must be able to load Indico storage URIs as lists or JSON strings/bytes.
+
+    Use `text`, `tokens`, and `tables` to specify what not to load.
+
+    ```
+    result = await results.load_async(submission.result_file, reader=read_uri)
+    etl_outputs = {
+        document: await etloutput.load_async(document.etl_output_uri, reader=read_uri)
+        for document in result.documents
+        if not document.failed
+    }
+    ```
+    """
+    if not isinstance(etl_output, dict):
+        if (isinstance(etl_output, str) and etl_output.startswith("{")) or (
+            isinstance(etl_output, bytes) and etl_output.startswith(b"{")
+        ):
+            etl_output = json_loaded(etl_output)
+        else:
+            etl_output = json_loaded(await reader(etl_output))  # type: ignore[arg-type]
+
+    pages = get(etl_output, list, "pages")
+
+    if text and has(pages, str, 0, "text"):
+        text_pages = [
+            str_decoded(await reader(get(page, str, "text"))) for page in pages  # type: ignore[arg-type]
+        ]
+    else:
+        text_pages = ()  # type: ignore[assignment]
+
+    if tokens and has(pages, str, 0, "tokens"):
+        token_dict_pages = [
+            json_loaded(await reader(get(page, str, "tokens"))) for page in pages
+        ]
+    else:
+        token_dict_pages = ()  # type: ignore[assignment]
+
+    if tables and has(pages, str, 0, "tables"):
+        table_dict_pages = [
+            json_loaded(await reader(get(page, str, "tables"))) for page in pages
+        ]
+    else:
+        table_dict_pages = ()  # type: ignore[assignment]
+
+    return EtlOutput.from_pages(text_pages, token_dict_pages, table_dict_pages)

{indico_toolkit-6.14.2 → indico_toolkit-7.2.0}/indico_toolkit/etloutput/cell.py

@@ -31,7 +31,7 @@ class Cell:
     @staticmethod
     def from_dict(cell: object, page: int) -> "Cell":
         """
-        Create a `Cell` from a v1 or v3 cell dictionary.
+        Create a `Cell` from a cell dictionary.
         """
         get(cell, dict, "position")["page_num"] = page
 

indico_toolkit-7.2.0/indico_toolkit/etloutput/errors.py

@@ -0,0 +1,16 @@
+class EtlOutputError(Exception):
+    """
+    Raised when an error occurs accessing `EtlOutput` values.
+    """
+
+
+class TokenNotFoundError(EtlOutputError):
+    """
+    Raised when a `Token` can't be found for a `Span`.
+    """
+
+
+class TableCellNotFoundError(EtlOutputError):
+    """
+    Raised when a `Table` and `Cell` can't be found for a `Token`.
+    """

{indico_toolkit-6.14.2 → indico_toolkit-7.2.0}/indico_toolkit/etloutput/etloutput.py

@@ -28,30 +28,30 @@ class EtlOutput:
 
     @staticmethod
     def from_pages(
-        text_by_page: "Iterable[str]",
-        token_dicts_by_page: "Iterable[Iterable[object]]",
-        table_dicts_by_page: "Iterable[Iterable[object]]",
+        text_pages: "Iterable[str]",
+        token_dict_pages: "Iterable[Iterable[object]]",
+        table_dict_pages: "Iterable[Iterable[object]]",
     ) -> "EtlOutput":
         """
-        Create an `EtlOutput` from v1 or v3 page lists.
+        Create an `EtlOutput` from pages of text, tokens, and tables.
         """
-        text_by_page = tuple(text_by_page)
-        tokens_by_page = tuple(
+        text_pages = tuple(text_pages)
+        token_pages = tuple(
             tuple(sorted(map(Token.from_dict, token_dict_page), key=attrgetter("span")))
-            for token_dict_page in token_dicts_by_page
+            for token_dict_page in token_dict_pages
         )
-        tables_by_page = tuple(
+        table_pages = tuple(
             tuple(sorted(map(Table.from_dict, table_dict_page), key=attrgetter("box")))
-            for table_dict_page in table_dicts_by_page
+            for table_dict_page in table_dict_pages
         )
 
         return EtlOutput(
-            text="\n".join(text_by_page),
-            text_on_page=text_by_page,
-            tokens=tuple(itertools.chain.from_iterable(tokens_by_page)),
-            tokens_on_page=tokens_by_page,
-            tables=tuple(itertools.chain.from_iterable(tables_by_page)),
-            tables_on_page=tables_by_page,
+            text="\n".join(text_pages),
+            text_on_page=text_pages,
+            tokens=tuple(itertools.chain.from_iterable(token_pages)),
+            tokens_on_page=token_pages,
+            tables=tuple(itertools.chain.from_iterable(table_pages)),
+            tables_on_page=table_pages,
         )
 
     def token_for(self, span: Span) -> Token:
@@ -64,21 +64,21 @@ class EtlOutput:
             first = bisect_right(tokens, span.start, key=attrgetter("span.end"))
             last = bisect_left(tokens, span.end, lo=first, key=attrgetter("span.start"))
             tokens = tokens[first:last]
-
-            return Token(
-                text=self.text[span.slice],
-                span=span,
-                box=Box(
-                    page=min(token.box.page for token in tokens),
-                    top=min(token.box.top for token in tokens),
-                    left=min(token.box.left for token in tokens),
-                    right=max(token.box.right for token in tokens),
-                    bottom=max(token.box.bottom for token in tokens),
-                ),
-            )
         except (IndexError, ValueError) as error:
             raise TokenNotFoundError(f"no token contains {span!r}") from error
 
+        return Token(
+            text=self.text[span.slice],
+            box=Box(
+                page=span.page,
+                top=min(token.box.top for token in tokens),
+                left=min(token.box.left for token in tokens),
+                right=max(token.box.right for token in tokens),
+                bottom=max(token.box.bottom for token in tokens),
+            ),
+            span=span,
+        )
+
     def table_cell_for(self, token: Token) -> "tuple[Table, Cell]":
         """
         Return the `Table` and `Cell` that contain the midpoint of `token`.

{indico_toolkit-6.14.2 → indico_toolkit-7.2.0}/indico_toolkit/etloutput/range.py

@@ -15,7 +15,7 @@ class Range:
     @staticmethod
     def from_dict(cell: object) -> "Range":
         """
-        Create a `Range` from a v1 or v3 cell dictionary.
+        Create a `Range` from a cell dictionary.
         """
         rows = get(cell, list, "rows")
         columns = get(cell, list, "columns")

{indico_toolkit-6.14.2 → indico_toolkit-7.2.0}/indico_toolkit/etloutput/table.py

@@ -16,7 +16,7 @@ class Table:
     @staticmethod
     def from_dict(table: object) -> "Table":
         """
-        Create a `Table` from a v1 or v3 table dictionary.
+        Create a `Table` from a table dictionary.
         """
         page = get(table, int, "page_num")
         get(table, dict, "position")["page_num"] = page

{indico_toolkit-6.14.2 → indico_toolkit-7.2.0}/indico_toolkit/etloutput/token.py

@@ -13,7 +13,7 @@ class Token:
     @staticmethod
     def from_dict(token: object) -> "Token":
         """
-        Create a `Token` from a v1 or v3 token dictionary.
+        Create a `Token` from a token dictionary.
         """
         get(token, dict, "position")["page_num"] = get(token, int, "page_num")
         get(token, dict, "doc_offset")["page_num"] = get(token, int, "page_num")

{indico_toolkit-6.14.2 → indico_toolkit-7.2.0}/indico_toolkit/polling/autoreview.py

@@ -53,7 +53,7 @@ class AutoReviewPoller:
         load_etl_output: bool = True,
         load_text: bool = True,
         load_tokens: bool = True,
-        load_tables: bool = False,
+        load_tables: bool = True,
         retry_count: int = 4,
         retry_wait: float = 1,
         retry_backoff: float = 4,
@@ -96,8 +96,8 @@ class AutoReviewPoller:
                 *(self._reap_workers() for _ in range(self._worker_count)),
             )
 
-    async def _retrieve_storage_object(self, url: str) -> object:
-        return await self._client_call(RetrieveStorageObject(url))
+    async def _retrieve_storage_object(self, uri: str) -> "Any":
+        return await self._client_call(RetrieveStorageObject(uri))
 
     async def _spawn_workers(self) -> None:
         """

{indico_toolkit-6.14.2 → indico_toolkit-7.2.0}/indico_toolkit/results/__init__.py

@@ -1,9 +1,7 @@
-import json
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeAlias, TypeVar, overload
 
 from .document import Document
 from .errors import ResultError
-from .model import ModelGroup, ModelGroupType
 from .predictionlist import PredictionList
 from .predictions import (
     NULL_BOX,
@@ -23,7 +21,8 @@ from .predictions import (
 )
 from .result import Result
 from .review import Review, ReviewType
-from .utils import get
+from .task import Task, TaskType
+from .utils import json_loaded
 
 if TYPE_CHECKING:
     from collections.abc import Awaitable, Callable
@@ -40,8 +39,6 @@ __all__ = (
     "Group",
     "load",
     "load_async",
-    "ModelGroup",
-    "ModelGroupType",
     "NULL_BOX",
     "NULL_CITATION",
     "NULL_SPAN",
@@ -53,16 +50,24 @@ __all__ = (
     "ReviewType",
     "Span",
     "Summarization",
+    "Task",
+    "TaskType",
     "Unbundling",
 )
 
+Loadable: TypeAlias = "dict[str, object] | str | bytes"
+Readable = TypeVar("Readable")
 
+
+@overload
+def load(result: Loadable) -> Result: ...
+@overload
+def load(result: Readable, *, reader: "Callable[[Readable], Loadable]") -> Result: ...
 def load(result: object, *, reader: "Callable[..., object] | None" = None) -> Result:
     """
-    Load `result` as a Result dataclass.
+    Load `result` as a `Result` dataclass. `result` can be a dict or JSON string/bytes.
 
-    `result` can be a dict, a JSON string, or something that can be read with `reader`
-    to produce either.
+    If `reader` is provided, it should read `result` to produce a loadable type.
 
     ```
     for result_file in result_folder.glob("*.json"):
@@ -72,33 +77,28 @@ def load(result: object, *, reader: "Callable[..., object] | None" = None) -> Re
     if reader:
         result = reader(result)
 
-    return _load(result)
+    return Result.from_dict(json_loaded(result))
 
 
+@overload
+async def load_async(result: Loadable) -> Result: ...
+@overload
+async def load_async(
+    result: Readable, *, reader: "Callable[[Readable], Awaitable[Loadable]]"
+) -> Result: ...
 async def load_async(
     result: object, *, reader: "Callable[..., Awaitable[object]] | None" = None
 ) -> Result:
     """
-    Load `result` as a Result dataclass.
+    Load `result` as a `Result` dataclass. `result` can be a dict or JSON string/bytes.
+
+    If `reader` is provided, it should read `result` to produce a loadable type.
 
-    `result` can be a dict, a JSON string, or something that can be read with `reader`
-    to produce either.
+    ```
+    result = await results.load_async(submission.result_file, reader=read_uri)
+    ```
     """
     if reader:
         result = await reader(result)
 
-    return _load(result)
-
-
-def _load(result: object) -> Result:
-    if isinstance(result, str) and result.strip().startswith("{"):
-        result = json.loads(result)
-
-    file_version = get(result, int, "file_version")
-
-    if file_version == 1:
-        return Result.from_v1_dict(result)
-    elif file_version == 3:
-        return Result.from_v3_dict(result)
-    else:
-        raise ResultError(f"unsupported file version `{file_version}`")
+    return Result.from_dict(json_loaded(result))

indico_toolkit-7.2.0/indico_toolkit/results/document.py

@@ -0,0 +1,57 @@
+from dataclasses import dataclass
+
+from .utils import get
+
+
+@dataclass(frozen=True, order=True)
+class Document:
+    id: int
+    name: str
+    etl_output_uri: str
+    failed: bool
+    error: str
+    traceback: str
+
+    # Auto review changes must reproduce all model and component sections that were
+    # present in the original result file. This may not be possible from the
+    # predictions alone--if a model or component had an empty section because it didn't
+    # produce predictions or if all of the predictions for that section were dropped.
+    # As such, the model and component IDs seen when parsing a result file are tracked
+    # per-document so that the empty sections can be reproduced later.
+    _model_ids: "frozenset[str]"
+    _component_ids: "frozenset[str]"
+
+    @staticmethod
+    def from_dict(document: object) -> "Document":
+        """
+        Create a `Document` from a document dictionary.
+        """
+        model_results = get(document, dict, "model_results", "ORIGINAL")
+        component_results = get(document, dict, "component_results", "ORIGINAL")
+
+        return Document(
+            id=get(document, int, "submissionfile_id"),
+            name=get(document, str, "input_filename"),
+            etl_output_uri=get(document, str, "etl_output"),
+            failed=False,
+            error="",
+            traceback="",
+            _model_ids=frozenset(model_results.keys()),
+            _component_ids=frozenset(component_results.keys()),
+        )
+
+    @staticmethod
+    def from_errored_file_dict(errored_file: object) -> "Document":
+        """
+        Create a `Document` from an errored file dictionary.
+        """
+        return Document(
+            id=get(errored_file, int, "submissionfile_id"),
+            name=get(errored_file, str, "input_filename"),
+            etl_output_uri="",
+            failed=True,
+            error=get(errored_file, str, "error"),
+            traceback=get(errored_file, str, "traceback"),
+            _model_ids=frozenset(),
+            _component_ids=frozenset(),
+        )

indico_toolkit-7.2.0/indico_toolkit/results/normalization.py

@@ -0,0 +1,82 @@
+import re
+from typing import TYPE_CHECKING
+
+from .task import TaskType
+from .utils import get, has
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+def normalize_result_dict(result: "Any") -> None:
+    """
+    Fix inconsistencies observed in result file structure.
+    """
+    for errored_file in get(result, dict, "errored_files").values():
+        # Parse filenames for errored files.
+        if not has(errored_file, str, "input_filename"):
+            reason = get(errored_file, str, "reason")
+            match = re.search(r"file '([^']*)' with id", reason)
+            errored_file["input_filename"] = match.group(1) if match else ""
+
+        # Parse error from trackback for errored files.
+        if not has(errored_file, str, "traceback"):
+            traceback = get(errored_file, str, "error")
+            error = traceback.split("\n")[-1].strip()
+            errored_file["traceback"] = traceback
+            errored_file["error"] = error
+
+    # Convert `None` review notes to "".
+    for review in get(result, dict, "reviews").values():
+        if not has(review, str, "review_notes"):
+            review["review_notes"] = ""
+
+
+def normalize_prediction_dict(task_type: TaskType, prediction: "Any") -> None:
+    """
+    Fix inconsistencies observed in prediction structure.
+    """
+    # Predictions added in review lack a `confidence` section.
+    # (And should theoretically have 100% confidence.)
+    if not has(prediction, dict, "confidence"):
+        prediction["confidence"] = {get(prediction, str, "label"): 1.0}
+
+    # Extractions added in review may lack a `normalized` section.
+    if task_type in (
+        TaskType.DOCUMENT_EXTRACTION,
+        TaskType.FORM_EXTRACTION,
+        TaskType.GENAI_EXTRACTION,
+    ) and not has(prediction, dict, "normalized"):
+        prediction["normalized"] = {"formatted": get(prediction, str, "text")}
+
+    # Document Extractions added in review may lack a `spans` section.
+    # This value will match `NULL_SPAN`.
+    if task_type in (
+        TaskType.DOCUMENT_EXTRACTION,
+        TaskType.GENAI_EXTRACTION,
+    ) and not has(prediction, list, "spans"):
+        prediction["spans"] = []
+
+    # Form Extractions added in review may lack bounding box information.
+    # These values will match `NULL_BOX`.
+    if task_type == TaskType.FORM_EXTRACTION and not has(prediction, int, "top"):
+        prediction["page_num"] = 0
+        prediction["top"] = 0
+        prediction["left"] = 0
+        prediction["right"] = 0
+        prediction["bottom"] = 0
+
+    # Document Extractions that didn't go through a linked labels transformer
+    # lack a `groupings` section.
+    if task_type in (
+        TaskType.DOCUMENT_EXTRACTION,
+        TaskType.GENAI_EXTRACTION,
+    ) and not has(prediction, list, "groupings"):
+        prediction["groupings"] = []
+
+    # Summarizations added in review may lack a `citations` section.
+    # This value will match `NULL_CITATION`.
+    if task_type == TaskType.GENAI_SUMMARIZATION and not has(
+        prediction, list, "citations"
+    ):
+        prediction["citations"] = []