indico_toolkit 6.14.0__py3-none-any.whl

indico_toolkit/__init__.py

@@ -0,0 +1,24 @@
+"""Classes, functions, and abstractions for Indico IPA"""
+
+from .client import create_client
+from .errors import (
+    ToolkitAuthError,
+    ToolkitError,
+    ToolkitInputError,
+    ToolkitInstantiationError,
+    ToolkitPopulationError,
+    ToolkitStaggeredLoopError,
+    ToolkitStatusError,
+)
+
+__all__ = (
+    "create_client",
+    "ToolkitAuthError",
+    "ToolkitError",
+    "ToolkitInputError",
+    "ToolkitInstantiationError",
+    "ToolkitPopulationError",
+    "ToolkitStaggeredLoopError",
+    "ToolkitStatusError",
+)
+__version__ = "6.14.0"

indico_toolkit/association/__init__.py

@@ -0,0 +1,15 @@
+from .association import Association, sequences_exact, sequences_overlap
+from .extracted_tokens import ExtractedTokens
+from .line_items import LineItems
+from .positioning import Positioning
+from .splitting import split_prediction
+
+__all__ = (
+    "Association",
+    "ExtractedTokens",
+    "LineItems",
+    "Positioning",
+    "sequences_exact",
+    "sequences_overlap",
+    "split_prediction",
+)

indico_toolkit/association/association.py

@@ -0,0 +1,65 @@
+from abc import ABC, abstractmethod
+from collections import defaultdict
+from typing import Dict, List, Union
+
+from ..types import Extractions
+
+
+class Association(ABC):
+    """
+    Base class for matching tokens to extraction predictions.
+    """
+
+    def __init__(self, predictions: Union[List[dict], Extractions]):
+        self._predictions = self.validate_prediction_formatting(predictions)
+        self._mapped_positions = []
+        self._manually_added_preds = []
+        self._errored_predictions = []
+
+    @abstractmethod
+    def match_pred_to_token(self):
+        pass
+
+    @staticmethod
+    def sort_predictions_by_start_index(predictions: List[dict]) -> List[dict]:
+        return sorted(predictions, key=lambda x: x["start"])
+
+    @property
+    def mapped_positions_by_page(self) -> Dict[int, List[dict]]:
+        """
+        Return mapped positions by page on which they first appear
+        """
+        page_map = defaultdict(list)
+        for position in self._mapped_positions:
+            page_map[position["page_num"]].append(position)
+        return page_map
+
+    def validate_prediction_formatting(
+        self, predictions: Union[List[dict], Extractions]
+    ) -> List[dict]:
+        if isinstance(predictions, Extractions):
+            predictions = predictions.to_list()
+        return predictions
+
+    def _is_manually_added_pred(self, prediction: dict) -> bool:
+        return Extractions.is_manually_added_prediction(prediction)
+
+
+def sequences_overlap(x: dict, y: dict) -> bool:
+    """
+    Boolean return value indicates whether or not seqs overlap
+    """
+    return x["start"] < y["end"] and y["start"] < x["end"]
+
+
+def sequences_exact(x: dict, y: dict) -> bool:
+    """
+    Boolean return value indicates whether or not seqs are exact
+    """
+    return x["start"] == y["start"] and x["end"] == y["end"]
+
+
+def _check_if_token_match_found(pred: dict, no_match_indicator: bool):
+    if no_match_indicator:
+        pred["error"] = "No matching token found for extraction"
+        raise ValueError(f"Couldn't match a token to this prediction:\n{pred}")

indico_toolkit/association/extracted_tokens.py

@@ -0,0 +1,106 @@
+from typing import List, Union
+
+from ..types import Extractions
+from .association import Association, _check_if_token_match_found, sequences_overlap
+
+# TODO: add test that 'errored_predictions' actually works with ValueErrors
+
+
+class ExtractedTokens(Association):
+    """
+    Class to collect all extracted tokens, e.g. to enable highlighting of source
+    document
+    """
+
+    def __init__(self, predictions: Union[List[dict], Extractions]):
+        super().__init__(predictions)
+
+    def match_pred_to_token(self, pred: dict, ocr_tokens: List[dict], pred_index: int):
+        """
+        Append matching token positions to self.mapped_positions, if no matches for
+        pred, raise ValueError
+
+        Args:
+            pred (dict): Indico extraction model prediction
+            ocr_tokens (List[dict]): List of OCR tokens
+            pred_index (int): unique number for each prediction so that tokens can be
+                linked to it
+
+        Raises:
+            ValueError: No matching token was found
+
+        Returns:
+            [int]: index in ocr tokens where prediction matched
+        """
+        no_match = True
+        match_token_index = 0
+        pred["prediction_index"] = pred_index
+        for ind, token in enumerate(ocr_tokens):
+            if sequences_overlap(token["doc_offset"], pred):
+                self._add_pred_meta_to_token(token, pred, pred_index)
+                self._remove_unneeded_token_keys(token)
+                self._mapped_positions.append(token)
+                match_token_index = ind
+                no_match = False
+            elif token["doc_offset"]["start"] > pred["end"]:
+                break
+        _check_if_token_match_found(pred, no_match)
+        return match_token_index
+
+    def collect_tokens(self, ocr_tokens: List[dict], raise_for_no_match: bool = True):
+        """
+        Collect all extracted tokens and with pred text and label added to dictionaries
+        Args:
+        ocr_tokens (List[dict]): Tokens from 'ondocument' OCR config output
+        raise_for_no_match (bool): raise exception if a matching token isn't found for a
+            prediction
+        """
+        self._separate_manually_added_predictions()
+        self._predictions = self.sort_predictions_by_start_index(self._predictions)
+        match_index = 0
+        for pred_ind, pred in enumerate(self._predictions):
+            try:
+                match_index = self.match_pred_to_token(
+                    pred, ocr_tokens[match_index:], pred_ind
+                )
+            except ValueError as e:
+                if raise_for_no_match:
+                    raise ValueError(e)
+                else:
+                    print(f"Ignoring Error: {e}")
+                    self._errored_predictions.append(pred)
+
+    @property
+    def extracted_tokens(self) -> List[dict]:
+        return self._mapped_positions
+
+    @property
+    def predictions(self) -> Extractions:
+        return Extractions(
+            self._predictions + self._errored_predictions + self._manually_added_preds
+        )
+
+    @property
+    def unmapped_predictions(self) -> Extractions:
+        return Extractions(self._errored_predictions + self._manually_added_preds)
+
+    def _remove_unneeded_token_keys(
+        self, token: dict, keys: List[str] = ["style", "block_offset", "page_offset"]
+    ):
+        for key in keys:
+            token.pop(key, None)
+
+    def _separate_manually_added_predictions(self):
+        predictions = []
+        for pred in self._predictions:
+            if self._is_manually_added_pred(pred):
+                pred["error"] = "Can't match tokens for manually added prediction"
+                self._manually_added_preds.append(pred)
+            else:
+                predictions.append(pred)
+        self._predictions = predictions
+
+    def _add_pred_meta_to_token(self, token: dict, pred: dict, pred_index: int):
+        token["label"] = pred["label"]
+        token["text"] = pred["text"]
+        token["prediction_index"] = pred_index

indico_toolkit/association/line_items.py

@@ -0,0 +1,189 @@
+"""Associate row items"""
+
+from collections import defaultdict
+from copy import deepcopy
+from typing import Iterable, List, Union
+
+from ..types import Extractions
+from .association import Association, _check_if_token_match_found, sequences_overlap
+
+
+class LineItems(Association):
+    """
+    Class for associating line items given extraction predictions and ondocument OCR
+    tokens
+
+    Example Usage:
+
+    litems = LineItems(
+            predictions=[{"label": "line_date", "start": 12, "text": "1/2/2021".....}],
+            line_item_fields=["line_value", "line_date"],
+        )
+    litems.get_bounding_boxes(ocr_tokens=[{"postion"...,},])
+    litems.assign_row_number()
+
+    # Get your updated predictions
+    print(litems.updated_predictions)
+    """
+
+    def __init__(
+        self,
+        predictions: Union[List[dict], Extractions],
+        line_item_fields: Iterable[str],
+    ):
+        """
+        Args:
+        predictions (List[dict]): List of extraction predictions
+        line_item_fields (Iterable[str]): Fields/labels to include as line item values,
+            other values will not be assigned a row_number.
+        """
+        self.predictions = self.validate_prediction_formatting(predictions)
+        self.line_item_fields: Iterable[str] = line_item_fields
+        self._mapped_positions: List[dict] = []
+        self._unmapped_positions: List[dict] = []
+        self._errored_predictions: List[dict] = []
+
+    @property
+    def updated_predictions(self) -> Extractions:
+        return Extractions(
+            self._mapped_positions
+            + self._unmapped_positions
+            + self._errored_predictions
+        )
+
+    @staticmethod
+    def match_pred_to_token(pred: dict, ocr_tokens: List[dict]):
+        """
+        Match and add bounding box metadata to prediction.
+
+        Args:
+            pred (dict): Indico extraction model prediction
+            ocr_tokens (List[dict]): List of OCR tokens
+
+        Raises:
+            ValueError: No matching token was found
+
+        Returns:
+            [int]: index in ocr tokens where prediction matched
+        """
+        no_match = True
+        match_token_index = 0
+        for ind, token in enumerate(ocr_tokens):
+            if no_match and sequences_overlap(token["doc_offset"], pred):
+                _add_bounding_metadata_to_pred(pred, token)
+                no_match = False
+                match_token_index = ind
+            elif sequences_overlap(token["doc_offset"], pred):
+                _update_bounding_metadata_for_pred(pred, token)
+            elif token["doc_offset"]["start"] > pred["end"]:
+                break
+        _check_if_token_match_found(pred, no_match)
+        return match_token_index
+
+    def get_bounding_boxes(
+        self, ocr_tokens: List[dict], raise_for_no_match: bool = True
+    ):
+        """
+        Adds keys for bounding box top/bottom/left/right and page number to line item
+        predictions
+
+        Args:
+        ocr_tokens (List[dict]): Tokens from 'ondocument' OCR config output
+        raise_for_no_match (bool): raise exception if a matching token isn't found for a
+            prediction
+        """
+        predictions = deepcopy(self.predictions)
+        predictions = self._remove_unneeded_predictions(predictions)
+        predictions = self.sort_predictions_by_start_index(predictions)
+        match_index = 0
+        for pred in predictions:
+            try:
+                match_index = self.match_pred_to_token(pred, ocr_tokens[match_index:])
+                self._mapped_positions.append(pred)
+            except ValueError as e:
+                if raise_for_no_match:
+                    raise e
+                else:
+                    print(f"Ignoring Error: {e}")
+                    self._errored_predictions.append(pred)
+
+    def assign_row_number(self):
+        """
+        Adds a row_number:int key/val pair based on bounding box position and page
+
+        Updates:
+        self._mapped_positions (list of dicts): predictions with row_number added
+        """
+        self._mapped_positions = sorted(
+            self._mapped_positions,
+            key=lambda x: (x["page_num"], x["bbTop"], x["bbLeft"]),
+        )
+        starting_pred = self._get_first_valid_line_item_pred()
+        max_bot = starting_pred["bbBot"]
+        page_number = starting_pred["page_num"]
+        row_number = 1
+        for pred in self._mapped_positions:
+            # if the top of one box equals the bottom of another, we still want a new
+            # line
+            if pred["bbTop"] >= max_bot or pred["page_num"] != page_number:
+                row_number += 1
+                page_number = pred["page_num"]
+                max_bot = pred["bbBot"]
+            else:
+                max_bot = max(pred["bbBot"], max_bot)
+            pred["row_number"] = row_number
+
+    @property
+    def grouped_line_items(self) -> List[List[dict]]:
+        """
+        After row number has been assigned to predictions, returns line item predictions
+        as a list of lists where each list is a row.
+        """
+        rows = defaultdict(list)
+        for pred in self._mapped_positions:
+            rows[pred["row_number"]].append(pred)
+        return list(rows.values())
+
+    def _remove_unneeded_predictions(self, predictions: List[dict]) -> List[dict]:
+        """
+        Remove predictions that are not line item fields or don't have valid start/end
+        index data
+        """
+        valid_line_item_preds = []
+        for pred in predictions:
+            if not self.is_line_item_pred(pred):
+                self._unmapped_positions.append(pred)
+            elif self._is_manually_added_pred(pred):
+                pred["error"] = "Can't match tokens for manually added prediction"
+                self._errored_predictions.append(pred)
+            else:
+                valid_line_item_preds.append(pred)
+        return valid_line_item_preds
+
+    def is_line_item_pred(self, pred: dict):
+        if pred["label"] in self.line_item_fields:
+            return True
+        return False
+
+    def _get_first_valid_line_item_pred(self) -> dict:
+        if len(self._mapped_positions) == 0:
+            raise RuntimeError(
+                "Whoops! You have no line_item_fields predictions. "
+                "Did you run get_bounding_boxes?"
+            )
+        return self._mapped_positions[0]
+
+
+def _add_bounding_metadata_to_pred(pred: dict, token: dict):
+    pred["bbTop"] = token["position"]["bbTop"]
+    pred["bbBot"] = token["position"]["bbBot"]
+    pred["bbLeft"] = token["position"]["bbLeft"]
+    pred["bbRight"] = token["position"]["bbRight"]
+    pred["page_num"] = token["page_num"]
+
+
+def _update_bounding_metadata_for_pred(pred: dict, token: dict):
+    pred["bbTop"] = min(token["position"]["bbTop"], pred["bbTop"])
+    pred["bbBot"] = max(token["position"]["bbBot"], pred["bbBot"])
+    pred["bbLeft"] = min(token["position"]["bbLeft"], pred["bbLeft"])
+    pred["bbRight"] = max(token["position"]["bbRight"], pred["bbRight"])