From d32ce2c8df7053c19061b709465cdcc765e45a15 Mon Sep 17 00:00:00 2001 From: Stefan Schweter Date: Mon, 18 Nov 2019 14:14:19 +0100 Subject: [PATCH] camembert: add wrapper for CamembertForTokenClassification --- transformers/modeling_camembert.py | 38 +++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/transformers/modeling_camembert.py b/transformers/modeling_camembert.py index 05538926e2..f302346f2d 100644 --- a/transformers/modeling_camembert.py +++ b/transformers/modeling_camembert.py @@ -20,7 +20,7 @@ from __future__ import (absolute_import, division, print_function, import logging -from .modeling_roberta import RobertaModel, RobertaForMaskedLM, RobertaForSequenceClassification, RobertaForMultipleChoice +from .modeling_roberta import RobertaModel, RobertaForMaskedLM, RobertaForSequenceClassification, RobertaForMultipleChoice, RobertaForTokenClassification from .configuration_camembert import CamembertConfig from .file_utils import add_start_docstrings @@ -255,3 +255,39 @@ class CamembertForMultipleChoice(RobertaForMultipleChoice): """ config_class = CamembertConfig pretrained_model_archive_map = CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_MAP + + +@add_start_docstrings("""CamemBERT 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. """, + CAMEMBERT_START_DOCSTRING, CAMEMBERT_INPUTS_DOCSTRING) +class CamembertForTokenClassification(RobertaForTokenClassification): + r""" + **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``: + Labels for computing the token classification loss. + Indices should be in ``[0, ..., config.num_labels - 1]``. + + Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs: + **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``: + Classification loss. + **scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.num_labels)`` + Classification scores (before SoftMax). + **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``) + list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings) + of shape ``(batch_size, sequence_length, hidden_size)``: + Hidden-states of the model at the output of each layer plus the initial embedding outputs. + **attentions**: (`optional`, returned when ``config.output_attentions=True``) + list of ``torch.FloatTensor`` (one for each layer) of shape ``(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:: + + tokenizer = CamembertTokenizer.from_pretrained('camembert-base') + model = CamembertForTokenClassification.from_pretrained('camembert-base') + input_ids = torch.tensor(tokenizer.encode("J'aime le camembert !", 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] + + """ + config_class = CamembertConfig + pretrained_model_archive_map = CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_MAP