feat: new eds.relation_detector_ffn trainable component

Author: percevalw

edsnlp/data/converters.py

@@ -321,7 +321,6 @@ def __call__(self, obj, tokenizer=None):
             for head in entities[from_entity_id]:
                 for tail in entities[to_entity_id]:
                     head._.rel.setdefault(relation_label, set()).add(tail)
-                    print("NEW REL FROM", head, "TO", tail, "WITH", relation_label)
 
         return doc
 

edsnlp/metrics/relations.py

@@ -4,15 +4,13 @@
 
 from edsnlp import registry
 from edsnlp.metrics import Examples, make_examples, prf
-from edsnlp.utils.span_getters import SpanGetterArg, get_spans
+from edsnlp.utils.span_getters import RelationCandidateGetter, get_spans
 from edsnlp.utils.typing import AsList
 
 
 def relations_scorer(
     examples: Examples,
-    head_getter: SpanGetterArg,
-    tail_getter: SpanGetterArg,
-    labels: AsList[str],
+    candidate_getter: AsList[RelationCandidateGetter],
     micro_key: str = "micro",
     filter_expr: Optional[str] = None,
 ):
@@ -24,12 +22,12 @@ def relations_scorer(
     examples : Examples
         The examples to score, either a tuple of (golds, preds) or a list of
         spacy.training.Example objects
-    head_getter : SpanGetterArg
-        The span getter to use to extract the relation heads from the document
-    tail_getter : SpanGetterArg
-        The span getter to use to extract the relation tails from the document
-    labels : Sequence[str]
-        The labels of the relations to evaluate
+    candidate_getter : AsList[RelationCandidateGetter]
+        The candidate getters to use to extract the possible relations from the
+        documents. Each candidate getter should be a dictionary with the keys
+        "head", "tail", and "labels". The "head" and "tail" keys should be
+        SpanGetterArg objects, and the "labels" key should be a list of strings
+        for these head-tail pairs.
     micro_key : str
         The key to use to store the micro-averaged results for spans of all types
     filter_expr : Optional[str]
@@ -49,36 +47,40 @@ def relations_scorer(
     total_pred_count = 0
     total_gold_count = 0
 
-    for eg_idx, eg in enumerate(examples):
-        pred_heads = [
-            ((h.start, h.end, h.label_), h)
-            for h in get_spans(eg.predicted, head_getter)
-        ]
-        pred_tails = [
-            ((t.start, t.end, t.label_), t)
-            for t in get_spans(eg.predicted, tail_getter)
-        ]
-        for (h_key, head), (t_key, tail) in product(pred_heads, pred_tails):
-            total_pred_count += 1
-            for label in labels:
-                if tail in head._.rel.get(label, ()):
-                    annotations[label][0].add((eg_idx, h_key, t_key, label))
-                    annotations[micro_key][0].add((eg_idx, h_key, t_key, label))
+    for candidate in candidate_getter:
+        head_getter = candidate["head"]
+        tail_getter = candidate["tail"]
+        labels = candidate["labels"]
+        for eg_idx, eg in enumerate(examples):
+            pred_heads = [
+                ((h.start, h.end, h.label_), h)
+                for h in get_spans(eg.predicted, head_getter)
+            ]
+            pred_tails = [
+                ((t.start, t.end, t.label_), t)
+                for t in get_spans(eg.predicted, tail_getter)
+            ]
+            for (h_key, head), (t_key, tail) in product(pred_heads, pred_tails):
+                total_pred_count += 1
+                for label in labels:
+                    if tail in head._.rel.get(label, ()):
+                        annotations[label][0].add((eg_idx, h_key, t_key, label))
+                        annotations[micro_key][0].add((eg_idx, h_key, t_key, label))
 
-        gold_heads = [
-            ((h.start, h.end, h.label_), h)
-            for h in get_spans(eg.reference, head_getter)
-        ]
-        gold_tails = [
-            ((t.start, t.end, t.label_), t)
-            for t in get_spans(eg.reference, tail_getter)
-        ]
-        for (h_key, head), (t_key, tail) in product(gold_heads, gold_tails):
-            total_gold_count += 1
-            for label in labels:
-                if tail in head._.rel.get(label, ()):
-                    annotations[label][1].add((eg_idx, h_key, t_key, label))
-                    annotations[micro_key][1].add((eg_idx, h_key, t_key, label))
+            gold_heads = [
+                ((h.start, h.end, h.label_), h)
+                for h in get_spans(eg.reference, head_getter)
+            ]
+            gold_tails = [
+                ((t.start, t.end, t.label_), t)
+                for t in get_spans(eg.reference, tail_getter)
+            ]
+            for (h_key, head), (t_key, tail) in product(gold_heads, gold_tails):
+                total_gold_count += 1
+                for label in labels:
+                    if tail in head._.rel.get(label, ()):
+                        annotations[label][1].add((eg_idx, h_key, t_key, label))
+                        annotations[micro_key][1].add((eg_idx, h_key, t_key, label))
 
     if total_pred_count != total_gold_count:
         raise ValueError(
@@ -101,15 +103,11 @@ def relations_scorer(
 class RelationsMetric:
     def __init__(
         self,
-        head_getter: SpanGetterArg,
-        tail_getter: SpanGetterArg,
-        labels: AsList[str],
+        candidate_getter: AsList[RelationCandidateGetter],
         micro_key: str = "micro",
         filter_expr: Optional[str] = None,
     ):
-        self.head_getter = head_getter
-        self.tail_getter = tail_getter
-        self.labels = labels
+        self.candidate_getter = candidate_getter
         self.micro_key = micro_key
         self.filter_expr = filter_expr
 
@@ -118,9 +116,7 @@ def __init__(
     def __call__(self, *examples: Any):
         return relations_scorer(
             examples,
-            head_getter=self.head_getter,
-            tail_getter=self.tail_getter,
-            labels=self.labels,
+            candidate_getter=self.candidate_getter,
             micro_key=self.micro_key,
             filter_expr=self.filter_expr,
         )

edsnlp/pipes/base.py

@@ -14,6 +14,7 @@
 from edsnlp.core import PipelineProtocol
 from edsnlp.core.registries import DraftPipe
 from edsnlp.utils.span_getters import (
+    RelationCandidateGetter,
     SpanGetter,  # noqa: F401
     SpanGetterArg,  # noqa: F401
     SpanSetter,
@@ -23,6 +24,7 @@
     validate_span_getter,  # noqa: F401
     validate_span_setter,
 )
+from edsnlp.utils.typing import AsList
 
 
 def value_getter(span: Span):
@@ -215,12 +217,22 @@ def __init__(
         nlp: PipelineProtocol = None,
         name: str = None,
         *args,
-        head_getter: SpanGetterArg,
-        tail_getter: SpanGetterArg,
-        labels: List[str],
+        candidate_getter: AsList[RelationCandidateGetter],
         **kwargs,
     ):
         super().__init__(nlp, name, *args, **kwargs)
-        self.head_getter: SpanGetter = validate_span_getter(head_getter)  # type: ignore
-        self.tail_getter: SpanGetter = validate_span_getter(tail_getter)  # type: ignore
-        self.labels = labels
+        self.candidate_getter = [
+            {
+                "head": validate_span_getter(candidate["head"]),
+                "tail": validate_span_getter(candidate["tail"]),
+                "labels": candidate["labels"],
+            }
+            for candidate in candidate_getter
+        ]
+        self.labels = sorted(
+            {
+                label
+                for candidate in self.candidate_getter
+                for label in candidate["labels"]
+            }
+        )

edsnlp/pipes/trainable/relation_detector_ffn/relation_detector_ffn.py

@@ -25,7 +25,7 @@
     SpanEmbeddingComponent,
     WordEmbeddingComponent,
 )
-from edsnlp.utils.span_getters import SpanGetterArg, get_spans
+from edsnlp.utils.span_getters import RelationCandidateGetter, get_spans
 from edsnlp.utils.typing import AsList
 
 
@@ -57,13 +57,17 @@ def make_ranges(starts, ends):
     if 0 in ends.shape:
         return ends
     sizes = ends - starts
+    mask = sizes > 0
     offsets = sizes.cumsum(0)
     offsets = offsets.roll(1)
     res = torch.ones(offsets[0], dtype=torch.long)
     offsets[0] = 0
-    res[offsets] = starts
-    res[offsets[1:]] -= ends[:-1] - 1
-    return res.cumsum(0)
+    masked_offsets = offsets[mask]
+    starts = starts[mask]
+    ends = ends[mask]
+    res[masked_offsets] = starts
+    res[masked_offsets[1:]] -= ends[:-1] - 1
+    return res.cumsum(0), offsets
 
 
 logger = logging.getLogger(__name__)
@@ -85,41 +89,56 @@ def make_ranges(starts, ends):
 """
 
 
+class MLP(torch.nn.Module):
+    def __init__(
+        self, input_dim: int, hidden_dim: int, output_dim: int, dropout_p: float = 0.0
+    ):
+        super().__init__()
+        self.hidden = torch.nn.Linear(input_dim, hidden_dim)
+        self.output = torch.nn.Linear(hidden_dim, output_dim)
+        self.dropout = torch.nn.Dropout(dropout_p)
+
+    def forward(self, x):
+        x = self.hidden(x)
+        x = F.relu(x)
+        x = self.dropout(x)
+        x = self.output(x)
+        return x
+
+
 class RelationDetectorFFN(
     TorchComponent[BatchOutput, FrameBatchInput],
     BaseRelationDetectorComponent,
 ):
     def __init__(
         self,
         nlp: Optional[PipelineProtocol] = None,
-        name: str = "rel_scope",
+        name: str = "relation_detector_ffn",
         *,
         span_embedding: SpanEmbeddingComponent,
         word_embedding: WordEmbeddingComponent,
-        head_getter: SpanGetterArg,
-        tail_getter: SpanGetterArg,
-        labels: AsList[str],
-        symmetric: bool = True,
+        candidate_getter: AsList[RelationCandidateGetter],
+        symmetric: bool = False,
+        hidden_size: int = 128,
+        dropout_p: float = 0.0,
     ):
         super().__init__(
             nlp=nlp,
             name=name,
-            head_getter=head_getter,
-            tail_getter=tail_getter,
-            labels=labels,
+            candidate_getter=candidate_getter,
         )
         self.span_embedding = span_embedding
         self.word_embedding = word_embedding
         self.symmetric = symmetric
-        # self.merge_mode = merge_mode
 
-        hidden_size = (
+        embed_size = (
             self.span_embedding.output_size * 2 + self.word_embedding.output_size
         )
         with warnings.catch_warnings():
             warnings.simplefilter("ignore", UserWarning)
             # self.head_projection = torch.nn.Linear(hidden_size, hidden_size)
             # self.tail_projection = torch.nn.Linear(hidden_size, hidden_size)
+            self.mlp = MLP(embed_size, hidden_size, hidden_size, dropout_p)
             self.classifier = torch.nn.Linear(hidden_size, len(self.labels))
 
     @property
@@ -155,22 +174,23 @@ def preprocess(self, doc: Doc, supervised: int = False) -> Dict[str, Any]:
         rel_labels = []
 
         all_spans = defaultdict(lambda: len(all_spans))
-        head_spans = list(get_spans(doc, self.head_getter))
-        tail_spans = list(get_spans(doc, self.tail_getter))
-
-        for head, tail in product(head_spans, tail_spans):
-            rel_head_idx.append(all_spans[head])
-            rel_tail_idx.append(all_spans[tail])
-            if supervised:
-                rel_labels.append(
-                    [
-                        (
-                            tail in head._.rel.get(lab, ())
-                            or (self.symmetric and head in tail._.rel.get(lab, ()))
-                        )
-                        for lab in self.labels
-                    ]
-                )
+
+        for candidate in self.candidate_getter:
+            head_spans = list(get_spans(doc, candidate["head"]))
+            tail_spans = list(get_spans(doc, candidate["tail"]))
+            for head, tail in product(head_spans, tail_spans):
+                rel_head_idx.append(all_spans[head])
+                rel_tail_idx.append(all_spans[tail])
+                if supervised:
+                    rel_labels.append(
+                        [
+                            (
+                                tail in head._.rel.get(lab, ())
+                                or (self.symmetric and head in tail._.rel.get(lab, ()))
+                            )
+                            for lab in self.labels
+                        ]
+                    )
 
         result = {
             "num_spans": len(all_spans),
@@ -231,17 +251,14 @@ def compute_inter_span_embeds(self, word_embeds, begins, ends, head_idx, tail_id
                 0, dim, dtype=word_embeds.dtype, device=word_embeds.device
             )
 
-        flat_begins = torch.minimum(
-            ends[head_idx],
-            ends[tail_idx],
-        )
-        flat_ends = torch.maximum(
-            begins[head_idx],
-            begins[tail_idx],
+        flat_begins = torch.minimum(ends[head_idx], ends[tail_idx])
+        flat_ends = torch.maximum(begins[head_idx], begins[tail_idx])
+        flat_begins, flat_ends = (
+            torch.minimum(flat_begins, flat_ends),
+            torch.maximum(flat_begins, flat_ends),
         )
         flat_embeds = word_embeds.view(-1, dim)
-        flat_indices = make_ranges(flat_begins, flat_ends)
-        flat_offsets = (flat_ends - flat_begins).cumsum(0).roll(1)
+        flat_indices, flat_offsets = make_ranges(flat_begins, flat_ends)
         flat_offsets[0] = 0
         inter_span_embeds = torch.nn.functional.embedding_bag(  # type: ignore
             input=flat_indices,
@@ -285,6 +302,7 @@ def forward(self, batch: FrameBatchInput) -> BatchOutput:
             ],
             dim=-1,
         )
+        rel_embeds = self.mlp(rel_embeds)
         logits = self.classifier(rel_embeds)
 
         losses = pred = None
@@ -321,12 +339,8 @@ def postprocess(
         Returns
         -------
         """
-        all_heads = [
-            prep["$spans"][idx] for prep in inputs for idx in prep["rel_heads"]
-        ]
-        all_tails = [
-            prep["$spans"][idx] for prep in inputs for idx in prep["rel_tails"]
-        ]
+        all_heads = [p["$spans"][idx] for p in inputs for idx in p["rel_heads"]]
+        all_tails = [p["$spans"][idx] for p in inputs for idx in p["rel_tails"]]
         for pair_idx, label_idx in results["pred"].nonzero(as_tuple=False).tolist():
             head = all_heads[pair_idx]
             tail = all_tails[pair_idx]

edsnlp/utils/span_getters.py

@@ -14,6 +14,7 @@
 
 from pydantic import NonNegativeInt
 from spacy.tokens import Doc, Span
+from typing_extensions import TypedDict
 
 from edsnlp import registry
 from edsnlp.utils.filter import filter_spans
@@ -321,3 +322,13 @@ def __call__(self, span: Union[Doc, Span]) -> Union[Span, List[Span]]:
             end = max(end, max_end_sent)
 
         return span.doc[start:end]
+
+
+RelationCandidateGetter = TypedDict(
+    "CandidateGetter",
+    {
+        "head": SpanGetterArg,
+        "tail": SpanGetterArg,
+        "labels": AsList[str],
+    },
+)