wuenlp_tools.models.suspense.aggregation

  1from __future__ import annotations
  2
  3from typing import Type
  4
  5from loguru import logger
  6from wuenlp.impl.UIMANLPStructs import UIMADocument, UIMASystemScene
  7from wuenlp.impl.uima import UIMASpan
  8from wuenlp.impl.uima.extensions.danger import DangerMixin
  9
 10from wuenlp_tools.pipeline import (
 11    PipelineCapability,
 12    PipelineProcessor,
 13    PipelineStep,
 14    UnitFilter,
 15    get_units,
 16)
 17
 18# Feature keys used by the classification time-series builder
 19# (`experiments/classification/happy_end_2016/rich_timeseries.py`).
 20SUSPENSE_FEATURE_KEYS: tuple[str, ...] = (
 21    "bert_suspense_anysuspense",
 22    "bert_suspense_dangerous",
 23    "bert_suspense_fear",
 24)
 25
 26SUSPENSE_ANNOTATION_TYPES: dict[str, Type] = {
 27    "bert_suspense_anysuspense": DangerMixin.SystemAnySuspense,
 28    "bert_suspense_dangerous": DangerMixin.SystemDangerousSituation,
 29    "bert_suspense_fear": DangerMixin.SystemFearDescription,
 30}
 31
 32
 33def _spans_overlap(left_begin: int, left_end: int, right_begin: int, right_end: int) -> bool:
 34    return left_begin < right_end and right_begin < left_end
 35
 36
 37def _overlap_length(left_begin: int, left_end: int, right_begin: int, right_end: int) -> int:
 38    if not _spans_overlap(left_begin, left_end, right_begin, right_end):
 39        return 0
 40    return max(0, min(left_end, right_end) - max(left_begin, right_begin))
 41
 42
 43def _index_suspense_annotations(doc: UIMADocument) -> dict[str, list]:
 44    return {
 45        key: list(doc.get_annos_of_type(annotation_type))
 46        for key, annotation_type in SUSPENSE_ANNOTATION_TYPES.items()
 47    }
 48
 49
 50def scene_suspense_coverage(
 51        unit: UIMASpan,
 52        suspense_index: dict[str, list],
 53) -> dict[str, float]:
 54    """Covered fraction of ``unit`` per suspense type, clipped at 1."""
 55    unit_len = max(1, unit.end - unit.begin)
 56    values: dict[str, float] = {}
 57    for key in SUSPENSE_FEATURE_KEYS:
 58        covered = 0
 59        for annotation in suspense_index.get(key, []):
 60            covered += _overlap_length(unit.begin, unit.end, annotation.begin, annotation.end)
 61        values[key] = min(1.0, covered / unit_len)
 62    return values
 63
 64
 65class SuspenseAggregationProcessor(PipelineProcessor):
 66    """Aggregate BERT suspense spans to one coverage value per unit.
 67
 68    For each unit (default: scene), writes the fraction of the unit covered by
 69    ``SystemAnySuspense``, ``SystemDangerousSituation`` and
 70    ``SystemFearDescription`` spans, clipped at 1. This matches the thesis
 71    pipeline construction of one coverage value per suspense type per scene.
 72    """
 73
 74    def __call__(
 75            self,
 76            doc: UIMADocument,
 77            unit_type: Type[UIMASpan] | None,
 78            overwrite: bool = False,
 79            unit_filter: UnitFilter | None = None,
 80            **kwargs,
 81    ) -> UIMADocument:
 82        if unit_type is None:
 83            raise ValueError("SuspenseAggregationProcessor requires a non-None unit_type.")
 84
 85        units = get_units(doc, unit_type, unit_filter)
 86        if not units:
 87            logger.warning(f"No units found for {unit_type.__name__}; skipping suspense aggregation.")
 88            return doc
 89
 90        suspense_index = _index_suspense_annotations(doc)
 91        if not any(suspense_index.values()):
 92            logger.warning("No suspense span annotations found; writing zero coverage per unit.")
 93
 94        for unit in units:
 95            coverage = scene_suspense_coverage(unit, suspense_index)
 96            for key, value in coverage.items():
 97                if not overwrite and key in unit.additional_features:
 98                    continue
 99                unit.additional_features[key] = value
100
101        return doc
102
103
104def SuspenseAggregationStep(
105        unit_type: Type[UIMASpan] = UIMASystemScene,
106        name: str = "Suspense Aggregation",
107) -> PipelineStep[SuspenseAggregationProcessor]:
108    """Factory for the suspense coverage aggregator used in the thesis pipeline example."""
109    return PipelineStep(
110        name=name,
111        processor=SuspenseAggregationProcessor(),
112        unit_type=unit_type,
113        modified_types=[unit_type],
114        added_additional_features=list(SUSPENSE_FEATURE_KEYS),
115        added_mixins=DangerMixin,
116        requires=[PipelineCapability.SCENES],
117    )
118
119
120SuspenseAggregator = SuspenseAggregationStep