indico_toolkit 6.14.0__py3-none-any.whl

indico_toolkit/results/predictions/summarization.py

@@ -0,0 +1,103 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+from ..review import Review
+from ..utilities import get, has, omit
+from .citation import NULL_CITATION, Citation
+from .extraction import Extraction
+
+if TYPE_CHECKING:
+    from typing import Any
+
+    from ..document import Document
+    from ..model import ModelGroup
+    from .span import Span
+
+
+@dataclass
+class Summarization(Extraction):
+    citations: "list[Citation]"
+
+    @property
+    def citation(self) -> Citation:
+        """
+        Return the first `Citation` the summarization covers else `NULL_CITATION`.
+
+        Post-review, summarizations have no citations.
+        """
+        return self.citations[0] if self.citations else NULL_CITATION
+
+    @citation.setter
+    def citation(self, citation: Citation) -> None:
+        """
+        Overwrite all `citations` with the one provided.
+
+        This is implemented under the assumption that if you're setting the single
+        citation, you want it to be the only one. And if you're working in a context
+        that's multiple-citation sensetive, you'll set `extraction.citations` instead.
+        """
+        self.citations = [citation]
+
+    @property
+    def span(self) -> "Span":
+        return self.citation.span
+
+    @span.setter
+    def span(self, span: "Span") -> None:
+        self.citations = [Citation(self.citation.start, self.citation.end, span)]
+
+    @staticmethod
+    def from_v3_dict(
+        document: "Document",
+        model: "ModelGroup",
+        review: "Review | None",
+        prediction: object,
+    ) -> "Summarization":
+        """
+        Create a `Summarization` from a v3 prediction dictionary.
+        """
+        return Summarization(
+            document=document,
+            model=model,
+            review=review,
+            label=get(prediction, str, "label"),
+            confidences=get(prediction, dict, "confidence"),
+            text=get(prediction, str, "text"),
+            accepted=(
+                has(prediction, bool, "accepted") and get(prediction, bool, "accepted")
+            ),
+            rejected=(
+                has(prediction, bool, "rejected") and get(prediction, bool, "rejected")
+            ),
+            citations=sorted(
+                map(Citation.from_dict, get(prediction, list, "citations"))
+            ),
+            extras=omit(
+                prediction,
+                "label",
+                "confidence",
+                "text",
+                "accepted",
+                "rejected",
+                "citations",
+            ),
+        )
+
+    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,
+            "text": self.text,
+            "citations": [citation.to_dict() for citation in self.citations],
+        }
+
+        if self.accepted:
+            prediction["accepted"] = True
+        elif self.rejected:
+            prediction["rejected"] = True
+
+        return prediction

indico_toolkit/results/predictions/unbundling.py

@@ -0,0 +1,51 @@
+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 Unbundling(Prediction):
+    pages: "list[int]"
+
+    @staticmethod
+    def from_v3_dict(
+        document: "Document",
+        model: "ModelGroup",
+        review: "Review | None",
+        prediction: object,
+    ) -> "Unbundling":
+        """
+        Create an `Unbundling` from a v3 prediction dictionary.
+        """
+        return Unbundling(
+            document=document,
+            model=model,
+            review=review,
+            label=get(prediction, str, "label"),
+            confidences=get(prediction, dict, "confidence"),
+            pages=[
+                get(span, int, "page_num")
+                for span in get(prediction, list, "spans")  # fmt: skip
+            ],
+            extras=omit(prediction, "confidence", "label", "spans"),
+        )
+
+    def to_v3_dict(self) -> "dict[str, Any]":
+        """
+        Create a prediction dictionary for v3 auto review changes.
+        """
+        return {
+            **self.extras,
+            "label": self.label,
+            "confidence": self.confidences,
+            "spans": [{"page_num": page} for page in self.pages],
+        }

indico_toolkit/results/result.py

@@ -0,0 +1,152 @@
+from dataclasses import dataclass
+from functools import partial
+from itertools import chain
+from typing import TYPE_CHECKING
+
+from . import predictions as prediction
+from .document import Document
+from .model import ModelGroup
+from .normalization import normalize_v1_result, normalize_v3_result
+from .predictionlist import PredictionList
+from .predictions import Prediction
+from .review import Review, ReviewType
+from .utilities import get
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+@dataclass(frozen=True, order=True)
+class Result:
+    version: int
+    submission_id: int
+    documents: "tuple[Document, ...]"
+    models: "tuple[ModelGroup, ...]"
+    predictions: "PredictionList[Prediction]"
+    reviews: "tuple[Review, ...]"
+
+    @property
+    def rejected(self) -> bool:
+        return len(self.reviews) > 0 and self.reviews[-1].rejected
+
+    @property
+    def pre_review(self) -> "PredictionList[Prediction]":
+        return self.predictions.where(review=None)
+
+    @property
+    def auto_review(self) -> "PredictionList[Prediction]":
+        return self.predictions.where(review=ReviewType.AUTO)
+
+    @property
+    def manual_review(self) -> "PredictionList[Prediction]":
+        return self.predictions.where(review=ReviewType.MANUAL)
+
+    @property
+    def admin_review(self) -> "PredictionList[Prediction]":
+        return self.predictions.where(review=ReviewType.ADMIN)
+
+    @property
+    def final(self) -> "PredictionList[Prediction]":
+        return self.predictions.where(review=self.reviews[-1] if self.reviews else None)
+
+    @staticmethod
+    def from_v1_dict(result: object) -> "Result":
+        """
+        Create a `Result` from a v1 result file dictionary.
+        """
+        normalize_v1_result(result)
+
+        version = get(result, int, "file_version")
+        submission_id = get(result, int, "submission_id")
+        submission_results = get(result, dict, "results", "document", "results")
+        review_metadata = get(result, list, "reviews_meta")
+
+        document = Document.from_v1_dict(result)
+        models = sorted(map(ModelGroup.from_v1_section, submission_results.items()))
+        predictions: "PredictionList[Prediction]" = PredictionList()
+        # Reviews must be sorted after parsing predictions, as they match positionally
+        # with prediction lists in `post_reviews`.
+        reviews = list(map(Review.from_dict, review_metadata))
+
+        for model_name, model_predictions in submission_results.items():
+            model = next(filter(lambda model: model.name == model_name, models))
+            reviewed_model_predictions: "list[tuple[Review | None, Any]]" = [
+                (None, get(model_predictions, list, "pre_review")),
+                *filter(
+                    lambda review_predictions: not review_predictions[0].rejected,
+                    zip(reviews, get(model_predictions, list, "post_reviews")),
+                ),
+            ]
+
+            for review, model_predictions in reviewed_model_predictions:
+                predictions.extend(
+                    map(
+                        partial(prediction.from_v1_dict, document, model, review),
+                        model_predictions,
+                    )
+                )
+
+        return Result(
+            version=version,
+            submission_id=submission_id,
+            documents=(document,),
+            models=tuple(models),
+            predictions=predictions,
+            reviews=tuple(sorted(reviews)),
+        )
+
+    @staticmethod
+    def from_v3_dict(result: object) -> "Result":
+        """
+        Create a `Result` from a v3 result file dictionary.
+        """
+        normalize_v3_result(result)
+
+        version = get(result, int, "file_version")
+        submission_id = get(result, int, "submission_id")
+        submission_results = get(result, list, "submission_results")
+        modelgroup_metadata = get(result, dict, "modelgroup_metadata")
+        review_metadata = get(result, dict, "reviews")
+
+        processed_documents = map(Document.from_v3_dict, submission_results)
+        errored_files = get(result, dict, "errored_files").values()
+        failed_documents = map(Document.from_v3_errored_file, errored_files)
+        documents = sorted(chain(processed_documents, failed_documents))
+        models = sorted(map(ModelGroup.from_v3_dict, modelgroup_metadata.values()))
+        predictions: "PredictionList[Prediction]" = PredictionList()
+        reviews = sorted(map(Review.from_dict, review_metadata.values()))
+
+        for document_dict in submission_results:
+            document_id = get(document_dict, int, "submissionfile_id")
+            document = next(
+                filter(lambda document: document.id == document_id, documents)
+            )
+            reviewed_model_predictions: "list[tuple[Review | None, Any]]" = [
+                (None, get(document_dict, dict, "model_results", "ORIGINAL"))
+            ]
+
+            if reviews:
+                reviewed_model_predictions.append(
+                    (reviews[-1], get(document_dict, dict, "model_results", "FINAL"))
+                )
+
+            for review, model_section in reviewed_model_predictions:
+                for model_id, model_predictions in model_section.items():
+                    model = next(
+                        filter(lambda model: model.id == int(model_id), models)
+                    )
+                    predictions.extend(
+                        map(
+                            partial(prediction.from_v3_dict, document, model, review),
+                            model_predictions,
+                        )
+                    )
+
+        return Result(
+            version=version,
+            submission_id=submission_id,
+            documents=tuple(documents),
+            models=tuple(models),
+            predictions=predictions,
+            reviews=tuple(reviews),
+        )

indico_toolkit/results/review.py

@@ -0,0 +1,32 @@
+from dataclasses import dataclass
+from enum import Enum
+
+from .utilities import get
+
+
+class ReviewType(Enum):
+    ADMIN = "admin"
+    AUTO = "auto"
+    MANUAL = "manual"
+
+
+@dataclass(frozen=True, order=True)
+class Review:
+    id: int
+    reviewer_id: int
+    notes: str
+    rejected: bool
+    type: ReviewType
+
+    @staticmethod
+    def from_dict(review: object) -> "Review":
+        """
+        Create a `Review` from a review dictionary.
+        """
+        return Review(
+            id=get(review, int, "review_id"),
+            reviewer_id=get(review, int, "reviewer_id"),
+            notes=get(review, str, "review_notes"),
+            rejected=get(review, bool, "review_rejected"),
+            type=ReviewType(get(review, str, "review_type")),
+        )

indico_toolkit/results/utilities.py

@@ -0,0 +1,73 @@
+from collections.abc import Iterable, Iterator
+from typing import Callable, TypeVar
+
+from .errors import ResultError
+
+Value = TypeVar("Value")
+
+
+def get(result: object, value_type: "type[Value]", *keys: "str | int") -> Value:
+    """
+    Return the value obtained by traversing `result` using `keys` as indices if that
+    value has type `value_type`. Raise a `ResultError` otherwise.
+    """
+    for key in keys:
+        if isinstance(key, str) and isinstance(result, dict) and key in result:
+            result = result[key]
+        elif isinstance(key, int) and isinstance(result, list) and key < len(result):
+            result = result[key]
+        else:
+            raise ResultError(
+                f"result object `{type(result)!r}` does not contain key `{key!r}`"
+            )
+
+    if isinstance(result, value_type):
+        return result
+    else:
+        raise ResultError(
+            f"result object `{type(result)!r}` does not have a value for "
+            f"key `{key!r}` of type `{value_type}`"
+        )
+
+
+def has(result: object, value_type: "type[Value]", *keys: "str | int") -> bool:
+    """
+    Check if `result` can be traversed using `keys` to a value of type `value_type`.
+    """
+    for key in keys:
+        if isinstance(key, str) and isinstance(result, dict) and key in result:
+            result = result[key]
+        elif isinstance(key, int) and isinstance(result, list) and key < len(result):
+            result = result[key]
+        else:
+            return False
+
+    return isinstance(result, value_type)
+
+
+def nfilter(
+    predicates: "Iterable[Callable[[Value], bool]]", values: "Iterable[Value]"
+) -> "Iterator[Value]":
+    """
+    Apply multiple filter predicates to an iterable of values.
+
+    `nfilter([first, second, third], values)` is equivalent to
+    `filter(third, filter(second, filter(first, values)))`.
+    """
+    for predicate in predicates:
+        values = filter(predicate, values)
+
+    yield from values
+
+
+def omit(dictionary: object, *keys: str) -> "dict[str, Value]":
+    """
+    Return a shallow copy of `dictionary` with `keys` omitted.
+    """
+    if not isinstance(dictionary, dict):
+        return {}
+    return {
+        key: value
+        for key, value in dictionary.items()
+        if key not in keys
+    }  # fmt: skip

indico_toolkit/retry.py

@@ -0,0 +1,94 @@
+import asyncio
+import time
+from functools import wraps
+from inspect import iscoroutinefunction
+from random import random
+from typing import TYPE_CHECKING, overload
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable, Callable
+    from typing import ParamSpec, TypeVar
+
+    ArgumentsType = ParamSpec("ArgumentsType")
+    OuterReturnType = TypeVar("OuterReturnType")
+    InnerReturnType = TypeVar("InnerReturnType")
+
+
+def retry(
+    *errors: "type[Exception]",
+    count: int = 4,
+    wait: float = 1,
+    backoff: float = 4,
+    jitter: float = 0.5,
+) -> "Callable[[Callable[ArgumentsType, OuterReturnType]], Callable[ArgumentsType, OuterReturnType]]":  # noqa: E501
+    """
+    Decorate a function or coroutine to retry when it raises specified errors,
+    apply exponential backoff and jitter to the wait time,
+    and raise the last error if it retries too many times.
+
+    By default, the decorated function or coroutine will be retried up to 4 times over
+    the course of ~2 minutes (waiting 1, 4, 16, and 64 seconds; plus up to 50% jitter).
+
+    Arguments:
+        errors:  Retry the function when it raises one of these errors.
+        count:   Retry the function this many times.
+        wait:    Wait this many seconds after the first error before retrying.
+        backoff: Multiply the wait time by this amount for each additional error.
+        jitter:  Add a random amount of time (up to this percent as a decimal)
+                 to the wait time to prevent simultaneous retries.
+    """
+
+    def wait_time(times_retried: int) -> float:
+        """
+        Calculate the sleep time based on number of times retried.
+        """
+        return wait * backoff**times_retried * (1 + jitter * random())
+
+    @overload
+    def retry_decorator(
+        decorated: "Callable[ArgumentsType, Awaitable[InnerReturnType]]",
+    ) -> "Callable[ArgumentsType, Awaitable[InnerReturnType]]": ...
+    @overload
+    def retry_decorator(
+        decorated: "Callable[ArgumentsType, InnerReturnType]",
+    ) -> "Callable[ArgumentsType, InnerReturnType]": ...
+    def retry_decorator(
+        decorated: "Callable[ArgumentsType, InnerReturnType]",
+    ) -> "Callable[ArgumentsType, Awaitable[InnerReturnType]] | Callable[ArgumentsType, InnerReturnType]":  # noqa: E501
+        """
+        Decorate either a function or coroutine as appropriate.
+        """
+        if iscoroutinefunction(decorated):
+
+            @wraps(decorated)
+            async def retrying_coroutine(  # type: ignore[return]
+                *args: "ArgumentsType.args", **kwargs: "ArgumentsType.kwargs"
+            ) -> "InnerReturnType":
+                for times_retried in range(count + 1):
+                    try:
+                        return await decorated(*args, **kwargs)  # type: ignore[no-any-return]
+                    except errors:
+                        if times_retried >= count:
+                            raise
+
+                    await asyncio.sleep(wait_time(times_retried))
+
+            return retrying_coroutine
+        else:
+
+            @wraps(decorated)
+            def retrying_function(  # type: ignore[return]
+                *args: "ArgumentsType.args", **kwargs: "ArgumentsType.kwargs"
+            ) -> "InnerReturnType":
+                for times_retried in range(count + 1):
+                    try:
+                        return decorated(*args, **kwargs)
+                    except errors:
+                        if times_retried >= count:
+                            raise
+
+                    time.sleep(wait_time(times_retried))
+
+            return retrying_function
+
+    return retry_decorator

indico_toolkit/snapshots/__init__.py

@@ -0,0 +1,3 @@
+from .snapshot import Snapshot
+
+__all__ = ("Snapshot",)