From 2b096508852344c1d361f36e251b1b74253120fd Mon Sep 17 00:00:00 2001 From: gilad19 Date: Tue, 26 Jul 2022 11:11:32 +0300 Subject: [PATCH] Add ViltForTokenClassification e.g. for Named-Entity-Recognition (NER) (#17924) * Add ViltForTokenClassification e.g. for Named-Entity-Recognition (NER) * Add ViltForTokenClassification e.g. for Named-Entity-Recognition (NER) * provide classifier only text hidden states * add test_for_token_classification * Update src/transformers/models/vilt/modeling_vilt.py Co-authored-by: NielsRogge <48327001+NielsRogge@users.noreply.github.com> * Update src/transformers/models/vilt/modeling_vilt.py Co-authored-by: NielsRogge <48327001+NielsRogge@users.noreply.github.com> * Update src/transformers/models/vilt/modeling_vilt.py Co-authored-by: NielsRogge <48327001+NielsRogge@users.noreply.github.com> * Update src/transformers/models/vilt/modeling_vilt.py Co-authored-by: NielsRogge <48327001+NielsRogge@users.noreply.github.com> * add test_for_token_classification Co-authored-by: gfuchs Co-authored-by: NielsRogge <48327001+NielsRogge@users.noreply.github.com> --- docs/source/en/model_doc/vilt.mdx | 5 ++ src/transformers/__init__.py | 2 + src/transformers/models/vilt/__init__.py | 2 + src/transformers/models/vilt/modeling_vilt.py | 88 +++++++++++++++++++ src/transformers/utils/dummy_pt_objects.py | 7 ++ tests/models/vilt/test_modeling_vilt.py | 32 ++++++- utils/check_repo.py | 1 + 7 files changed, 133 insertions(+), 4 deletions(-) diff --git a/docs/source/en/model_doc/vilt.mdx b/docs/source/en/model_doc/vilt.mdx index 34397e7b3c..b6b87e7aa5 100644 --- a/docs/source/en/model_doc/vilt.mdx +++ b/docs/source/en/model_doc/vilt.mdx @@ -87,3 +87,8 @@ This model was contributed by [nielsr](https://huggingface.co/nielsr). The origi [[autodoc]] ViltForImageAndTextRetrieval - forward + +## ViltForTokenClassification + +[[autodoc]] ViltForTokenClassification + - forward diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py index 700636bb86..8f4e4840f9 100755 --- a/src/transformers/__init__.py +++ b/src/transformers/__init__.py @@ -1816,6 +1816,7 @@ else: "VILT_PRETRAINED_MODEL_ARCHIVE_LIST", "ViltForImageAndTextRetrieval", "ViltForImagesAndTextClassification", + "ViltForTokenClassification", "ViltForMaskedLM", "ViltForQuestionAnswering", "ViltLayer", @@ -4317,6 +4318,7 @@ if TYPE_CHECKING: ViltForImagesAndTextClassification, ViltForMaskedLM, ViltForQuestionAnswering, + ViltForTokenClassification, ViltLayer, ViltModel, ViltPreTrainedModel, diff --git a/src/transformers/models/vilt/__init__.py b/src/transformers/models/vilt/__init__.py index 3861b081be..d05318202b 100644 --- a/src/transformers/models/vilt/__init__.py +++ b/src/transformers/models/vilt/__init__.py @@ -42,6 +42,7 @@ else: "VILT_PRETRAINED_MODEL_ARCHIVE_LIST", "ViltForImageAndTextRetrieval", "ViltForImagesAndTextClassification", + "ViltForTokenClassification", "ViltForMaskedLM", "ViltForQuestionAnswering", "ViltLayer", @@ -74,6 +75,7 @@ if TYPE_CHECKING: ViltForImagesAndTextClassification, ViltForMaskedLM, ViltForQuestionAnswering, + ViltForTokenClassification, ViltLayer, ViltModel, ViltPreTrainedModel, diff --git a/src/transformers/models/vilt/modeling_vilt.py b/src/transformers/models/vilt/modeling_vilt.py index 6e7e9f154e..f20573d0d5 100755 --- a/src/transformers/models/vilt/modeling_vilt.py +++ b/src/transformers/models/vilt/modeling_vilt.py @@ -32,6 +32,7 @@ from ...modeling_outputs import ( MaskedLMOutput, ModelOutput, SequenceClassifierOutput, + TokenClassifierOutput, ) from ...modeling_utils import PreTrainedModel from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer @@ -1402,3 +1403,90 @@ class ViltForImagesAndTextClassification(ViltPreTrainedModel): hidden_states=hidden_states, attentions=attentions, ) + + +@add_start_docstrings( + """ + ViLT Model with a token classification head on top (a linear layer on top of the final hidden-states of the text + tokens) e.g. for Named-Entity-Recognition (NER) tasks. + """, + VILT_START_DOCSTRING, +) +class ViltForTokenClassification(ViltPreTrainedModel): + + _keys_to_ignore_on_load_unexpected = [r"pooler"] + + def __init__(self, config): + super().__init__(config) + + self.num_labels = config.num_labels + self.vilt = ViltModel(config, add_pooling_layer=False) + + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.classifier = nn.Linear(config.hidden_size, config.num_labels) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(VILT_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=TokenClassifierOutput, config_class=_CONFIG_FOR_DOC) + def forward( + self, + input_ids=None, + attention_mask=None, + token_type_ids=None, + pixel_values=None, + pixel_mask=None, + head_mask=None, + inputs_embeds=None, + image_embeds=None, + labels=None, + output_attentions=None, + output_hidden_states=None, + return_dict=None, + ): + r""" + labels (`torch.LongTensor` of shape `(batch_size, text_sequence_length)`, *optional*): + Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`. + + Returns: + """ + + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + outputs = self.vilt( + input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + pixel_values=pixel_values, + pixel_mask=pixel_mask, + head_mask=head_mask, + inputs_embeds=inputs_embeds, + image_embeds=image_embeds, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] + + text_input_size = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] + + sequence_output = self.dropout(sequence_output) + logits = self.classifier(sequence_output[:, :text_input_size]) + + loss = None + if labels is not None: + loss_fct = CrossEntropyLoss() + 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 + + return TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/src/transformers/utils/dummy_pt_objects.py b/src/transformers/utils/dummy_pt_objects.py index 043f4cdb9d..fb32886659 100644 --- a/src/transformers/utils/dummy_pt_objects.py +++ b/src/transformers/utils/dummy_pt_objects.py @@ -4725,6 +4725,13 @@ class ViltForQuestionAnswering(metaclass=DummyObject): requires_backends(self, ["torch"]) +class ViltForTokenClassification(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + class ViltLayer(metaclass=DummyObject): _backends = ["torch"] diff --git a/tests/models/vilt/test_modeling_vilt.py b/tests/models/vilt/test_modeling_vilt.py index 1a2f95d0e6..82aa076747 100644 --- a/tests/models/vilt/test_modeling_vilt.py +++ b/tests/models/vilt/test_modeling_vilt.py @@ -37,6 +37,7 @@ if is_torch_available(): ViltForImagesAndTextClassification, ViltForMaskedLM, ViltForQuestionAnswering, + ViltForTokenClassification, ViltModel, ) from transformers.models.vilt.modeling_vilt import VILT_PRETRAINED_MODEL_ARCHIVE_LIST @@ -173,6 +174,23 @@ class ViltModelTester: result.last_hidden_state.shape, (self.batch_size, self.expected_seq_len, self.hidden_size) ) + def create_and_check_for_token_classification( + self, + config, + input_ids, + token_type_ids, + input_mask, + pixel_values, + token_labels, + ): + model = ViltForTokenClassification(config=config) + model.to(torch_device) + model.eval() + result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, pixel_values=pixel_values) + result = model(input_ids, token_type_ids=token_type_ids, pixel_values=pixel_values) + result = model(input_ids, pixel_values=pixel_values) + self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels)) + def prepare_config_and_inputs_for_common(self): config_and_inputs = self.prepare_config_and_inputs() ( @@ -204,6 +222,7 @@ class ViltModelTest(ModelTesterMixin, unittest.TestCase): ViltForQuestionAnswering, ViltForImageAndTextRetrieval, ViltForMaskedLM, + ViltForTokenClassification, ) if is_torch_available() else () @@ -216,15 +235,12 @@ class ViltModelTest(ModelTesterMixin, unittest.TestCase): def _prepare_for_class(self, inputs_dict, model_class, return_labels=False): inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels) - # if model_class.__name__ == "ViltForNaturalLanguageVisualReasonining": - # inputs_dict["pixel_values"] = floats_tensor([self.model_tester.batch_size, self.model_tester.num_images, self.model_tester.num_channels, self.model_tester.image_size, self.model_tester.image_size]) - if return_labels: if model_class.__name__ == "ViltForQuestionAnswering": inputs_dict["labels"] = torch.zeros( self.model_tester.batch_size, self.model_tester.num_labels, device=torch_device ) - elif model_class.__name__ == "ViltForMaskedLM": + elif model_class.__name__ in ["ViltForMaskedLM", "ViltForTokenClassification"]: inputs_dict["labels"] = torch.zeros( (self.model_tester.batch_size, self.model_tester.seq_length), dtype=torch.long, device=torch_device ) @@ -246,6 +262,10 @@ class ViltModelTest(ModelTesterMixin, unittest.TestCase): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_model(*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_for_token_classification(*config_and_inputs) + def test_training(self): if not self.model_tester.is_training: return @@ -503,6 +523,10 @@ class ViltForImagesAndTextClassificationModelTest(ViltModelTest, unittest.TestCa def test_model(self): pass + @unittest.skip("We only test the model that takes in multiple images") + def test_for_token_classification(self): + pass + # We will verify our results on an image of cute cats def prepare_img(): diff --git a/utils/check_repo.py b/utils/check_repo.py index 47fee16313..00cc6a048b 100644 --- a/utils/check_repo.py +++ b/utils/check_repo.py @@ -131,6 +131,7 @@ IGNORE_NON_AUTO_CONFIGURED = PRIVATE_MODELS.copy() + [ "ViltForQuestionAnswering", "ViltForImagesAndTextClassification", "ViltForImageAndTextRetrieval", + "ViltForTokenClassification", "ViltForMaskedLM", "XGLMEncoder", "XGLMDecoder",