Upload BERTWithDateLayerTokenClassification.py

BERTWithDateLayerTokenClassification.py

@@ -0,0 +1,154 @@
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers.models.bert.modeling_bert import BertPreTrainedModel, BertModel, \
+    BERT_INPUTS_DOCSTRING, _TOKENIZER_FOR_DOC, _CHECKPOINT_FOR_DOC, TokenClassifierOutput, _CONFIG_FOR_DOC
+from transformers.file_utils import (
+    add_code_sample_docstrings,
+    add_start_docstrings_to_model_forward,
+)
+
+
+class DateEmebdding(nn.Module):
+    """Construct the embeddings the creation date"""
+
+    def __init__(self, config):
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.date_vocab_size, config.date_hidden_size,
+                                            padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.date_max_position_embeddings, config.date_hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = nn.LayerNorm(config.date_hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer("position_ids", torch.arange(config.date_max_position_embeddings).expand((1, -1)))
+        self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
+        self.dense = nn.Linear(config.date_hidden_size, config.date_hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(
+            self, input_ids=None, position_ids=None, inputs_embeds=None, past_key_values_length=0
+    ):
+
+        try:
+
+            if input_ids is not None:
+                input_shape = input_ids.shape
+            else:
+                input_shape = inputs_embeds.size()[:-1]
+
+            seq_length = input_shape[1]
+
+            if position_ids is None:
+                position_ids = self.position_ids[:, past_key_values_length: seq_length + past_key_values_length]
+
+            if inputs_embeds is None:
+                inputs_embeds = self.word_embeddings(input_ids)
+
+            embeddings = inputs_embeds
+            if self.position_embedding_type == "absolute":
+                position_embeddings = self.position_embeddings(position_ids)
+                embeddings += position_embeddings
+            embeddings = self.LayerNorm(embeddings)
+            embeddings = self.dropout(embeddings)
+            max_over_time = torch.max(embeddings, 1)[0]
+        except Exception as ex:
+            print(type(ex).__name__, ex.args)
+            import pdb
+            pdb.set_trace()
+        return max_over_time
+
+
+class BERTWithDateLayerTokenClassification(BertPreTrainedModel):
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.bert = BertModel(config, add_pooling_layer=False)
+        self.date_embedding = DateEmebdding(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.date_hidden_size + config.hidden_size, config.num_labels)
+        # self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=TokenClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+            self,
+            input_ids=None,
+            input_date_ids=None,
+            attention_mask=None,
+            token_type_ids=None,
+            position_ids=None,
+            head_mask=None,
+            inputs_embeds=None,
+            labels=None,
+            output_attentions=None,
+            output_hidden_states=None,
+            return_dict=None,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for computing the token classification loss. Indices should be in ``[0, ..., config.num_labels -
+            1]``.
+        """
+        try:
+            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+            _, seq_length = input_ids.shape
+            outputs = self.bert(
+                input_ids,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids,
+                position_ids=position_ids,
+                head_mask=head_mask,
+                inputs_embeds=inputs_embeds,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+            )
+            date_output = self.date_embedding(input_date_ids)
+            sequence_output = torch.cat((outputs[0], date_output.unsqueeze(1).repeat(1, seq_length, 1)), 2)
+
+            sequence_output = self.dropout(sequence_output)
+            logits = self.classifier(sequence_output)
+
+            loss = None
+            if labels is not None:
+                loss_fct = CrossEntropyLoss()
+                # Only keep active parts of the loss
+                if attention_mask is not None:
+                    active_loss = attention_mask.view(-1) == 1
+                    active_logits = logits.view(-1, self.num_labels)
+                    active_labels = torch.where(
+                        active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+                    )
+                    loss = loss_fct(active_logits, active_labels)
+                else:
+                    loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+            if not return_dict:
+                output = (logits,) + outputs[2:]
+                return ((loss,) + output) if loss is not None else output
+        except:
+            import pdb
+            pdb.set_trace()
+            raise BrokenPipeError("Problems in forward pass")
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )