From e444648a302dc8520beec96356a4bf500944355c Mon Sep 17 00:00:00 2001 From: Suraj Patil Date: Thu, 28 May 2020 16:18:18 +0530 Subject: [PATCH] LongformerForTokenClassification (#4638) --- src/transformers/__init__.py | 1 + src/transformers/modeling_auto.py | 2 + src/transformers/modeling_longformer.py | 99 +++++++++++++++++++++++++ tests/test_modeling_longformer.py | 20 +++++ 4 files changed, 122 insertions(+) diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 0787232e76..33907741f9 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -326,6 +326,7 @@ if is_torch_available(): LongformerModel, LongformerForMaskedLM, LongformerForSequenceClassification, + LongformerForTokenClassification, LongformerForQuestionAnswering, LONGFORMER_PRETRAINED_MODEL_ARCHIVE_MAP, ) diff --git a/src/transformers/modeling_auto.py b/src/transformers/modeling_auto.py index db44e2e6ed..cc8604f560 100644 --- a/src/transformers/modeling_auto.py +++ b/src/transformers/modeling_auto.py @@ -106,6 +106,7 @@ from .modeling_longformer import ( LongformerForMaskedLM, LongformerForQuestionAnswering, LongformerForSequenceClassification, + LongformerForTokenClassification, LongformerModel, ) from .modeling_marian import MarianMTModel @@ -282,6 +283,7 @@ MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = OrderedDict( (CamembertConfig, CamembertForTokenClassification), (XLMConfig, XLMForTokenClassification), (XLMRobertaConfig, XLMRobertaForTokenClassification), + (LongformerConfig, LongformerForTokenClassification), (RobertaConfig, RobertaForTokenClassification), (BertConfig, BertForTokenClassification), (XLNetConfig, XLNetForTokenClassification), diff --git a/src/transformers/modeling_longformer.py b/src/transformers/modeling_longformer.py index 4524e7b21f..fcfbeecc99 100644 --- a/src/transformers/modeling_longformer.py +++ b/src/transformers/modeling_longformer.py @@ -971,3 +971,102 @@ class LongformerForQuestionAnswering(BertPreTrainedModel): outputs = (total_loss,) + outputs return outputs # (loss), start_logits, end_logits, (hidden_states), (attentions) + + +@add_start_docstrings( + """Longformer Model with a token classification head on top (a linear layer on top of + the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """, + LONGFORMER_START_DOCSTRING, +) +class LongformerForTokenClassification(BertPreTrainedModel): + config_class = LongformerConfig + pretrained_model_archive_map = LONGFORMER_PRETRAINED_MODEL_ARCHIVE_MAP + base_model_prefix = "longformer" + + def __init__(self, config): + super().__init__(config) + self.num_labels = config.num_labels + + self.longformer = LongformerModel(config) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + self.init_weights() + + @add_start_docstrings_to_callable(LONGFORMER_INPUTS_DOCSTRING) + def forward( + self, + input_ids=None, + attention_mask=None, + token_type_ids=None, + position_ids=None, + inputs_embeds=None, + labels=None, + ): + r""" + labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`): + Labels for computing the token classification loss. + Indices should be in ``[0, ..., config.num_labels - 1]``. + + Returns: + :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.LongformerConfig`) and inputs: + loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided) : + Classification loss. + scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`) + Classification scores (before SoftMax). + hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``): + Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer) + of shape :obj:`(batch_size, sequence_length, hidden_size)`. + + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``): + Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape + :obj:`(batch_size, num_heads, sequence_length, sequence_length)`. + + Attentions weights after the attention softmax, used to compute the weighted average in the self-attention + heads. + + Examples:: + + from transformers import LongformerTokenizer, LongformerForTokenClassification + import torch + + tokenizer = LongformerTokenizer.from_pretrained('longformer-base-4096') + model = LongformerForTokenClassification.from_pretrained('longformer-base-4096') + input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0) # Batch size 1 + labels = torch.tensor([1] * input_ids.size(1)).unsqueeze(0) # Batch size 1 + outputs = model(input_ids, labels=labels) + loss, scores = outputs[:2] + + """ + + outputs = self.longformer( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + ) + + sequence_output = outputs[0] + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output) + + outputs = (logits,) + outputs[2:] # add hidden states and attention if they are here + + 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)) + outputs = (loss,) + outputs + + return outputs # (loss), scores, (hidden_states), (attentions) diff --git a/tests/test_modeling_longformer.py b/tests/test_modeling_longformer.py index 9f0a4cbd36..dbf5e39ab3 100644 --- a/tests/test_modeling_longformer.py +++ b/tests/test_modeling_longformer.py @@ -30,6 +30,7 @@ if is_torch_available(): LongformerModel, LongformerForMaskedLM, LongformerForSequenceClassification, + LongformerForTokenClassification, LongformerForQuestionAnswering, ) @@ -212,6 +213,21 @@ class LongformerModelTester(object): self.parent.assertListEqual(list(result["logits"].size()), [self.batch_size, self.num_labels]) self.check_loss_output(result) + def create_and_check_longformer_for_token_classification( + self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels + ): + config.num_labels = self.num_labels + model = LongformerForTokenClassification(config=config) + model.to(torch_device) + model.eval() + loss, logits = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels) + result = { + "loss": loss, + "logits": logits, + } + self.parent.assertListEqual(list(result["logits"].size()), [self.batch_size, self.seq_length, self.num_labels]) + self.check_loss_output(result) + def prepare_config_and_inputs_for_common(self): config_and_inputs = self.prepare_config_and_inputs() ( @@ -278,6 +294,10 @@ class LongformerModelTest(ModelTesterMixin, unittest.TestCase): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_longformer_for_sequence_classification(*config_and_inputs) + def test_for_token_classification(self): + config_and_inputs = self.model_tester.prepare_config_and_inputs() + self.model_tester.create_and_check_longformer_for_token_classification(*config_and_inputs) + class LongformerModelIntegrationTest(unittest.TestCase): @slow