Create model.py

model.py

@@ -0,0 +1,54 @@
+from imports import *
+from all_datasets import *
+
+class PhoBertLstmCrf(RobertaForTokenClassification):
+    def __init__(self, config):
+        super(PhoBertLstmCrf, self).__init__(config=config)
+        self.num_labels = config.num_labels
+        self.lstm = nn.LSTM(input_size=config.hidden_size,
+                            hidden_size=config.hidden_size // 2,
+                            num_layers=1,
+                            batch_first=True,
+                            bidirectional=True)
+        self.crf = CRF(config.num_labels, batch_first=True)
+
+    @staticmethod
+    def sort_batch(src_tensor, lengths):
+        """
+        Sort a minibatch by the length of the sequences with the longest sequences first
+        return the sorted batch targes and sequence lengths.
+        This way the output can be used by pack_padd ed_sequences(...)
+        """
+        seq_lengths, perm_idx = lengths.sort(0, descending=True)
+        seq_tensor = src_tensor[perm_idx]
+        _, reversed_idx = perm_idx.sort(0, descending=False)
+        return seq_tensor, seq_lengths, reversed_idx
+
+    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, valid_ids=None,
+                label_masks=None):
+        seq_outputs = self.roberta(input_ids=input_ids,
+                                   token_type_ids=token_type_ids,
+                                   attention_mask=attention_mask,
+                                   head_mask=None)[0]
+
+        batch_size, max_len, feat_dim = seq_outputs.shape
+        seq_lens = torch.sum(label_masks, dim=-1)
+        range_vector = torch.arange(0, batch_size, dtype=torch.long, device=seq_outputs.device).unsqueeze(1)
+        seq_outputs = seq_outputs[range_vector, valid_ids]
+
+        sorted_seq_outputs, sorted_seq_lens, reversed_idx = self.sort_batch(src_tensor=seq_outputs,
+                                                                            lengths=seq_lens)
+        packed_words = pack_padded_sequence(sorted_seq_outputs, sorted_seq_lens.cpu(), True)
+        lstm_outs, _ = self.lstm(packed_words)
+        lstm_outs, _ = pad_packed_sequence(lstm_outs, batch_first=True, total_length=max_len)
+        seq_outputs = lstm_outs[reversed_idx]
+
+        seq_outputs = self.dropout(seq_outputs)
+        logits = self.classifier(seq_outputs)
+        seq_tags = self.crf.decode(logits, mask=label_masks != 0)
+
+        if labels is not None:
+            log_likelihood = self.crf(logits, labels, mask=label_masks.type(torch.uint8))
+            return NerOutput(loss=-1.0 * log_likelihood, tags=seq_tags, cls_metrics=seq_tags)
+        else:
+            return NerOutput(tags=seq_tags)