wuenlp_tools.models.suspense.api

  1from __future__ import annotations
  2
  3import os
  4import tempfile
  5from enum import Enum
  6from pathlib import Path
  7from typing import Type
  8
  9from loguru import logger
 10from wuenlp.impl.UIMANLPStructs import UIMADocument
 11from wuenlp.impl.uima import UIMAAnnotation, UIMASpan
 12from wuenlp.impl.uima.extensions.danger import DangerDocument, DangerMixin
 13from wuenlp.utils.internals import generate_xml_typesystem
 14
 15from wuenlp_tools.models.ssc.local import is_ssc_local_enabled, suspense_annotate_local
 16from wuenlp_tools.pipeline import PipelineStep, PipelineProcessor
 17from wuenlp_tools.utils.api import query_wuenlp_api
 18
 19
 20def _ensure_danger_typesystem() -> None:
 21    ts_path = DangerMixin.EXTRA_TS
 22    if ts_path.exists() and ts_path.read_text().strip():
 23        return
 24    ts_path.parent.mkdir(parents=True, exist_ok=True)
 25    xml = generate_xml_typesystem(DangerMixin)
 26    tmp = ts_path.with_suffix(".xml.tmp")
 27    tmp.write_text(xml)
 28    os.replace(tmp, ts_path)
 29
 30
 31def _load_suspense_response(path: str | Path) -> DangerDocument:
 32    _ensure_danger_typesystem()
 33    return DangerDocument.from_xmi(path, lenient=True, ignore_unknown_types=True)
 34
 35
 36class SuspenseTaskType(str, Enum):
 37    ANY_SUSPENSE = "anysuspense"
 38    DANGEROUS_SITUATION = "dangeroussituation"
 39    FEAR_DESCRIPTION = "feardescription"
 40
 41
 42TASK_ANNOTATION_TYPES: dict[SuspenseTaskType, Type[UIMAAnnotation]] = {
 43    SuspenseTaskType.ANY_SUSPENSE: DangerMixin.SystemAnySuspense,
 44    SuspenseTaskType.DANGEROUS_SITUATION: DangerMixin.SystemDangerousSituation,
 45    SuspenseTaskType.FEAR_DESCRIPTION: DangerMixin.SystemFearDescription,
 46}
 47
 48
 49def _iter_exact_annos(doc: UIMADocument, annotation_type: Type[UIMAAnnotation]):
 50    """Yield annotations of exactly ``annotation_type`` (no UIMA/Python subtypes)."""
 51    expected = getattr(annotation_type, "uima_type", None)
 52    for anno in doc.get_annos_of_type(annotation_type):
 53        if expected is None or getattr(anno, "uima_type", None) == expected:
 54            yield anno
 55
 56
 57def _has_annotations(doc: UIMADocument, annotation_type: Type[UIMAAnnotation]) -> bool:
 58    return any(True for _ in _iter_exact_annos(doc, annotation_type))
 59
 60
 61def _snapshot_annos(doc: UIMADocument, annotation_type: Type[UIMAAnnotation]) -> list[tuple[int, int, dict]]:
 62    snapped = []
 63    for anno in list(_iter_exact_annos(doc, annotation_type)):
 64        feats = dict(anno.additional_features) if anno.additional_features else {}
 65        snapped.append((anno.begin, anno.end, feats))
 66        doc.remove_annotation(anno)
 67    return snapped
 68
 69
 70def _restore_annos(doc: UIMADocument, annotation_type: Type[UIMAAnnotation], snapped: list[tuple[int, int, dict]]) -> None:
 71    for begin, end, feats in snapped:
 72        anno = doc.create_anno(annotation_type, begin, end, add_to_document=True)
 73        if feats:
 74            anno.additional_features.update(feats)
 75
 76
 77def _remap_api_any_suspense_to_task(doc: UIMADocument, task: SuspenseTaskType) -> UIMADocument:
 78    """Cluster suspense API may write all tasks as SystemAnySuspense; remap to the requested type."""
 79    expected = TASK_ANNOTATION_TYPES[task]
 80    if expected is DangerMixin.SystemAnySuspense:
 81        return doc
 82    if _has_annotations(doc, expected):
 83        return doc
 84    wrong = list(_iter_exact_annos(doc, DangerMixin.SystemAnySuspense))
 85    if not wrong:
 86        return doc
 87    logger.warning(
 88        f"Suspense API returned SystemAnySuspense for {task.value}; "
 89        f"remapping {len(wrong)} spans to {expected.__name__}"
 90    )
 91    for anno in wrong:
 92        begin, end = anno.begin, anno.end
 93        feats = dict(anno.additional_features) if anno.additional_features else {}
 94        doc.remove_annotation(anno)
 95        new_anno = doc.create_anno(expected, begin, end, add_to_document=True)
 96        if feats:
 97            new_anno.additional_features.update(feats)
 98    return doc
 99
100
101def suspense_api(doc: UIMADocument, task: SuspenseTaskType, overwrite: bool = False,
102                 base_url: str = "/suspense/"):
103    annotation_type = TASK_ANNOTATION_TYPES[task]
104    if _has_annotations(doc, annotation_type) and not overwrite:
105        logger.info(f"Document {doc} already has {task.value} annotations. Skipping")
106        return doc
107
108    if is_ssc_local_enabled():
109        logger.info(f"Running local SSC suspense model for {task.value}")
110        return suspense_annotate_local(doc, task.value, overwrite=overwrite)
111
112    # Danger/fear API currently emits SystemAnySuspense and may no-op when sibling
113    # suspense types are already present. Strip all suspense types for the request;
114    # restore siblings after remapping the response.
115    preserved: dict[SuspenseTaskType, list[tuple[int, int, dict]]] = {}
116    for sibling_task, sibling_type in TASK_ANNOTATION_TYPES.items():
117        if sibling_task is task:
118            # Drop existing target-type spans; this call replaces them.
119            _snapshot_annos(doc, sibling_type)
120            continue
121        preserved[sibling_task] = _snapshot_annos(doc, sibling_type)
122
123    document_full_extension = ".xmi.zip"
124    with tempfile.NamedTemporaryFile(suffix=document_full_extension, delete=False) as temp_file:
125        doc.serialize(temp_file.name)
126
127    if not base_url.endswith("/"):
128        base_url += "/"
129    with open(temp_file.name, "rb") as file:
130        files = {"file": (file.name, file)}
131        logger.info(f"Sending file to {task.value} suspense model")
132        response = query_wuenlp_api(url=f"{base_url}{task.value}", files=files)
133
134    with tempfile.NamedTemporaryFile(suffix=document_full_extension, delete=False) as temp_out_file:
135        with open(temp_out_file.name, "wb") as f:
136            f.write(response.content)
137
138    result = _load_suspense_response(temp_out_file.name)
139    result = _remap_api_any_suspense_to_task(result, task)
140    # Drop any echoed sibling suspense before restoring the caller's originals.
141    for sibling_task, sibling_type in TASK_ANNOTATION_TYPES.items():
142        if sibling_task is task:
143            continue
144        _snapshot_annos(result, sibling_type)
145        _restore_annos(result, sibling_type, preserved.get(sibling_task, []))
146    return result
147
148
149class SuspenseApiProcessor(PipelineProcessor):
150    def __init__(self, task: SuspenseTaskType):
151        self.task = task
152
153    def __call__(self, doc: UIMADocument, unit_type: Type[UIMASpan] | None, overwrite: bool = False, **kwargs):
154        return suspense_api(doc, task=self.task, overwrite=overwrite)
155
156
157def _make_step(name: str, task: SuspenseTaskType) -> PipelineStep:
158    return PipelineStep(
159        name,
160        SuspenseApiProcessor(task),
161        unit_type=None,
162        added_mixins=DangerMixin,
163        modified_types=[TASK_ANNOTATION_TYPES[task]],
164        needs_manual_merge=True,
165    )
166
167
168BERTAnySuspenseAnnotator = _make_step("BERT Any Suspense", SuspenseTaskType.ANY_SUSPENSE)
169BERTDangerousSituationAnnotator = _make_step("BERT Dangerous Situation", SuspenseTaskType.DANGEROUS_SITUATION)
170BERTFearDescriptionAnnotator = _make_step("BERT Fear Description", SuspenseTaskType.FEAR_DESCRIPTION)
171
172if __name__ == '__main__':
173    from argparse import ArgumentParser
174
175    parser = ArgumentParser()
176    parser.add_argument("--xmi", type=str, help="Path to the xmi file")
177    parser.add_argument("--task", type=str, default="dangeroussituation",
178                        choices=[t.value for t in SuspenseTaskType],
179                        help="Suspense task to run")
180    parser.add_argument("--overwrite", action="store_true", help="Overwrite existing annotations")
181    parser.add_argument("--output_folder", type=str, default=None, help="Output folder for the annotated file")
182
183    args = parser.parse_args()
184
185    xmi = Path(args.xmi)
186    doc = UIMADocument.from_xmi(xmi, ignore_unknown_types=True)
187    task = SuspenseTaskType(args.task)
188    output_folder = Path(args.output_folder) if args.output_folder else xmi.parent
189
190    logger.info(args)
191
192    doc = suspense_api(doc, task=task, overwrite=args.overwrite)
193    output_folder.mkdir(parents=True, exist_ok=True)
194    doc.serialize(output_folder / xmi.name)