indico_toolkit 6.14.0__py3-none-any.whl

indico_toolkit/results/document.py

@@ -0,0 +1,78 @@
+from dataclasses import dataclass
+
+from .utilities 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 sections that were present in the
+    # original result file. This may not be possible from the predictions alone--if a
+    # model had an empty section because it didn't produce predictions or if all of
+    # the predictions were removed to reject them. As such, the models seen when
+    # parsing result files are tracked per-document so that the empty sections can be
+    # reproduced later.
+    _model_sections: "frozenset[str]"
+
+    @staticmethod
+    def from_v1_dict(result: object) -> "Document":
+        """
+        Create a `Document` from a v1 result file dictionary.
+        """
+        document_results = get(result, dict, "results", "document", "results")
+        model_names = frozenset(document_results.keys())
+        etl_output_uri = get(result, str, "etl_output")
+
+        return Document(
+            # v1 result files don't include document IDs or filenames.
+            id=None,  # type: ignore[arg-type]
+            name=None,  # type: ignore[arg-type]
+            etl_output_uri=etl_output_uri,
+            failed=False,
+            error="",
+            traceback="",
+            _model_sections=model_names,
+        )
+
+    @staticmethod
+    def from_v3_dict(document: object) -> "Document":
+        """
+        Create a `Document` from a v3 document dictionary.
+        """
+        model_results = get(document, dict, "model_results", "ORIGINAL")
+        model_ids = frozenset(model_results.keys())
+        etl_output_uri = get(document, str, "etl_output")
+
+        return Document(
+            id=get(document, int, "submissionfile_id"),
+            name=get(document, str, "input_filename"),
+            etl_output_uri=etl_output_uri,
+            failed=False,
+            error="",
+            traceback="",
+            _model_sections=model_ids,
+        )
+
+    @staticmethod
+    def from_v3_errored_file(errored_file: object) -> "Document":
+        """
+        Create a `Document` from a v3 errored file dictionary.
+        """
+        traceback = get(errored_file, str, "error")
+        error = traceback.split("\n")[-1].strip()
+
+        return Document(
+            id=get(errored_file, int, "submissionfile_id"),
+            name=get(errored_file, str, "input_filename"),
+            etl_output_uri="",
+            failed=True,
+            error=error,
+            traceback=traceback,
+            _model_sections=frozenset(),
+        )

indico_toolkit/results/errors.py

@@ -0,0 +1,4 @@
+class ResultError(Exception):
+    """
+    Raised when an error occurs while loading or dumping a result file.
+    """

indico_toolkit/results/model.py

@@ -0,0 +1,60 @@
+from dataclasses import dataclass
+from enum import Enum
+
+from .utilities import get, has
+
+
+class ModelGroupType(Enum):
+    CLASSIFICATION = "classification"
+    DOCUMENT_EXTRACTION = "annotation"
+    FORM_EXTRACTION = "form_extraction"
+    GENAI_CLASSIFICATION = "genai_classification"
+    GENAI_EXTRACTION = "genai_annotation"
+    GENAI_SUMMARIZATION = "summarization"
+    UNBUNDLING = "classification_unbundling"
+
+
+@dataclass(frozen=True, order=True)
+class ModelGroup:
+    id: int
+    name: str
+    type: ModelGroupType
+
+    @staticmethod
+    def from_v1_section(section: "tuple[str, object]") -> "ModelGroup":
+        """
+        Create a `ModelGroup` from a v1 prediction section.
+        Use a heuristic on the first prediction of the model to determine its type.
+        """
+        name, predictions = section
+
+        if has(predictions, dict, "pre_review", 0):
+            prediction = get(predictions, dict, "pre_review", 0)
+
+            if has(prediction, str, "type"):
+                type = ModelGroupType.FORM_EXTRACTION
+            elif has(prediction, str, "text"):
+                type = ModelGroupType.DOCUMENT_EXTRACTION
+            else:
+                type = ModelGroupType.CLASSIFICATION
+        else:
+            # Likely an extraction model that produced no predictions.
+            type = ModelGroupType.DOCUMENT_EXTRACTION
+
+        return ModelGroup(
+            # v1 result files don't include model IDs.
+            id=None,  # type: ignore[arg-type]
+            name=name,
+            type=type,
+        )
+
+    @staticmethod
+    def from_v3_dict(model_group: object) -> "ModelGroup":
+        """
+        Create a `ModelGroup` from a v3 model group dictionary.
+        """
+        return ModelGroup(
+            id=get(model_group, int, "id"),
+            name=get(model_group, str, "name"),
+            type=ModelGroupType(get(model_group, str, "task_type")),
+        )

indico_toolkit/results/normalization.py

@@ -0,0 +1,159 @@
+import re
+from typing import TYPE_CHECKING
+
+from .utilities import get, has
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+    from typing import Any
+
+
+def normalize_v1_result(result: "Any") -> None:
+    """
+    Fix inconsistencies observed in v1 result files.
+    """
+    submission_results = get(result, dict, "results", "document", "results")
+
+    for model_name, model_results in tuple(submission_results.items()):
+        # Incomplete and unreviewed submissions don't have review sections.
+        if not has(model_results, list, "post_reviews"):
+            submission_results[model_name] = model_results = {
+                "pre_review": model_results,
+                "post_reviews": [],
+            }
+
+        # Classifications aren't wrapped in lists like other prediction types.
+        if has(model_results, dict, "pre_review"):
+            model_results["pre_review"] = [model_results["pre_review"]]
+            model_results["post_reviews"] = [
+                [prediction] for prediction in model_results["post_reviews"]
+            ]
+
+        predictions: "Any" = (
+            prediction
+            for review in (
+                [get(model_results, list, "pre_review")]
+                + get(model_results, list, "post_reviews")
+            )
+            if review is not None
+            for prediction in review
+            if prediction is not None
+        )
+
+        for prediction in predictions:
+            # Predictions added in review lack a `confidence` section.
+            if "confidence" not in prediction:
+                prediction["confidence"] = {prediction["label"]: 0}
+
+            # Form Extractions added in review may lack bounding boxes.
+            # Set values that will equal `NULL_BOX`.
+            if "type" in prediction and "top" not in prediction:
+                prediction["page_num"] = 0
+                prediction["top"] = 0
+                prediction["left"] = 0
+                prediction["right"] = 0
+                prediction["bottom"] = 0
+
+            # Prior to 6.11, some Extractions lack a `normalized` section after review.
+            if "text" in prediction and "normalized" not in prediction:
+                prediction["normalized"] = {"formatted": prediction["text"]}
+
+            # Document Extractions that didn't go through a linked labels transformer
+            # lack a `groupings` section.
+            if (
+                "text" in prediction
+                and "type" not in prediction
+                and "groupings" not in prediction
+            ):
+                prediction["groupings"] = []
+
+    # Incomplete and unreviewed submissions don't include a `reviews_meta` section.
+    if not has(result, list, "reviews_meta"):
+        result["reviews_meta"] = []
+
+    # Incomplete and unreviewed submissions retrieved with `SubmissionResult()` have a
+    # single `{"review_id": None}` review.
+    if not has(result, int, "reviews_meta", 0, "review_id"):
+        result["reviews_meta"] = []
+
+    # Review notes are `None` unless the reviewer enters a reason for rejection.
+    for review_dict in get(result, list, "reviews_meta"):
+        if not has(review_dict, str, "review_notes"):
+            review_dict["review_notes"] = ""
+
+
+def normalize_v3_result(result: "Any") -> None:
+    """
+    Fix inconsistencies observed in v3 result files.
+    """
+    task_type_by_model_group_id = {
+        model_group_id: model_group["task_type"]
+        for model_group_id, model_group in result["modelgroup_metadata"].items()
+    }
+    predictions_with_task_type: "Iterator[tuple[Any, str]]" = (
+        (prediction, task_type_by_model_group_id[model_group_id])
+        for submission_result in get(result, list, "submission_results")
+        for review_result in get(submission_result, dict, "model_results").values()
+        for model_group_id, model_results in review_result.items()
+        for prediction in model_results
+    )
+
+    for prediction, task_type in predictions_with_task_type:
+        # Predictions added in review may lack a `confidence` section.
+        if "confidence" not in prediction:
+            prediction["confidence"] = {prediction["label"]: 0}
+
+        # Document Extractions added in review may lack spans.
+        if (
+            task_type in ("annotation", "genai_annotation")
+            and "spans" not in prediction
+        ):
+            prediction["spans"] = []
+
+        # Form Extractions added in review may lack bounding boxes.
+        # Set values that will equal `NULL_BOX`.
+        if task_type == "form_extraction" and "top" not in prediction:
+            prediction["page_num"] = 0
+            prediction["top"] = 0
+            prediction["left"] = 0
+            prediction["right"] = 0
+            prediction["bottom"] = 0
+
+        # Prior to 6.11, some Extractions lack a `normalized` section after
+        # review.
+        if (
+            task_type in ("annotation", "form_extraction", "genai_annotation")
+            and "normalized" not in prediction
+        ):
+            prediction["normalized"] = {"formatted": prediction["text"]}
+
+        # Document Extractions that didn't go through a linked labels
+        # transformer lack a `groupings` section.
+        if (
+            task_type in ("annotation", "genai_annotation")
+            and "groupings" not in prediction
+        ):
+            prediction["groupings"] = []
+
+        # Summarizations may lack citations after review.
+        if task_type == "summarization" and "citations" not in prediction:
+            prediction["citations"] = []
+
+    # Prior to 6.8, v3 result files don't include a `reviews` section.
+    if not has(result, dict, "reviews"):
+        result["reviews"] = {}
+
+    # Review notes are `None` unless the reviewer enters a reason for rejection.
+    for review_dict in get(result, dict, "reviews").values():
+        if not has(review_dict, str, "review_notes"):
+            review_dict["review_notes"] = ""
+
+    # Prior to 7.0, v3 result files don't include an `errored_files` section.
+    if not has(result, dict, "errored_files"):
+        result["errored_files"] = {}
+
+    # Prior to 7.X, errored files may lack filenames.
+    for file in get(result, dict, "errored_files").values():
+        if not has(file, str, "input_filename") and has(file, str, "reason"):
+            match = re.search(r"file '([^']*)' with id", get(file, str, "reason"))
+            file["input_filename"] = match.group(1) if match else ""

indico_toolkit/results/predictionlist.py

@@ -0,0 +1,395 @@
+from collections import defaultdict
+from operator import attrgetter
+from typing import TYPE_CHECKING, List, TypeVar, overload
+
+from .model import ModelGroupType
+from .predictions import (
+    Classification,
+    DocumentExtraction,
+    Extraction,
+    FormExtraction,
+    Prediction,
+    Summarization,
+    Unbundling,
+)
+from .review import Review, ReviewType
+from .utilities import nfilter
+
+if TYPE_CHECKING:
+    from collections.abc import Callable, Collection, Container, Iterable
+    from typing import Any, Final, SupportsIndex
+
+    from typing_extensions import Self
+
+    from .document import Document
+    from .model import ModelGroup
+    from .result import Result
+
+PredictionType = TypeVar("PredictionType", bound=Prediction)
+OfType = TypeVar("OfType", bound=Prediction)
+KeyType = TypeVar("KeyType")
+
+# Non-None sentinel value to support `PredictionList.where(review=None)`.
+REVIEW_UNSPECIFIED: "Final" = Review(
+    id=None, reviewer_id=None, notes=None, rejected=None, type=None  # type: ignore[arg-type]
+)
+
+
+class PredictionList(List[PredictionType]):
+    @property
+    def classifications(self) -> "PredictionList[Classification]":
+        return self.oftype(Classification)
+
+    @property
+    def document_extractions(self) -> "PredictionList[DocumentExtraction]":
+        return self.oftype(DocumentExtraction)
+
+    @property
+    def extractions(self) -> "PredictionList[Extraction]":
+        return self.oftype(Extraction)
+
+    @property
+    def form_extractions(self) -> "PredictionList[FormExtraction]":
+        return self.oftype(FormExtraction)
+
+    @property
+    def summarizations(self) -> "PredictionList[Summarization]":
+        return self.oftype(Summarization)
+
+    @property
+    def unbundlings(self) -> "PredictionList[Unbundling]":
+        return self.oftype(Unbundling)
+
+    @overload
+    def __getitem__(self, index: "SupportsIndex", /) -> PredictionType: ...
+
+    @overload
+    def __getitem__(self, index: slice, /) -> "PredictionList[PredictionType]": ...
+
+    def __getitem__(
+        self, index: "SupportsIndex | slice"
+    ) -> "PredictionType | PredictionList[PredictionType]":
+        if isinstance(index, slice):
+            return type(self)(super().__getitem__(index))
+        else:
+            return super().__getitem__(index)
+
+    def apply(self, function: "Callable[[PredictionType], None]") -> "Self":
+        """
+        Apply `function` to all predictions.
+        """
+        for prediction in self:
+            function(prediction)
+
+        return self
+
+    def groupby(
+        self, key: "Callable[[PredictionType], KeyType]"
+    ) -> "dict[KeyType, Self]":
+        """
+        Group predictions into a dictionary using `key` to derive each prediction's key.
+        E.g. `key=attrgetter("label")` or `key=attrgetter("model")`.
+
+        If a derived key is an unhashable mutable collection (like set),
+        it's automatically converted to its hashable immutable variant (like frozenset).
+        This makes it easy to group by linked labels or unbundling pages.
+        """
+        grouped = defaultdict(type(self))  # type: ignore[var-annotated]
+
+        for prediction in self:
+            derived_key = key(prediction)
+
+            if isinstance(derived_key, list):
+                derived_key = tuple(derived_key)  # type: ignore[assignment]
+            elif isinstance(derived_key, set):
+                derived_key = frozenset(derived_key)  # type: ignore[assignment]
+
+            grouped[derived_key].append(prediction)
+
+        return grouped
+
+    def groupbyiter(
+        self, keys: "Callable[[PredictionType], Iterable[KeyType]]"
+    ) -> "dict[KeyType, Self]":
+        """
+        Group predictions into a dictionary using `key` to derive an iterable of keys.
+        E.g. `key=attrgetter("groups")` or `key=attrgetter("pages")`.
+
+        Each prediction is associated with every key in the iterable individually.
+        If the iterable is empty, the prediction is not included in any group.
+        """
+        grouped = defaultdict(type(self))  # type: ignore[var-annotated]
+
+        for prediction in self:
+            for derived_key in keys(prediction):
+                grouped[derived_key].append(prediction)
+
+        return grouped
+
+    def oftype(self, type: "type[OfType]") -> "PredictionList[OfType]":
+        """
+        Return a new prediction list containing predictions of type `type`.
+        """
+        return self.where(lambda prediction: isinstance(prediction, type))  # type: ignore[return-value]
+
+    def orderby(
+        self,
+        key: "Callable[[PredictionType], Any]",
+        *,
+        reverse: bool = False,
+    ) -> "Self":
+        """
+        Return a new prediction list with predictions sorted by `key`.
+        """
+        return type(self)(sorted(self, key=key, reverse=reverse))
+
+    def where(
+        self,
+        predicate: "Callable[[PredictionType], bool] | None" = None,
+        *,
+        document: "Document | None" = None,
+        document_in: "Container[Document] | None" = None,
+        model: "ModelGroup | ModelGroupType | str | None" = None,
+        model_in: "Container[ModelGroup | ModelGroupType | str] | None" = None,
+        review: "Review | ReviewType | None" = REVIEW_UNSPECIFIED,
+        review_in: "Container[Review | ReviewType | None]" = {REVIEW_UNSPECIFIED},
+        label: "str | None" = None,
+        label_in: "Container[str] | None" = None,
+        page: "int | None" = None,
+        page_in: "Collection[int] | None" = None,
+        min_confidence: "float | None" = None,
+        max_confidence: "float | None" = None,
+        accepted: "bool | None" = None,
+        rejected: "bool | None" = None,
+        checked: "bool | None" = None,
+        signed: "bool | None" = None,
+    ) -> "Self":
+        """
+        Return a new prediction list containing predictions that match
+        all of the specified filters.
+
+        predicate: predictions for which this function returns True,
+        document: predictions from this document,
+        document_in: predictions from these documents,
+        model: predictions from this model, task type, or name,
+        model_in: predictions from these models, task types, or names,
+        review: predictions from this review or review type,
+        review_in: predictions from these reviews or review types,
+        label: predictions with this label,
+        label_in: predictions with these labels,
+        page: extractions/unbundlings on this page,
+        page_in: extractions/unbundlings on these pages,
+        min_confidence: predictions with confidence >= this threshold,
+        max_confidence: predictions with confidence <= this threshold,
+        accepted: extractions that have or haven't been accepted,
+        rejected: extractions that have or haven't been rejected,
+        checked: form extractions that are or aren't checked,
+        signed: form extractions that are or aren't signed.
+        """
+        predicates = []
+
+        if predicate is not None:
+            predicates.append(predicate)
+
+        if document is not None:
+            predicates.append(lambda prediction: prediction.document == document)
+
+        if document_in is not None:
+            predicates.append(lambda prediction: prediction.document in document_in)
+
+        if model is not None:
+            predicates.append(
+                lambda prediction: (
+                    prediction.model == model
+                    or prediction.model.type == model
+                    or prediction.model.name == model
+                )
+            )
+
+        if model_in is not None:
+            predicates.append(
+                lambda prediction: (
+                    prediction.model in model_in
+                    or prediction.model.type in model_in
+                    or prediction.model.name in model_in
+                )
+            )
+
+        if review is not REVIEW_UNSPECIFIED:
+            predicates.append(
+                lambda prediction: (
+                    prediction.review == review
+                    or (
+                        prediction.review is not None
+                        and prediction.review.type == review
+                    )
+                )
+            )
+
+        if review_in != {REVIEW_UNSPECIFIED}:
+            predicates.append(
+                lambda prediction: (
+                    prediction.review in review_in
+                    or (
+                        prediction.review is not None
+                        and prediction.review.type in review_in
+                    )
+                )
+            )
+
+        if label is not None:
+            predicates.append(lambda prediction: prediction.label == label)
+
+        if label_in is not None:
+            predicates.append(lambda prediction: prediction.label in label_in)
+
+        if min_confidence is not None:
+            predicates.append(
+                lambda prediction: prediction.confidence >= min_confidence
+            )
+
+        if max_confidence is not None:
+            predicates.append(
+                lambda prediction: prediction.confidence <= max_confidence
+            )
+
+        if page is not None:
+            predicates.append(
+                lambda prediction: (
+                    (isinstance(prediction, Extraction) and prediction.page == page)
+                    or (isinstance(prediction, Unbundling) and page in prediction.pages)
+                )
+            )
+
+        if page_in is not None:
+            page_in = set(page_in)
+            predicates.append(
+                lambda prediction: (
+                    (isinstance(prediction, Extraction) and prediction.page in page_in)
+                    or (
+                        isinstance(prediction, Unbundling)
+                        and bool(page_in & set(prediction.pages))
+                    )
+                )
+            )
+
+        if accepted is not None:
+            predicates.append(
+                lambda prediction: isinstance(prediction, Extraction)
+                and prediction.accepted == accepted
+            )
+
+        if rejected is not None:
+            predicates.append(
+                lambda prediction: isinstance(prediction, Extraction)
+                and prediction.rejected == rejected
+            )
+
+        if checked is not None:
+            predicates.append(
+                lambda prediction: isinstance(prediction, FormExtraction)
+                and prediction.checked == checked
+            )
+
+        if signed is not None:
+            predicates.append(
+                lambda prediction: isinstance(prediction, FormExtraction)
+                and prediction.signed == signed
+            )
+
+        return type(self)(nfilter(predicates, self))
+
+    def accept(self) -> "Self":
+        """
+        Mark extractions as accepted for auto review.
+        """
+        self.oftype(Extraction).apply(Extraction.accept)
+        return self
+
+    def unaccept(self) -> "Self":
+        """
+        Mark extractions as not accepted for auto review.
+        """
+        self.oftype(Extraction).apply(Extraction.unaccept)
+        return self
+
+    def reject(self) -> "Self":
+        """
+        Mark extractions as rejected for auto review.
+        """
+        self.oftype(Extraction).apply(Extraction.reject)
+        return self
+
+    def unreject(self) -> "Self":
+        """
+        Mark extractions as not rejected for auto review.
+        """
+        self.oftype(Extraction).apply(Extraction.unreject)
+        return self
+
+    def to_changes(self, result: "Result") -> "Any":
+        """
+        Create a dict or list for the `changes` argument of `SubmitReview` based on the
+        predictions in this prediction list and the documents and version of `result`.
+        """
+        if result.version == 1:
+            return self.to_v1_changes(result.documents[0])
+        elif result.version == 3:
+            return self.to_v3_changes(result.documents)
+        else:
+            raise ValueError(f"unsupported file version `{result.version!r}`")
+
+    def to_v1_changes(self, document: "Document") -> "dict[str, Any]":
+        """
+        Create a v1 dict for the `changes` argument of `SubmitReview`.
+        """
+        changes: "dict[str, Any]" = {}
+
+        for model, predictions in self.groupby(attrgetter("model")).items():
+            if model.type == ModelGroupType.CLASSIFICATION:
+                changes[model.name] = predictions[0].to_v1_dict()
+            else:
+                changes[model.name] = [
+                    prediction.to_v1_dict() for prediction in predictions
+                ]
+
+        for model_name in document._model_sections:
+            if model_name not in changes:
+                changes[model_name] = []
+
+        return changes
+
+    def to_v3_changes(self, documents: "Iterable[Document]") -> "list[dict[str, Any]]":
+        """
+        Create a v3 list for the `changes` argument of `SubmitReview`.
+        """
+        changes: "list[dict[str, Any]]" = []
+
+        for document in documents:
+            if document.failed:
+                continue
+
+            model_results: "dict[str, Any]" = {}
+            changes.append(
+                {
+                    "submissionfile_id": document.id,
+                    "model_results": model_results,
+                    "component_results": {},
+                }
+            )
+            predictions_by_model = self.where(
+                document=document,
+            ).groupby(
+                attrgetter("model"),
+            )
+
+            for model, predictions in predictions_by_model.items():
+                model_results[str(model.id)] = [
+                    prediction.to_v3_dict() for prediction in predictions
+                ]
+
+            for model_id in document._model_sections:
+                if model_id not in model_results:
+                    model_results[model_id] = []
+
+        return changes