indico_toolkit 6.14.0__py3-none-any.whl

indico_toolkit/results/predictions/__init__.py

@@ -0,0 +1,88 @@
+from typing import TYPE_CHECKING
+
+from ..model import ModelGroupType
+from .box import NULL_BOX, Box
+from .citation import NULL_CITATION, Citation
+from .classification import Classification
+from .documentextraction import DocumentExtraction
+from .extraction import Extraction
+from .formextraction import FormExtraction, FormExtractionType
+from .group import Group
+from .prediction import Prediction
+from .span import NULL_SPAN, Span
+from .summarization import Summarization
+from .unbundling import Unbundling
+
+if TYPE_CHECKING:
+    from ..document import Document
+    from ..errors import ResultError
+    from ..model import ModelGroup
+    from ..review import Review
+
+__all__ = (
+    "Box",
+    "Citation",
+    "Classification",
+    "DocumentExtraction",
+    "Extraction",
+    "FormExtraction",
+    "FormExtractionType",
+    "Group",
+    "NULL_BOX",
+    "NULL_CITATION",
+    "NULL_SPAN",
+    "Prediction",
+    "Span",
+    "Summarization",
+    "Unbundling",
+)
+
+CLASSIFICATION = ModelGroupType.CLASSIFICATION
+DOCUMENT_EXTRACTION = ModelGroupType.DOCUMENT_EXTRACTION
+FORM_EXTRACTION = ModelGroupType.FORM_EXTRACTION
+GENAI_CLASSIFICATION = ModelGroupType.GENAI_CLASSIFICATION
+GENAI_EXTRACTION = ModelGroupType.GENAI_EXTRACTION
+GENAI_SUMMARIZATION = ModelGroupType.GENAI_SUMMARIZATION
+UNBUNDLING = ModelGroupType.UNBUNDLING
+
+
+def from_v1_dict(
+    document: "Document",
+    model: "ModelGroup",
+    review: "Review | None",
+    prediction: object,
+) -> "Prediction":
+    """
+    Create a `Prediction` subclass from a v1 prediction dictionary.
+    """
+    if model.type == CLASSIFICATION:
+        return Classification.from_v1_dict(document, model, review, prediction)
+    elif model.type == DOCUMENT_EXTRACTION:
+        return DocumentExtraction.from_v1_dict(document, model, review, prediction)
+    elif model.type == FORM_EXTRACTION:
+        return FormExtraction.from_v1_dict(document, model, review, prediction)
+    else:
+        raise ResultError(f"unsupported v1 model type `{model.type!r}`")
+
+
+def from_v3_dict(
+    document: "Document",
+    model: "ModelGroup",
+    review: "Review | None",
+    prediction: object,
+) -> "Prediction":
+    """
+    Create a `Prediction` subclass from a v3 prediction dictionary.
+    """
+    if model.type in (CLASSIFICATION, GENAI_CLASSIFICATION):
+        return Classification.from_v3_dict(document, model, review, prediction)
+    elif model.type in (DOCUMENT_EXTRACTION, GENAI_EXTRACTION):
+        return DocumentExtraction.from_v3_dict(document, model, review, prediction)
+    elif model.type == FORM_EXTRACTION:
+        return FormExtraction.from_v3_dict(document, model, review, prediction)
+    elif model.type == GENAI_SUMMARIZATION:
+        return Summarization.from_v3_dict(document, model, review, prediction)
+    elif model.type == UNBUNDLING:
+        return Unbundling.from_v3_dict(document, model, review, prediction)
+    else:
+        raise ResultError(f"unsupported v3 model type `{model.type!r}`")

indico_toolkit/results/predictions/box.py

@@ -0,0 +1,58 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..utilities import get
+
+if TYPE_CHECKING:
+    from typing import Final
+
+
+@dataclass(frozen=True)
+class Box:
+    page: int
+    top: int
+    left: int
+    right: int
+    bottom: int
+
+    def __lt__(self, other: "Box") -> bool:
+        """
+        Bounding boxes are sorted with vertical hysteresis. Those on the same line are
+        sorted left-to-right, even when later tokens are higher than earlier ones,
+        as long as they overlap vertically.
+
+        ┌──────────────────┐ ┌───────────────────┐
+        │        1         │ │         2         │
+        └──────────────────┘ │                   │
+                             └───────────────────┘
+                        ┌────────────────┐
+        ┌─────────────┐ │        4       │ ┌─────┐
+        │      3      │ └────────────────┘ │  5  │
+        └─────────────┘                    └─────┘
+        """
+        return (
+            self.page < other.page
+            or (self.page == other.page and self.bottom < other.top)
+            or (
+                self.page == other.page
+                and self.top < other.bottom
+                and self.left < other.left
+            )
+        )
+
+    @staticmethod
+    def from_dict(box: object) -> "Box":
+        return Box(
+            page=get(box, int, "page_num"),
+            top=get(box, int, "top"),
+            left=get(box, int, "left"),
+            right=get(box, int, "right"),
+            bottom=get(box, int, "bottom"),
+        )
+
+
+# It's more ergonomic to represent the lack of a bounding box with a special null box
+# object rather than using `None` or raising an error. This lets you e.g. sort by the
+# `box` attribute without having to constantly check for `None`, while still allowing
+# you do a "None check" with `extraction.box == NULL_BOX`.
+NULL_BOX: "Final" = Box(page=0, top=0, left=0, right=0, bottom=0)

indico_toolkit/results/predictions/citation.py

@@ -0,0 +1,47 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..utilities import get
+from .span import NULL_SPAN, Span
+
+if TYPE_CHECKING:
+    from typing import Any, Final
+
+
+@dataclass(order=True, frozen=True)
+class Citation:
+    start: int
+    end: int
+    span: Span
+
+    @property
+    def slice(self) -> slice:
+        return slice(self.start, self.end)
+
+    def __bool__(self) -> bool:
+        return self != NULL_CITATION
+
+    @staticmethod
+    def from_dict(span: object) -> "Citation":
+        return Citation(
+            start=get(span, int, "response", "start"),
+            end=get(span, int, "response", "end"),
+            span=Span.from_dict(get(span, dict, "document")),
+        )
+
+    def to_dict(self) -> "dict[str, Any]":
+        return {
+            "document": self.span.to_dict(),
+            "response": {
+                "start": self.start,
+                "end": self.end,
+            },
+        }
+
+
+# It's more ergonomic to represent the lack of citations with a special null citation
+# object rather than using `None` or raising an error. This lets you e.g. sort by the
+# `citation` attribute without having to constantly check for `None`, while still
+# allowing you do a "None check" with `summarization.citation == NULL_CITATION` or
+# `bool(summarization.citation)`.
+NULL_CITATION: "Final" = Citation(span=NULL_SPAN, start=0, end=0)

indico_toolkit/results/predictions/classification.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..review import Review
+from ..utilities import get, omit
+from .prediction import Prediction
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from ..document import Document
+    from ..model import ModelGroup
+
+
+@dataclass
+class Classification(Prediction):
+    @staticmethod
+    def _from_dict(
+        document: "Document",
+        model: "ModelGroup",
+        review: "Review | None",
+        prediction: object,
+    ) -> "Classification":
+        """
+        Create a `Classification` from a prediction dictionary.
+        """
+        return Classification(
+            document=document,
+            model=model,
+            review=review,
+            label=get(prediction, str, "label"),
+            confidences=get(prediction, dict, "confidence"),
+            extras=omit(prediction, "label", "confidence"),
+        )
+
+    from_v1_dict = _from_dict
+    from_v3_dict = _from_dict
+
+    def _to_dict(self) -> "dict[str, Any]":
+        """
+        Create a prediction dictionary for auto review changes.
+        """
+        return {
+            **self.extras,
+            "label": self.label,
+            "confidence": self.confidences,
+        }
+
+    to_v1_dict = _to_dict
+    to_v3_dict = _to_dict

indico_toolkit/results/predictions/documentextraction.py

@@ -0,0 +1,160 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..review import Review
+from ..utilities import get, has, omit
+from .extraction import Extraction
+from .group import Group
+from .span import NULL_SPAN, Span
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from ..document import Document
+    from ..model import ModelGroup
+
+
+@dataclass
+class DocumentExtraction(Extraction):
+    groups: "set[Group]"
+    spans: "list[Span]"
+
+    @property
+    def span(self) -> Span:
+        """
+        Return the first `Span` the document extraction covers else `NULL_SPAN`.
+
+        Post-review, document extractions have no spans.
+        """
+        return self.spans[0] if self.spans else NULL_SPAN
+
+    @span.setter
+    def span(self, span: Span) -> None:
+        """
+        Overwrite all `spans` with the one provided.
+
+        This is implemented under the assumption that if you're setting the single span,
+        you want it to be the only one. And if you're working in a context that's
+        multiple-span sensetive, you'll set `extraction.spans` instead.
+        """
+        self.spans = [span]
+
+    @staticmethod
+    def from_v1_dict(
+        document: "Document",
+        model: "ModelGroup",
+        review: "Review | None",
+        prediction: object,
+    ) -> "DocumentExtraction":
+        """
+        Create a `DocumentExtraction` from a v1 prediction dictionary.
+        """
+        return DocumentExtraction(
+            document=document,
+            model=model,
+            review=review,
+            label=get(prediction, str, "label"),
+            confidences=get(prediction, dict, "confidence"),
+            text=get(prediction, str, "normalized", "formatted"),
+            accepted=(
+                has(prediction, bool, "accepted") and get(prediction, bool, "accepted")
+            ),
+            rejected=(
+                has(prediction, bool, "rejected") and get(prediction, bool, "rejected")
+            ),
+            groups=set(map(Group.from_dict, get(prediction, list, "groupings"))),
+            spans=[Span.from_dict(prediction)] if has(prediction, int, "start") else [],
+            extras=omit(
+                prediction,
+                "label",
+                "confidence",
+                "accepted",
+                "rejected",
+                "groupings",
+                "page_num",
+                "start",
+                "end",
+            ),
+        )
+
+    @staticmethod
+    def from_v3_dict(
+        document: "Document",
+        model: "ModelGroup",
+        review: "Review | None",
+        prediction: object,
+    ) -> "DocumentExtraction":
+        """
+        Create a `DocumentExtraction` from a v3 prediction dictionary.
+        """
+        return DocumentExtraction(
+            document=document,
+            model=model,
+            review=review,
+            label=get(prediction, str, "label"),
+            confidences=get(prediction, dict, "confidence"),
+            text=get(prediction, str, "normalized", "formatted"),
+            accepted=(
+                has(prediction, bool, "accepted") and get(prediction, bool, "accepted")
+            ),
+            rejected=(
+                has(prediction, bool, "rejected") and get(prediction, bool, "rejected")
+            ),
+            groups=set(map(Group.from_dict, get(prediction, list, "groupings"))),
+            spans=sorted(map(Span.from_dict, get(prediction, list, "spans"))),
+            extras=omit(
+                prediction,
+                "label",
+                "confidence",
+                "accepted",
+                "rejected",
+                "groupings",
+                "spans",
+            ),
+        )
+
+    def to_v1_dict(self) -> "dict[str, Any]":
+        """
+        Create a prediction dictionary for v1 auto review changes.
+        """
+        prediction = {
+            **self.extras,
+            "label": self.label,
+            "confidence": self.confidences,
+            "groupings": [group.to_dict() for group in self.groups],
+            "page_num": self.span.page,
+            "start": self.span.start,
+            "end": self.span.end,
+        }
+
+        prediction["normalized"]["formatted"] = self.text
+        prediction["text"] = self.text  # 6.10 sometimes reverts to raw text in review.
+
+        if self.accepted:
+            prediction["accepted"] = True
+        elif self.rejected:
+            prediction["rejected"] = True
+
+        return prediction
+
+    def to_v3_dict(self) -> "dict[str, Any]":
+        """
+        Create a prediction dictionary for v3 auto review changes.
+        """
+        prediction = {
+            **self.extras,
+            "label": self.label,
+            "confidence": self.confidences,
+            "groupings": [group.to_dict() for group in self.groups],
+            "spans": [span.to_dict() for span in self.spans],
+        }
+
+        prediction["normalized"]["formatted"] = self.text
+        prediction["text"] = self.text  # 6.10 sometimes reverts to raw text in review.
+
+        if self.accepted:
+            prediction["accepted"] = True
+        elif self.rejected:
+            prediction["rejected"] = True
+
+        return prediction

indico_toolkit/results/predictions/extraction.py

@@ -0,0 +1,35 @@
+from dataclasses import dataclass
+
+from .prediction import Prediction
+
+
+@dataclass
+class Extraction(Prediction):
+    text: str
+    accepted: bool
+    rejected: bool
+
+    @property
+    def page(self) -> int:
+        """
+        Convenience property to get an extraction's page without knowing its subclass.
+        Allows you to do `predictions.extractions.groupby(attrgetter("page"))` et al.
+        """
+        if hasattr(self, "box"):
+            return self.box.page  # type: ignore
+        else:
+            return self.span.page  # type: ignore
+
+    def accept(self) -> None:
+        self.accepted = True
+        self.rejected = False
+
+    def unaccept(self) -> None:
+        self.accepted = False
+
+    def reject(self) -> None:
+        self.accepted = False
+        self.rejected = True
+
+    def unreject(self) -> None:
+        self.rejected = False

indico_toolkit/results/predictions/formextraction.py

@@ -0,0 +1,121 @@
+from dataclasses import dataclass
+from enum import Enum
+from typing import TYPE_CHECKING
+
+from ..review import Review
+from ..utilities import get, has, omit
+from .box import Box
+from .extraction import Extraction
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from ..document import Document
+    from ..model import ModelGroup
+
+
+class FormExtractionType(Enum):
+    CHECKBOX = "checkbox"
+    SIGNATURE = "signature"
+    TEXT = "text"
+
+
+@dataclass
+class FormExtraction(Extraction):
+    type: FormExtractionType
+    box: Box
+    checked: bool
+    signed: bool
+
+    @staticmethod
+    def _from_dict(
+        document: "Document",
+        model: "ModelGroup",
+        review: "Review | None",
+        prediction: object,
+    ) -> "FormExtraction":
+        """
+        Create a `FormExtraction` from a prediction dictionary.
+        """
+        return FormExtraction(
+            document=document,
+            model=model,
+            review=review,
+            label=get(prediction, str, "label"),
+            confidences=get(prediction, dict, "confidence"),
+            text=get(prediction, str, "normalized", "formatted"),
+            accepted=(
+                has(prediction, bool, "accepted") and get(prediction, bool, "accepted")
+            ),
+            rejected=(
+                has(prediction, bool, "rejected") and get(prediction, bool, "rejected")
+            ),
+            type=FormExtractionType(get(prediction, str, "type")),
+            box=Box.from_dict(prediction),
+            checked=(
+                has(prediction, bool, "normalized", "structured", "checked")
+                and get(prediction, bool, "normalized", "structured", "checked")
+            ),
+            signed=(
+                has(prediction, bool, "normalized", "structured", "signed")
+                and get(prediction, bool, "normalized", "structured", "signed")
+            ),
+            extras=omit(
+                prediction,
+                "label",
+                "confidence",
+                "accepted",
+                "rejected",
+                "type",
+                "page_num",
+                "top",
+                "left",
+                "right",
+                "bottom",
+                "checked",
+                "signed",
+            ),
+        )
+
+    from_v1_dict = _from_dict
+    from_v3_dict = _from_dict
+
+    def _to_dict(self) -> "dict[str, Any]":
+        """
+        Create a prediction dictionary for auto review changes.
+        """
+        prediction = {
+            **self.extras,
+            "label": self.label,
+            "confidence": self.confidences,
+            "type": self.type.value,
+            "page_num": self.box.page,
+            "top": self.box.top,
+            "left": self.box.left,
+            "right": self.box.right,
+            "bottom": self.box.bottom,
+        }
+
+        if self.type == FormExtractionType.CHECKBOX:
+            prediction["normalized"]["structured"]["checked"] = self.checked
+            prediction["normalized"]["formatted"] = (
+                "Checked" if self.checked else "Unchecked"
+            )
+        elif self.type == FormExtractionType.SIGNATURE:
+            prediction["normalized"]["structured"]["signed"] = self.signed
+            prediction["normalized"]["formatted"] = (
+                "Signed" if self.signed else "Unsigned"
+            )
+        elif self.type == FormExtractionType.TEXT:
+            prediction["normalized"]["formatted"] = self.text
+            prediction["text"] = self.text  # 6.10 sometimes reverts to text in review.
+
+        if self.accepted:
+            prediction["accepted"] = True
+        elif self.rejected:
+            prediction["rejected"] = True
+
+        return prediction
+
+    to_v1_dict = _to_dict
+    to_v3_dict = _to_dict

indico_toolkit/results/predictions/group.py

@@ -0,0 +1,29 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..utilities import get
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+@dataclass(frozen=True, order=True)
+class Group:
+    id: int
+    name: str
+    index: int
+
+    @staticmethod
+    def from_dict(group: object) -> "Group":
+        return Group(
+            id=int(get(group, str, "group_id").split(":")[0]),
+            name=get(group, str, "group_name"),
+            index=get(group, int, "group_index"),
+        )
+
+    def to_dict(self) -> "dict[str, Any]":
+        return {
+            "group_id": f"{self.id}:{self.name}",
+            "group_name": self.name,
+            "group_index": self.index,
+        }

indico_toolkit/results/predictions/prediction.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..review import Review
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from ..document import Document
+    from ..model import ModelGroup
+
+
+@dataclass
+class Prediction:
+    document: "Document"
+    model: "ModelGroup"
+    review: "Review | None"
+
+    label: str
+    confidences: "dict[str, float]"
+    extras: "dict[str, Any]"
+
+    @property
+    def confidence(self) -> float:
+        return self.confidences[self.label]
+
+    @confidence.setter
+    def confidence(self, value: float) -> None:
+        self.confidences[self.label] = value
+
+    def to_v1_dict(self) -> "dict[str, Any]":
+        """
+        Create a prediction dictionary for v1 auto review changes.
+        """
+        raise NotImplementedError()
+
+    def to_v3_dict(self) -> "dict[str, Any]":
+        """
+        Create a prediction dictionary for v3 auto review changes.
+        """
+        raise NotImplementedError()

indico_toolkit/results/predictions/span.py

@@ -0,0 +1,43 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..utilities import get
+
+if TYPE_CHECKING:
+    from typing import Any, Final
+
+
+@dataclass(order=True, frozen=True)
+class Span:
+    page: int
+    start: int
+    end: int
+
+    @property
+    def slice(self) -> slice:
+        return slice(self.start, self.end)
+
+    def __bool__(self) -> bool:
+        return self != NULL_SPAN
+
+    @staticmethod
+    def from_dict(span: object) -> "Span":
+        return Span(
+            page=get(span, int, "page_num"),
+            start=get(span, int, "start"),
+            end=get(span, int, "end"),
+        )
+
+    def to_dict(self) -> "dict[str, Any]":
+        return {
+            "page_num": self.page,
+            "start": self.start,
+            "end": self.end,
+        }
+
+
+# It's more ergonomic to represent the lack of spans with a special null span object
+# rather than using `None` or raising an error. This lets you e.g. sort by the `span`
+# attribute without having to constantly check for `None`, while still allowing you do
+# a "None check" with `extraction.span == NULL_SPAN` or `bool(extraction.span)`.
+NULL_SPAN: "Final" = Span(page=0, start=0, end=0)