wuenlp_tools.models.representation.moment 

  1from __future__ import annotations
  2
  3from collections import defaultdict
  4from typing import Type, List, Optional
  5
  6import numpy as np
  7import torch
  8from loguru import logger
  9
 10from wuenlp import UIMADocument
 11from wuenlp.impl.uima import UIMASpan, UIMASystemScene
 12
 13from wuenlp_tools.pipeline import PipelineProcessor, PipelineStep, PipelineCapability, UnitFilter, get_units
 14
 15_embedding_model = None
 16
 17
 18def _get_embedding_model():
 19    global _embedding_model
 20    if _embedding_model is None:
 21        from wuenlp_tools.vendored.momentfm import MOMENTPipeline
 22
 23        _embedding_model = MOMENTPipeline.from_pretrained(
 24            "AutonLab/MOMENT-1-small",
 25            model_kwargs={"task_name": "embedding"},
 26        )
 27        _embedding_model.init()
 28    return _embedding_model
 29
 30
 31def embed_ts(doc: UIMADocument, unit_type: Type[UIMASpan] | None, include_features: Optional[List] = None,
 32             overwrite: bool = False, unit_filter: UnitFilter | None = None):
 33    if include_features is None:
 34        raise NotImplementedError(
 35            f"Selecting all features automatically is not implemented yet. Pass an explicit list of features.")
 36
 37    if not overwrite and doc.document_annotation.additional_features["embedding"]:
 38        logger.info("Embedding already exists and overwrite is False. Returning unchanged.")
 39        return doc
 40
 41    unit: unit_type
 42
 43    features = defaultdict(list)
 44
 45    for unit in get_units(doc, unit_type, unit_filter):
 46        unit_additional_features = unit.additional_features
 47        for feature in include_features:
 48            if feature in unit_additional_features:
 49                feature_value = unit_additional_features[feature]
 50            else:
 51                feature_value = unit.__getattribute__(feature)
 52
 53            features[feature].append(feature_value)
 54
 55    numpy_features = {}
 56    for feature in features:
 57        numpy_features[feature] = np.array(features[feature])
 58
 59    numerical_features = {feature: values for feature, values in numpy_features.items() if
 60                          np.issubdtype(values.dtype, np.number) or np.issubdtype(values.dtype, bool)}
 61
 62    logger.info(f"Selected {len(numerical_features)} numerical features: {numerical_features.keys()}")
 63
 64    num_arrays = [numerical_features[feature] for feature in numerical_features]
 65    num_arrays = [
 66        x.reshape(-1, 1) if x.ndim == 1 else x
 67        for x in num_arrays
 68    ]
 69    value_array = torch.tensor(np.hstack(num_arrays)).unsqueeze(0).to(torch.float32)
 70
 71    embedding = list(_get_embedding_model()(x_enc=value_array).embeddings[0].numpy().astype(float))
 72
 73    doc.document_annotation.additional_features["embedding"] = embedding
 74
 75    return doc
 76
 77
 78class MomentEmbeddingProcessor(PipelineProcessor):
 79    include_features: Optional[List[str]] = None
 80
 81    def __init__(self, include_features: Optional[List[str]] = None) -> None:
 82        self.include_features = include_features
 83
 84    def set_include_features(self, include_features: List):
 85        self.include_features = include_features
 86
 87    def __call__(self, doc: UIMADocument, unit_type: Type[UIMASpan] | None, overwrite: bool = False,
 88                 unit_filter: UnitFilter | None = None, **kwargs) -> UIMADocument:
 89        if self.include_features is None:
 90            raise RuntimeError(f"You need to set include_features before calling this function.")
 91        return embed_ts(
 92            doc,
 93            unit_type,
 94            overwrite=overwrite,
 95            include_features=self.include_features,
 96            unit_filter=unit_filter,
 97        )
 98
 99
100MomentEmbedder = PipelineStep(
101    "MomentEmbedder",
102    MomentEmbeddingProcessor(),
103    unit_type=UIMASystemScene,
104    requires=[PipelineCapability.SCENES],
105    provides=[PipelineCapability.EMBEDDING],
106)
def embed_ts( doc: wuenlp.impl.uima.UIMANLPStructs.UIMADocument, unit_type: Optional[Type[wuenlp.impl.uima.UIMANLPStructs.UIMASpan]], include_features: Optional[List] = None, overwrite: bool = False, unit_filter: Optional[Callable[[wuenlp.impl.uima.UIMANLPStructs.UIMASpan], bool]] = None): 
32def embed_ts(doc: UIMADocument, unit_type: Type[UIMASpan] | None, include_features: Optional[List] = None,
33             overwrite: bool = False, unit_filter: UnitFilter | None = None):
34    if include_features is None:
35        raise NotImplementedError(
36            f"Selecting all features automatically is not implemented yet. Pass an explicit list of features.")
37
38    if not overwrite and doc.document_annotation.additional_features["embedding"]:
39        logger.info("Embedding already exists and overwrite is False. Returning unchanged.")
40        return doc
41
42    unit: unit_type
43
44    features = defaultdict(list)
45
46    for unit in get_units(doc, unit_type, unit_filter):
47        unit_additional_features = unit.additional_features
48        for feature in include_features:
49            if feature in unit_additional_features:
50                feature_value = unit_additional_features[feature]
51            else:
52                feature_value = unit.__getattribute__(feature)
53
54            features[feature].append(feature_value)
55
56    numpy_features = {}
57    for feature in features:
58        numpy_features[feature] = np.array(features[feature])
59
60    numerical_features = {feature: values for feature, values in numpy_features.items() if
61                          np.issubdtype(values.dtype, np.number) or np.issubdtype(values.dtype, bool)}
62
63    logger.info(f"Selected {len(numerical_features)} numerical features: {numerical_features.keys()}")
64
65    num_arrays = [numerical_features[feature] for feature in numerical_features]
66    num_arrays = [
67        x.reshape(-1, 1) if x.ndim == 1 else x
68        for x in num_arrays
69    ]
70    value_array = torch.tensor(np.hstack(num_arrays)).unsqueeze(0).to(torch.float32)
71
72    embedding = list(_get_embedding_model()(x_enc=value_array).embeddings[0].numpy().astype(float))
73
74    doc.document_annotation.additional_features["embedding"] = embedding
75
76    return doc
class MomentEmbeddingProcessor(wuenlp_tools.pipeline.AbstractPipelineProcessor[wuenlp.impl.uima.UIMANLPStructs.UIMADocument]): 
79class MomentEmbeddingProcessor(PipelineProcessor):
80    include_features: Optional[List[str]] = None
81
82    def __init__(self, include_features: Optional[List[str]] = None) -> None:
83        self.include_features = include_features
84
85    def set_include_features(self, include_features: List):
86        self.include_features = include_features
87
88    def __call__(self, doc: UIMADocument, unit_type: Type[UIMASpan] | None, overwrite: bool = False,
89                 unit_filter: UnitFilter | None = None, **kwargs) -> UIMADocument:
90        if self.include_features is None:
91            raise RuntimeError(f"You need to set include_features before calling this function.")
92        return embed_ts(
93            doc,
94            unit_type,
95            overwrite=overwrite,
96            include_features=self.include_features,
97            unit_filter=unit_filter,
98        )

Base class for protocol classes.

Protocol classes are defined as::

class Proto(Protocol):
    def meth(self) -> int:
        ...

Such classes are primarily used with static type checkers that recognize structural subtyping (static duck-typing).

For example::

class C:
    def meth(self) -> int:
        return 0

def func(x: Proto) -> int:
    return x.meth()

func(C())  # Passes static type check

See PEP 544 for details. Protocol classes decorated with @typing.runtime_checkable act as simple-minded runtime protocols that check only the presence of given attributes, ignoring their type signatures. Protocol classes can be generic, they are defined as::

class GenProto[T](Protocol):
    def meth(self) -> T:
        ...
MomentEmbeddingProcessor(include_features: Optional[List[str]] = None) 
82    def __init__(self, include_features: Optional[List[str]] = None) -> None:
83        self.include_features = include_features
include_features: Optional[List[str]] = None
def set_include_features(self, include_features: List): 
85    def set_include_features(self, include_features: List):
86        self.include_features = include_features
MomentEmbedder = PipelineStep('MomentEmbedder', processor=MomentEmbeddingProcessor, provides=['embedding'], requires=['scenes'])

Pipeline step MomentEmbedder (MomentEmbeddingProcessor).

Provides: embedding

  • embedding: Provides document-level Moment embedding in doc.document_annotation.additional_features['embedding'].

Requires: scenes

unit type UIMASystemScene.