Update pytorch-partial-tagger (#41)

Author: yasufumy

spacy_partial_tagger/pipeline.py

@@ -2,7 +2,8 @@
 
 import srsly
 import torch
-from partial_tagger.training import compute_partially_supervised_loss, create_tag_bitmap
+from partial_tagger.crf import functional as F
+from partial_tagger.training import create_tag_bitmap
 from sequence_label import LabelSet, SequenceLabel
 from spacy import util
 from spacy.errors import Errors
@@ -18,6 +19,40 @@
 from thinc.types import Floats2d, Floats4d
 
 
+def compute_partially_supervised_loss(
+    log_potentials: torch.Tensor,
+    tag_bitmap: torch.Tensor,
+    mask: torch.Tensor,
+    outside_index: int,
+    target_entity_ratio: float = 0.15,
+    entity_ratio_margin: float = 0.05,
+    balancing_coefficient: int = 10,
+) -> torch.Tensor:
+    with torch.enable_grad():
+        # log partition
+        log_Z = F.forward_algorithm(log_potentials)
+
+        # marginal probabilities
+        p = torch.autograd.grad(log_Z.sum(), log_potentials, create_graph=True)[0].sum(
+            dim=-1
+        )
+    p *= mask[..., None]
+
+    expected_entity_count = (
+        p[:, :, :outside_index].sum() + p[:, :, outside_index + 1 :].sum()
+    )
+    expected_entity_ratio = expected_entity_count / p.sum()
+    expected_entity_ratio_loss = torch.clamp(
+        (expected_entity_ratio - target_entity_ratio).abs() - entity_ratio_margin,
+        min=0,
+    )
+
+    score = F.multitag_sequence_score(log_potentials, tag_bitmap, mask)
+    supervised_loss = (log_Z - score).mean()
+
+    return supervised_loss + balancing_coefficient * expected_entity_ratio_loss
+
+
 class PartialEntityRecognizer(TrainablePipe):
     def __init__(
         self,

spacy_partial_tagger/util.py

@@ -1,26 +1,80 @@
-from typing import List, Tuple
+from typing import Dict, List, Tuple, cast
 
 import spacy_alignments as tokenizations
-from partial_tagger.decoders.viterbi import Constrainer, ViterbiDecoder
+import torch
+from partial_tagger.crf import functional as F
+from partial_tagger.crf.nn import CRF
+from partial_tagger.encoders.base import BaseEncoder
 from partial_tagger.encoders.transformer import TransformerModelEncoderFactory
-from partial_tagger.tagger import SequenceTagger
 from sequence_label import LabelSet
+from torch import nn
 from transformers import PreTrainedTokenizer
 
 
+class SequenceTagger(nn.Module):
+    def __init__(
+        self,
+        encoder: BaseEncoder,
+        padding_index: int,
+        start_states: Tuple[bool, ...],
+        end_states: Tuple[bool, ...],
+        transitions: Tuple[Tuple[bool, ...], ...],
+    ):
+        super().__init__()
+
+        self.encoder = encoder
+        self.crf = CRF(encoder.get_hidden_size())
+        self.start_constraints = nn.Parameter(
+            torch.tensor(start_states), requires_grad=False
+        )
+        self.end_constraints = nn.Parameter(
+            torch.tensor(end_states), requires_grad=False
+        )
+        self.transition_constraints = nn.Parameter(
+            torch.tensor(transitions), requires_grad=False
+        )
+        self.padding_index = padding_index
+
+    def __constrain(
+        self, log_potentials: torch.Tensor, mask: torch.Tensor
+    ) -> torch.Tensor:
+        return F.constrain_log_potentials(
+            log_potentials,
+            mask,
+            self.start_constraints,
+            self.end_constraints,
+            self.transition_constraints,
+        )
+
+    def forward(
+        self, inputs: Dict[str, torch.Tensor], mask: torch.Tensor
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        log_potentials = self.crf(self.encoder(inputs), mask)
+
+        contrained = self.__constrain(log_potentials, mask)
+
+        contrained.requires_grad_()
+
+        with torch.enable_grad():
+            _, tag_indices = F.decode(contrained)
+
+        return log_potentials, tag_indices * mask + self.padding_index * (~mask)
+
+    def predict(
+        self, inputs: Dict[str, torch.Tensor], mask: torch.Tensor
+    ) -> torch.Tensor:
+        return cast(torch.Tensor, self(inputs, mask)[1])
+
+
 def create_tagger(
     model_name: str, label_set: LabelSet, padding_index: int
 ) -> SequenceTagger:
     return SequenceTagger(
         TransformerModelEncoderFactory(model_name).create(label_set),
-        ViterbiDecoder(
-            padding_index,
-            Constrainer(
-                label_set.start_states,
-                label_set.end_states,
-                label_set.transitions,
-            ),
-        ),
+        padding_index,
+        label_set.start_states,
+        label_set.end_states,
+        label_set.transitions,
     )