wuenlp_tools.models.scenes.segmentation

  1from __future__ import annotations
  2
  3import tempfile
  4from enum import Enum
  5from pathlib import Path
  6from typing import Callable, Type
  7
  8import requests
  9from loguru import logger
 10from wuenlp.impl.UIMANLPStructs import UIMADocument
 11from wuenlp.impl.uima import UIMASpan, UIMASystemScene
 12
 13from wuenlp_tools.models.ssc.local import is_ssc_local_enabled, scene_segment_local
 14from wuenlp_tools.pipeline import PipelineStep, PipelineProcessor, PipelineCapability
 15from wuenlp_tools.utils.api import query_wuenlp_api
 16
 17
 18class SceneModelType(str, Enum):
 19    BERT = "bert"
 20    LLAMA = "llama"
 21
 22
 23def scene_api(doc: UIMADocument, model: SceneModelType = SceneModelType.BERT, overwrite: bool = False,
 24              base_url: str = "/scenes/"):
 25    """Run remote scene segmentation and return a document with `UIMASystemScene` annotations."""
 26    if doc.system_scenes and not overwrite:
 27        logger.info(f"Document {doc} already has system scenes. Skipping")
 28        return doc
 29
 30    document_full_extension = ".xmi.zip"
 31    with tempfile.NamedTemporaryFile(suffix=document_full_extension, delete=False) as temp_file:
 32        doc.serialize(temp_file.name)
 33
 34    if not base_url.endswith("/"):
 35        base_url += "/"
 36    with open(temp_file.name, "rb") as file:
 37        files = {"file": (file.name, file)}
 38        logger.info(f"Sending file to {model.value} model")
 39        response = query_wuenlp_api(url=f"{base_url}{model.value}", files=files)
 40
 41    with tempfile.NamedTemporaryFile(suffix=document_full_extension, delete=False) as temp_out_file:
 42        with open(temp_out_file.name, "wb") as f:
 43            f.write(response.content)
 44
 45    return UIMADocument.from_xmi(temp_out_file.name)
 46
 47
 48class BERTSceneProcessor(PipelineProcessor):
 49    """Segment a document into scenes with the BERT scene model.
 50
 51    Uses the local SSC model when local inference is enabled, otherwise the
 52    remote `/scenes/bert` API. Writes `UIMASystemScene` annotations onto the
 53    document (`doc.system_scenes`).
 54    """
 55
 56    def __call__(self, doc: UIMADocument, unit_type: Type[UIMASpan] | None, overwrite: bool = False, **kwargs):
 57        if is_ssc_local_enabled():
 58            return scene_segment_local(doc, overwrite=overwrite)
 59        return scene_api(doc, model=SceneModelType.BERT, overwrite=overwrite)
 60
 61
 62class LlamaSceneProcessor(PipelineProcessor):
 63    """Segment a document into scenes with the Llama scene model via `/scenes/llama`."""
 64
 65    def __call__(self, doc: UIMADocument, unit_type=None, overwrite: bool = False, **kwargs):
 66        return scene_api(doc, model=SceneModelType.LLAMA, overwrite=overwrite)
 67
 68
 69BERTSceneSegmenter = PipelineStep("BERT Scene segmentation", BERTSceneProcessor(), unit_type=None,
 70                                  modified_types=[UIMASystemScene], needs_manual_merge=True, overwrite=True,
 71                                  provides=[PipelineCapability.SCENES], requires=[PipelineCapability.PREPROCESS])
 72LLAMASceneSegmenter = PipelineStep("Llama Scene segmentation", LlamaSceneProcessor(), unit_type=None,
 73                                   modified_types=[UIMASystemScene], needs_manual_merge=True, overwrite=True,
 74                                   provides=[PipelineCapability.SCENES], requires=[PipelineCapability.PREPROCESS])
 75
 76if __name__ == '__main__':
 77    from argparse import ArgumentParser
 78
 79    parser = ArgumentParser()
 80    parser.add_argument("--xmi", type=str, help="Path to the xmi file")
 81    parser.add_argument("--api", action="store_true", help="Use api scene segmenter")
 82    parser.add_argument("--model_type", type=str, default="bert", help="Model to use for scene segmentation")
 83    parser.add_argument("--overwrite", type=bool, help="Overwrite existing system scenes")
 84    parser.add_argument("--output_folder", type=str, default=None, help="Output folder for the annotated file")
 85
 86    args = parser.parse_args()
 87
 88    xmi = Path(args.xmi)
 89    doc = UIMADocument.from_xmi(xmi, ignore_unknown_types=True)
 90    api = args.api
 91    model_type = args.model_type
 92    overwrite = args.overwrite
 93    output_folder = Path(args.output_folder) if args.output_folder else xmi.parent
 94
 95    logger.info(args)
 96
 97    if doc.system_scenes and not overwrite:
 98        logger.info(f"Document {doc} already has system scenes and overwrite is not set. Exiting")
 99        exit(0)
100
101    if api:
102        doc = scene_api(doc, model=SceneModelType(model_type), overwrite=overwrite)
103    else:
104        doc = scene_segment_local(doc, overwrite=overwrite or False)
105    output_folder.mkdir(parents=True, exist_ok=True)
106    doc.serialize(output_folder / xmi.name)