add a CRF decode path for inference

forward() needs labels to locate the first subword of each word, so there was
no way to run the tagger without them

Author: Kaushik-Kumar-CEG

etc/requirements-ml.txt

@@ -13,6 +13,10 @@ accelerate>=0.33
 # CRF layer on top of the token classifier
 pytorch-crf>=0.7.2
 
+# reading the trained weights and pulling the model from the hub
+safetensors>=0.4
+huggingface-hub>=0.24
+
 # 8 bit optimizer to fit deberta-large on a 16gb gpu, optional at runtime
 bitsandbytes>=0.43
 

etc/scripts/dataset_pipeline/phrase_model.py

@@ -8,6 +8,7 @@
 from transformers import AutoModel
 from transformers import Trainer
 
+from train_model import first_subword_positions
 from train_model import IGNORE_INDEX
 from train_model import LABELS
 
@@ -125,6 +126,25 @@ def forward(self, input_ids, attention_mask, labels=None):
 
         return result
 
+    def predict_words(self, input_ids, attention_mask, word_ids):
+        """Label id per word for a single rule, without labels
+
+        forward() needs labels to find the first subword of each word, which we
+        do not have at inference time, so take those positions from the
+        tokenizer word_ids and decode from the CRF directly
+        """
+        positions = first_subword_positions(word_ids)
+        if not positions:
+            return []
+
+        emissions = self.emissions(input_ids, attention_mask)
+        word_emissions = emissions[0, positions].unsqueeze(0).float()
+        if not self.use_crf:
+            return word_emissions.argmax(dim=-1)[0].tolist()
+
+        mask = torch.ones(word_emissions.shape[:2], dtype=torch.bool, device=emissions.device)
+        return self.crf.decode(word_emissions, mask=mask)[0]
+
     def pad_decoded(self, decoded, width, device):
         """Turn the variable length CRF paths into a padded tensor"""
         preds = torch.full((len(decoded), width), IGNORE_INDEX, dtype=torch.long, device=device)

etc/scripts/dataset_pipeline/test_train_model.py

@@ -7,8 +7,10 @@
 from train_model import align_labels
 from train_model import decode_row
 from train_model import extract_spans
+from train_model import first_subword_positions
 from train_model import IGNORE_INDEX
 from train_model import LABEL2ID
+from train_model import LABELS
 
 
 class FakeEncoding(dict):
@@ -85,6 +87,67 @@ def test_all_outside(self):
         assert encoding['labels'] == [IGNORE_INDEX, 0, 0, IGNORE_INDEX]
 
 
+class TestFirstSubwordPositions:
+
+    def test_skips_specials_and_continuations(self):
+        # CLS w0 w1a w1b w2 SEP
+        assert first_subword_positions([None, 0, 1, 1, 2, None]) == [1, 2, 4]
+
+    def test_without_special_tokens(self):
+        # depending on its files a tokenizer may add no CLS or SEP at all
+        assert first_subword_positions([0, 1, 2]) == [0, 1, 2]
+
+    def test_long_word(self):
+        assert first_subword_positions([None, 0, 0, 0, 1, None]) == [1, 4]
+
+    def test_nothing_to_tag(self):
+        assert first_subword_positions([None, None]) == []
+        assert first_subword_positions([]) == []
+
+
+def make_crf_tagger():
+    """A PhraseTagger with only the CRF built, the backbone needs a download"""
+    import torch
+    from torchcrf import CRF
+    from phrase_model import PhraseTagger
+
+    tagger = PhraseTagger.__new__(PhraseTagger)
+    torch.nn.Module.__init__(tagger)
+    tagger.use_crf = True
+    tagger.num_labels = len(LABELS)
+    tagger.crf = CRF(len(LABELS), batch_first=True)
+    # zero transitions make the decode a plain argmax, so the expected path is obvious
+    with torch.no_grad():
+        for param in tagger.crf.parameters():
+            param.zero_()
+    return tagger
+
+
+def test_predict_words_uses_the_first_subword():
+    import torch
+
+    tagger = make_crf_tagger()
+    # CLS w0 w1a w1b SEP, the score on the continuation subword must be ignored
+    emissions = torch.zeros((1, 5, len(LABELS)))
+    emissions[0, 1, LABEL2ID['B-REQ']] = 9.0
+    emissions[0, 2, LABEL2ID['E-REQ']] = 9.0
+    emissions[0, 3, LABEL2ID['S-REQ']] = 9.0
+    tagger.emissions = lambda input_ids, attention_mask: emissions
+
+    ids = torch.zeros((1, 5), dtype=torch.long)
+    tags = tagger.predict_words(ids, ids, [None, 0, 1, 1, None])
+    assert tags == [LABEL2ID['B-REQ'], LABEL2ID['E-REQ']]
+
+
+def test_predict_words_with_no_words():
+    import torch
+
+    tagger = make_crf_tagger()
+    tagger.emissions = lambda input_ids, attention_mask: torch.zeros((1, 2, len(LABELS)))
+    ids = torch.zeros((1, 2), dtype=torch.long)
+    assert tagger.predict_words(ids, ids, [None, None]) == []
+
+
 def test_viterbi_with_zero_transitions_is_argmax():
     import numpy as np
     from export_onnx import viterbi_decode

etc/scripts/dataset_pipeline/train_model.py

@@ -113,6 +113,26 @@ def align_labels(tokens, word_labels, tokenizer, max_length):
     return encoding
 
 
+def first_subword_positions(word_ids):
+    """Subword positions that start a new word
+
+    word_ids comes from a fast tokenizer encoding and has None for special
+    tokens. Training locates words through the labels, inference has none, so
+    the tagger uses these positions instead
+    """
+    positions = []
+    previous = None
+    for i, word_id in enumerate(word_ids):
+        if word_id is None:
+            previous = None
+            continue
+        if word_id != previous:
+            positions.append(i)
+        previous = word_id
+
+    return positions
+
+
 class PhraseDataset:
     """Reads a BIOES JSONL split and encodes each rule for the model"""