wuenlp_tools.pipeline

  1from __future__ import annotations
  2
  3import inspect
  4from enum import Enum
  5from abc import ABC, abstractmethod
  6from collections import defaultdict
  7from hashlib import md5
  8from pathlib import Path
  9from typing import Callable, List, Type, Optional, Protocol, TypeVar, Dict, Tuple, Generic, Iterable, Union
 10
 11from wuenlp.impl.UIMANLPStructs import (
 12    UIMAAnnotation,
 13    UIMADocument,
 14    UIMASpan,
 15    UIMAEntity,
 16    UIMAEntityReference,
 17    UIMASystemScene,
 18)
 19from loguru import logger
 20from wuenlp.impl.uima import UIMACharacter
 21from wuenlp.impl.uima.extensions import ExtraTypesMixin, enrich_with_mixins
 22
 23from wuenlp_tools.keys import LazyKey
 24
 25InDocT = TypeVar("InDocT", UIMADocument, str, contravariant=True)
 26
 27UnitFilter = Callable[[UIMASpan], bool]
 28
 29
 30def get_units(
 31    doc: UIMADocument,
 32    unit_type: Type[UIMASpan] | None,
 33    unit_filter: UnitFilter | None = None,
 34) -> list[UIMASpan]:
 35    """Return units of ``unit_type``, optionally filtered, in CAS annotation-index order."""
 36    if unit_type is None:
 37        return []
 38    units = list(doc.get_annos_of_type(unit_type))
 39    if unit_filter is not None:
 40        units = [unit for unit in units if unit_filter(unit)]
 41    return units
 42
 43class PipelineCapability(str, Enum):
 44    """
 45    Declarative capabilities produced/required by pipeline steps.
 46
 47    These are used to auto-order steps and (when possible) auto-inject missing
 48    providers.
 49    """
 50
 51    def __new__(
 52        cls,
 53        value: str,
 54        description: str,
 55        affected_types: tuple[type, ...],
 56    ):
 57        obj = str.__new__(cls, value)
 58        obj._value_ = value
 59        obj._description = description
 60        obj._affected_types = affected_types
 61        return obj
 62
 63    PREPROCESS = (
 64        "preprocess",
 65        "Provides a fully initialized `UIMADocument` after preprocessing (tokenization, NER, rule-based setup).",
 66        (UIMADocument,),
 67    )
 68    COREF = (
 69        "coref",
 70        "Provides coreference annotations (`UIMAEntity`, `UIMAEntityReference`) needed by character-related steps.",
 71        (UIMAEntity, UIMAEntityReference),
 72    )
 73    CHARACTERS = (
 74        "characters",
 75        "Provides `UIMACharacter` annotations derived from entity/coreference information.",
 76        (UIMACharacter,),
 77    )
 78    MAIN_CHARACTERS = (
 79        "main_characters",
 80        "Provides main and important characters on document annotation (`main_character`, `important_characters`).",
 81        (UIMACharacter,),
 82    )
 83    SCENES = (
 84        "scenes",
 85        "Provides scene segmentation as `UIMASystemScene` annotations in `doc.system_scenes`.",
 86        (UIMASystemScene,),
 87    )
 88    SCENE_SUMMARIES = (
 89        "scene_summaries",
 90        "Provides per-scene summaries/keywords in additional features (e.g. `llama_summary`, `llama_keywords`).",
 91        (UIMASystemScene,),
 92    )
 93    SCENE_EMBEDDINGS = (
 94        "scene_embeddings",
 95        "Provides per-scene embeddings in additional features (e.g. `text_embedding`).",
 96        (UIMASystemScene,),
 97    )
 98    EMBEDDING = (
 99        "embedding",
100        "Provides document-level Moment embedding in `doc.document_annotation.additional_features['embedding']`.",
101        (UIMADocument,),
102    )
103
104    @property
105    def description(self) -> str:
106        return self._description
107
108    @property
109    def affected_types(self) -> tuple[type, ...]:
110        return self._affected_types
111
112
113_CapabilityLike = Union[PipelineCapability, str]
114
115
116def _normalize_capabilities(values: Optional[Iterable[_CapabilityLike]]) -> list[PipelineCapability]:
117    if not values:
118        return []
119    normalized: list[PipelineCapability] = []
120    for v in values:
121        if isinstance(v, PipelineCapability):
122            normalized.append(v)
123        else:
124            normalized.append(PipelineCapability(v))
125    return normalized
126
127
128class AbstractPipelineProcessor(Protocol[InDocT]):
129    @abstractmethod
130    def __call__(self, doc: InDocT, unit_type: Optional[Type[UIMASpan]], overwrite: bool = False,
131                 **kwargs) -> UIMADocument:
132        ...
133
134
135class PipelineProcessor(AbstractPipelineProcessor[UIMADocument]):
136    @abstractmethod
137    def __call__(self, doc: UIMADocument, unit_type: Type[UIMASpan] | None, overwrite: bool = False,
138                 **kwargs) -> UIMADocument:
139        ...
140
141
142class PipelinePreprocessor(AbstractPipelineProcessor[InDocT]):
143    @abstractmethod
144    def __call__(self, doc: UIMADocument | str, unit_type: Type[UIMASpan] | None, overwrite: bool = False,
145                 **kwargs) -> UIMADocument:
146        ...
147
148
149AbstractPipelineProcessorType = TypeVar('AbstractPipelineProcessorType', bound=AbstractPipelineProcessor)
150
151
152class PipelineStep(Generic[AbstractPipelineProcessorType]):
153    """A named processing unit that can be composed into a `Pipeline`.
154
155    Pass step instances into `Pipeline(steps=[...])`. Each step declares what it
156    `provides` and `requires` (`PipelineCapability`), which span type it operates
157    on, and which annotation features it writes.
158    """
159
160    name: Optional[str] = None
161    processor: AbstractPipelineProcessorType
162    unit_type: Optional[Type[UIMASpan]] = None
163    unit_filter: UnitFilter | None = None
164    overwrite: Optional[bool] = None
165    added_mixins: Optional[List[Type[ExtraTypesMixin]]] = None
166    added_features: Optional[List[str]] = None
167    added_additional_features: Optional[List[str]] = None
168
169    requires_api_key: Optional[LazyKey] = None
170    requires_paid_api_requests: Optional[bool] = None
171    provides: List[PipelineCapability]
172    requires: List[PipelineCapability]
173
174    def __init__(self,
175                 name: str,
176                 processor: AbstractPipelineProcessorType,
177                 unit_type: Optional[Type[UIMASpan]],
178                 unit_filter: UnitFilter | None = None,
179                 overwrite: Optional[bool] = None,
180                 added_mixins: Optional[List[Type[ExtraTypesMixin]] | Type[ExtraTypesMixin]] = None,
181                 added_features: Optional[List[str]] = None,
182                 added_additional_features: Optional[List[str]] = None,
183                 modified_types: Optional[List[Type[UIMAAnnotation]]] = None,
184                 needs_manual_merge: bool = False,
185                 requires_api_key: Optional[LazyKey] = None,
186                 requires_paid_api_requests: Optional[bool] = None,
187                 provides: Optional[List[_CapabilityLike]] = None,
188                 requires: Optional[List[_CapabilityLike]] = None,
189                 ):
190        self.name = name
191        self.processor = processor
192        self.unit_type = unit_type
193        self.unit_filter = unit_filter
194        self.overwrite = overwrite
195        self.added_mixins = [added_mixins] if isinstance(added_mixins, type) else (added_mixins or [])
196        self.added_features = added_features if added_features is not None else []
197        self.added_additional_features = added_additional_features if added_additional_features is not None else []
198        self.modified_types = modified_types if modified_types is not None else []
199        self.needs_manual_merge = needs_manual_merge
200        self.requires_api_key = requires_api_key
201        self.requires_paid_api_requests = requires_paid_api_requests
202        self.provides = _normalize_capabilities(provides)
203        self.requires = _normalize_capabilities(requires)
204        self.__doc__ = self._format_doc()
205
206    def __call__(self, doc: UIMADocument | str | Path, overwrite: Optional[bool] = None,
207                 DocType: Type[UIMADocument] = UIMADocument, **kwargs):
208        logger.info(f"Starting {self.name} pipeline step")
209
210        final_overwrite: bool
211
212        if overwrite is True:
213            final_overwrite = True
214        elif overwrite is None:
215            final_overwrite = self.overwrite if self.overwrite is not None else False
216        else:
217            # overwrite is False
218            final_overwrite = self.overwrite if self.overwrite is not None else False
219
220        if isinstance(doc, Path):
221            doc = DocType.from_xmi(doc)
222
223        unit_filter = kwargs.pop("unit_filter", self.unit_filter)
224        return self.processor(
225            doc,
226            unit_type=self.unit_type,
227            unit_filter=unit_filter,
228            overwrite=final_overwrite,
229            **kwargs,
230        )
231
232    def set_overwrite(self, overwrite: bool) -> "PipelineStep":
233        self.overwrite = overwrite
234        return self
235
236    def _format_doc(self) -> str:
237        lines = [f"Pipeline step **{self.name}** (`{type(self.processor).__name__}`)."]
238        proc_doc = type(self.processor).__dict__.get("__doc__")
239        if isinstance(proc_doc, str) and proc_doc.strip():
240            lines.append("")
241            lines.append(inspect.cleandoc(proc_doc).strip())
242        if self.provides:
243            lines.append("")
244            lines.append("**Provides:** " + ", ".join(f"`{c.value}`" for c in self.provides))
245            for cap in self.provides:
246                lines.append(f"- `{cap.value}`: {cap.description}")
247        if self.requires:
248            lines.append("")
249            lines.append("**Requires:** " + ", ".join(f"`{c.value}`" for c in self.requires))
250        details: list[str] = []
251        if self.unit_type is not None:
252            details.append(f"unit type `{self.unit_type.__name__}`")
253        if self.modified_types:
254            details.append(
255                "modifies " + ", ".join(f"`{t.__name__}`" for t in self.modified_types)
256            )
257        if self.added_additional_features:
258            details.append(
259                "additional features "
260                + ", ".join(f"`{f}`" for f in self.added_additional_features)
261            )
262        if self.added_features:
263            details.append("features " + ", ".join(f"`{f}`" for f in self.added_features))
264        if details:
265            lines.append("")
266            lines.append("; ".join(details) + ".")
267        return "\n".join(lines)
268
269    def __repr__(self) -> str:
270        parts = [repr(self.name), f"processor={type(self.processor).__name__}"]
271        if self.provides:
272            parts.append(f"provides={[c.value for c in self.provides]}")
273        if self.requires:
274            parts.append(f"requires={[c.value for c in self.requires]}")
275        return f"{type(self).__name__}({', '.join(parts)})"
276
277
278class Pipeline:
279    def __init__(self, steps: List[PipelineStep], intermediate_file_dir: Optional[Path] = None,
280                 overwrite: bool = False):
281        self.steps = steps
282        self.intermediate_file_dir = intermediate_file_dir
283        self.overwrite = overwrite
284
285        if intermediate_file_dir:
286            intermediate_file_dir.mkdir(exist_ok=True, parents=True)
287
288        force_overwrite = False
289
290        for step in self.steps:
291            if step.overwrite:
292                force_overwrite = True
293                logger.info(
294                    f"Step {step.name} is set to overwrite, setting all following steps (if any) to overwrite as well.")
295            else:
296                if force_overwrite:
297                    step.overwrite = force_overwrite
298                    logger.info(f"Setting step {step.name} to overwrite as a previous step is overwriting.")
299
300    def __call__(self, doc: UIMADocument | str | Path, DocType: Type = UIMADocument) -> UIMADocument | Path:
301        if isinstance(doc, Path):
302            logger.info(f"Processing {doc}")
303            current_class = DocType
304        elif isinstance(doc, UIMADocument):
305            logger.info(f"Processing {doc.path}")
306            current_class = type(doc)
307        else:
308            current_class = DocType
309        orig_overwrite = self.overwrite
310        result = None
311        for step in self.steps:
312            if step.added_mixins:
313                missing_mixins = [
314                    mixin for mixin in step.added_mixins if mixin not in current_class.__mro__
315                ]
316                if missing_mixins:
317                    current_class = enrich_with_mixins(current_class, missing_mixins)
318            if self.intermediate_file_dir:
319                text = doc.text if isinstance(doc, UIMADocument) else str(doc)
320                step_file = self.intermediate_file_dir / f"{md5(text.encode()).hexdigest()}_{step.name}.xmi.zip"
321            else:
322                step_file = None
323            if step_file and step_file.exists() and not self.overwrite and not step.overwrite:
324                logger.info(f"Using intermediate file for {step.name} from {self.intermediate_file_dir}")
325                result = step_file  # current_class.from_xmi(step_file)
326                continue
327            elif step_file and not step_file.exists():
328                self.overwrite = True
329                logger.info(
330                    f"Intermediate file for {step.name} does not exist, setting overwrite to True for all remaining steps.")
331            # TODO: This is highly experimental and needs to be tested/fixed
332            result = doc if result is None else result
333            step_output = step(result, overwrite=self.overwrite, DocType=current_class)
334            if not step.needs_manual_merge:
335                result = step_output
336            else:
337                if result is step_output:
338                    logger.warning("Document was unchanged by this step. Skipping manual merge.")
339                else:
340                    if isinstance(result, Path):
341                        result = current_class.from_xmi(result)
342                    for t in step.modified_types:
343                        for anno in result.get_annos_of_type(t):
344                            result.remove_annotation(anno)
345                        for anno in step_output.get_annos_of_type(t):
346                            result.add_anno(anno)
347            if self.intermediate_file_dir and isinstance(result, UIMADocument):
348                logger.info(f"Serializing intermediate file for {step.name} to {self.intermediate_file_dir}")
349                step_file.parent.mkdir(parents=True, exist_ok=True)
350                result.serialize(step_file)
351
352        self.overwrite = orig_overwrite
353
354        if result is not None:
355            if isinstance(result, UIMADocument):
356                return result
357            else:
358                return current_class.from_xmi(result)
359        else:
360            raise RuntimeError(f"No result produced")
361
362    @property
363    def added_features(self) -> Dict[Type[UIMAAnnotation], List]:
364        added_features: defaultdict = defaultdict(set)
365
366        for step in self.steps:
367            added_features[step.unit_type].update(step.added_features)
368
369        return dict(added_features)
370
371    @property
372    def added_additional_features(self) -> Dict[Type[UIMAAnnotation], List]:
373        added_additional_features: defaultdict = defaultdict(set)
374
375        for step in self.steps:
376            added_additional_features[step.unit_type].update(step.added_additional_features)
377
378        return dict(added_additional_features)
379
380    def __str__(self):
381        return " → ".join(step.name or type(step.processor).__name__ for step in self.steps)
382
383    def __repr__(self):
384        names = [step.name or type(step.processor).__name__ for step in self.steps]
385        return f"Pipeline({names})"
386
387    def __len__(self):
388        return len(self.steps)
389
390    steps: List[PipelineStep] = []