codingutf8 2018 the google ai language team s and the huggingface inc team c 2018 nvidia corporation licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache orglicenseslicense2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license named entity recognition finetuning utilities to work with conll2003 task import logging import os from dataclasses import dataclass from enum import enum from typing import list optional union from filelock import filelock from transformers import pretrainedtokenizer istfavailable istorchavailable logger logging getloggername dataclass class inputexample guid str words liststr labels optionalliststr dataclass class inputfeatures inputids listint attentionmask listint tokentypeids optionallistint none labelids optionallistint none class splitenum train train dev dev test test class tokenclassificationtask staticmethod def readexamplesfromfiledatadir mode unionsplit str listinputexample raise notimplementederror staticmethod def getlabelspath str liststr raise notimplementederror staticmethod def convertexamplestofeatures examples listinputexample labellist liststr maxseqlength int tokenizer pretrainedtokenizer clstokenatendfalse clstokencls clstokensegmentid1 septokensep septokenextrafalse padonleftfalse padtoken0 padtokensegmentid0 padtokenlabelid100 sequenceasegmentid0 maskpaddingwithzerotrue listinputfeatures todo clean up all this to leverage builtin features of tokenizers labelmap label i for i label in enumeratelabellist features for exindex example in enumerateexamples if exindex 10000 0 logger infowriting example d of d exindex lenexamples tokens labelids for word label in zipexample words example labels wordtokens tokenizer tokenizeword bertbasemultilingualcased sometimes output nothing when calling tokenize with just a space if lenwordtokens 0 tokens extendwordtokens use the real label id for the first token of the word and padding ids for the remaining tokens labelids extendlabelmaplabel padtokenlabelid lenwordtokens 1 account for cls and sep with 2 and with 3 for roberta specialtokenscount tokenizer numspecialtokenstoadd if lentokens maxseqlength specialtokenscount tokens tokens maxseqlength specialtokenscount labelids labelids maxseqlength specialtokenscount the convention in bert is a for sequence pairs tokens cls is this jack son ville sep no it is not sep typeids 0 0 0 0 0 0 0 0 1 1 1 1 1 1 b for single sequences tokens cls the dog is hairy sep typeids 0 0 0 0 0 0 0 where typeids are used to indicate whether this is the first sequence or the second sequence the embedding vectors for type0 and type1 were learned during pretraining and are added to the wordpiece embedding vector and position vector this is not strictly necessary since the sep token unambiguously separates the sequences but it makes it easier for the model to learn the concept of sequences for classification tasks the first vector corresponding to cls is used as the sentence vector note that this only makes sense because the entire model is finetuned tokens septoken labelids padtokenlabelid if septokenextra roberta uses an extra separator bw pairs of sentences tokens septoken labelids padtokenlabelid segmentids sequenceasegmentid lentokens if clstokenatend tokens clstoken labelids padtokenlabelid segmentids clstokensegmentid else tokens clstoken tokens labelids padtokenlabelid labelids segmentids clstokensegmentid segmentids inputids tokenizer converttokenstoidstokens the mask has 1 for real tokens and 0 for padding tokens only real tokens are attended to inputmask 1 if maskpaddingwithzero else 0 leninputids zeropad up to the sequence length paddinglength maxseqlength leninputids if padonleft inputids padtoken paddinglength inputids inputmask 0 if maskpaddingwithzero else 1 paddinglength inputmask segmentids padtokensegmentid paddinglength segmentids labelids padtokenlabelid paddinglength labelids else inputids padtoken paddinglength inputmask 0 if maskpaddingwithzero else 1 paddinglength segmentids padtokensegmentid paddinglength labelids padtokenlabelid paddinglength assert leninputids maxseqlength assert leninputmask maxseqlength assert lensegmentids maxseqlength assert lenlabelids maxseqlength if exindex 5 logger info example logger infoguid s example guid logger infotokens s joinstrx for x in tokens logger infoinputids s joinstrx for x in inputids logger infoinputmask s joinstrx for x in inputmask logger infosegmentids s joinstrx for x in segmentids logger infolabelids s joinstrx for x in labelids if tokentypeids not in tokenizer modelinputnames segmentids none features append inputfeatures inputidsinputids attentionmaskinputmask tokentypeidssegmentids labelidslabelids return features if istorchavailable import torch from torch import nn from torch utils data import dataset class tokenclassificationdatasetdataset features listinputfeatures padtokenlabelid int nn crossentropyloss ignoreindex use cross entropy ignoreindex as padding label id so that only real label ids contribute to the loss later def init self tokenclassificationtask tokenclassificationtask datadir str tokenizer pretrainedtokenizer labels liststr modeltype str maxseqlength optionalint none overwritecachefalse mode split split train load data features from cache or dataset file cachedfeaturesfile os path join datadir cached formatmode value tokenizer class name strmaxseqlength make sure only the first process in distributed training processes the dataset and the others will use the cache lockpath cachedfeaturesfile lock with filelocklockpath if os path existscachedfeaturesfile and not overwritecache logger infofloading features from cached file cachedfeaturesfile self features torch loadcachedfeaturesfile else logger infofcreating features from dataset file at datadir examples tokenclassificationtask readexamplesfromfiledatadir mode todo clean up all this to leverage builtin features of tokenizers self features tokenclassificationtask convertexamplestofeatures examples labels maxseqlength tokenizer clstokenatendboolmodeltype in xlnet xlnet has a cls token at the end clstokentokenizer clstoken clstokensegmentid2 if modeltype in xlnet else 0 septokentokenizer septoken septokenextrafalse roberta uses an extra separator bw pairs of sentences cf github compytorchfairseqcommit1684e166e3da03f5b600dbb7855cb98ddfcd0805 padonleftbooltokenizer paddingside left padtokentokenizer padtokenid padtokensegmentidtokenizer padtokentypeid padtokenlabelidself padtokenlabelid logger infofsaving features into cached file cachedfeaturesfile torch saveself features cachedfeaturesfile def lenself return lenself features def getitemself i inputfeatures return self featuresi if istfavailable import tensorflow as tf class tftokenclassificationdataset features listinputfeatures padtokenlabelid int 100 use cross entropy ignoreindex as padding label id so that only real label ids contribute to the loss later def init self tokenclassificationtask tokenclassificationtask datadir str tokenizer pretrainedtokenizer labels liststr modeltype str maxseqlength optionalint none overwritecachefalse mode split split train examples tokenclassificationtask readexamplesfromfiledatadir mode todo clean up all this to leverage builtin features of tokenizers self features tokenclassificationtask convertexamplestofeatures examples labels maxseqlength tokenizer clstokenatendboolmodeltype in xlnet xlnet has a cls token at the end clstokentokenizer clstoken clstokensegmentid2 if modeltype in xlnet else 0 septokentokenizer septoken septokenextrafalse roberta uses an extra separator bw pairs of sentences cf github compytorchfairseqcommit1684e166e3da03f5b600dbb7855cb98ddfcd0805 padonleftbooltokenizer paddingside left padtokentokenizer padtokenid padtokensegmentidtokenizer padtokentypeid padtokenlabelidself padtokenlabelid def gen for ex in self features if ex tokentypeids is none yield inputids ex inputids attentionmask ex attentionmask ex labelids else yield inputids ex inputids attentionmask ex attentionmask tokentypeids ex tokentypeids ex labelids if tokentypeids not in tokenizer modelinputnames self dataset tf data dataset fromgenerator gen inputids tf int32 attentionmask tf int32 tf int64 inputids tf tensorshapenone attentionmask tf tensorshapenone tf tensorshapenone else self dataset tf data dataset fromgenerator gen inputids tf int32 attentionmask tf int32 tokentypeids tf int32 tf int64 inputids tf tensorshapenone attentionmask tf tensorshapenone tokentypeids tf tensorshapenone tf tensorshapenone def getdatasetself self dataset self dataset applytf data experimental assertcardinalitylenself features return self dataset def lenself return lenself features def getitemself i inputfeatures return self featuresi coding utf 8 2018 the google ai language team s and the huggingface inc team c 2018 nvidia corporation licensed under the apache license version 2 0 the license you may not use this file except in compliance with the license you may obtain a copy of the license at http www apache org licenses license 2 0 unless required by applicable law or agreed to in writing software distributed under the license is distributed on an as is basis without warranties or conditions of any kind either express or implied see the license for the specific language governing permissions and limitations under the license named entity recognition fine tuning utilities to work with conll 2003 task a single training test example for token classification args guid unique id for the example words list the words of the sequence labels optional list the labels for each word of the sequence this should be specified for train and dev examples but not for test examples a single set of features of data property names are the same names as the corresponding inputs to a model loads a data file into a list of inputfeatures cls_token_at_end define the location of the cls token false default bert xlm pattern cls a sep b sep true xlnet gpt pattern a sep b sep cls cls_token_segment_id define the segment id associated to the cls token 0 for bert 2 for xlnet todo clean up all this to leverage built in features of tokenizers bert base multilingual cased sometimes output nothing when calling tokenize with just a space use the real label id for the first token of the word and padding ids for the remaining tokens account for cls and sep with 2 and with 3 for roberta the convention in bert is a for sequence pairs tokens cls is this jack son ville sep no it is not sep type_ids 0 0 0 0 0 0 0 0 1 1 1 1 1 1 b for single sequences tokens cls the dog is hairy sep type_ids 0 0 0 0 0 0 0 where type_ids are used to indicate whether this is the first sequence or the second sequence the embedding vectors for type 0 and type 1 were learned during pre training and are added to the wordpiece embedding vector and position vector this is not strictly necessary since the sep token unambiguously separates the sequences but it makes it easier for the model to learn the concept of sequences for classification tasks the first vector corresponding to cls is used as the sentence vector note that this only makes sense because the entire model is fine tuned roberta uses an extra separator b w pairs of sentences the mask has 1 for real tokens and 0 for padding tokens only real tokens are attended to zero pad up to the sequence length this will be superseded by a framework agnostic approach soon use cross entropy ignore_index as padding label id so that only real label ids contribute to the loss later load data features from cache or dataset file make sure only the first process in distributed training processes the dataset and the others will use the cache todo clean up all this to leverage built in features of tokenizers xlnet has a cls token at the end roberta uses an extra separator b w pairs of sentences cf github com pytorch fairseq commit 1684e166e3da03f5b600dbb7855cb98ddfcd0805 this will be superseded by a framework agnostic approach soon use cross entropy ignore_index as padding label id so that only real label ids contribute to the loss later todo clean up all this to leverage built in features of tokenizers xlnet has a cls token at the end roberta uses an extra separator b w pairs of sentences cf github com pytorch fairseq commit 1684e166e3da03f5b600dbb7855cb98ddfcd0805
import logging
import os
from dataclasses import dataclass
from enum import Enum
from typing import List, Optional, Union

from filelock import FileLock

from transformers import PreTrainedTokenizer, is_tf_available, is_torch_available


logger = logging.getLogger(__name__)


@dataclass
class InputExample:
    

    guid: str
    words: List[str]
    labels: Optional[List[str]]


@dataclass
class InputFeatures:
    

    input_ids: List[int]
    attention_mask: List[int]
    token_type_ids: Optional[List[int]] = None
    label_ids: Optional[List[int]] = None


class Split(Enum):
    train = "train"
    dev = "dev"
    test = "test"


class TokenClassificationTask:
    @staticmethod
    def read_examples_from_file(data_dir, mode: Union[Split, str]) -> List[InputExample]:
        raise NotImplementedError

    @staticmethod
    def get_labels(path: str) -> List[str]:
        raise NotImplementedError

    @staticmethod
    def convert_examples_to_features(
        examples: List[InputExample],
        label_list: List[str],
        max_seq_length: int,
        tokenizer: PreTrainedTokenizer,
        cls_token_at_end=False,
        cls_token="[CLS]",
        cls_token_segment_id=1,
        sep_token="[SEP]",
        sep_token_extra=False,
        pad_on_left=False,
        pad_token=0,
        pad_token_segment_id=0,
        pad_token_label_id=-100,
        sequence_a_segment_id=0,
        mask_padding_with_zero=True,
    ) -> List[InputFeatures]:
        
        

        label_map = {label: i for i, label in enumerate(label_list)}

        features = []
        for ex_index, example in enumerate(examples):
            if ex_index % 10_000 == 0:
                logger.info("Writing example %d of %d", ex_index, len(examples))

            tokens = []
            label_ids = []
            for word, label in zip(example.words, example.labels):
                word_tokens = tokenizer.tokenize(word)

                
                if len(word_tokens) > 0:
                    tokens.extend(word_tokens)
                    
                    label_ids.extend([label_map[label]] + [pad_token_label_id] * (len(word_tokens) - 1))

            
            special_tokens_count = tokenizer.num_special_tokens_to_add()
            if len(tokens) > max_seq_length - special_tokens_count:
                tokens = tokens[: (max_seq_length - special_tokens_count)]
                label_ids = label_ids[: (max_seq_length - special_tokens_count)]

            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            
            tokens += [sep_token]
            label_ids += [pad_token_label_id]
            if sep_token_extra:
                
                tokens += [sep_token]
                label_ids += [pad_token_label_id]
            segment_ids = [sequence_a_segment_id] * len(tokens)

            if cls_token_at_end:
                tokens += [cls_token]
                label_ids += [pad_token_label_id]
                segment_ids += [cls_token_segment_id]
            else:
                tokens = [cls_token] + tokens
                label_ids = [pad_token_label_id] + label_ids
                segment_ids = [cls_token_segment_id] + segment_ids

            input_ids = tokenizer.convert_tokens_to_ids(tokens)

            
            
            input_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)

            
            padding_length = max_seq_length - len(input_ids)
            if pad_on_left:
                input_ids = ([pad_token] * padding_length) + input_ids
                input_mask = ([0 if mask_padding_with_zero else 1] * padding_length) + input_mask
                segment_ids = ([pad_token_segment_id] * padding_length) + segment_ids
                label_ids = ([pad_token_label_id] * padding_length) + label_ids
            else:
                input_ids += [pad_token] * padding_length
                input_mask += [0 if mask_padding_with_zero else 1] * padding_length
                segment_ids += [pad_token_segment_id] * padding_length
                label_ids += [pad_token_label_id] * padding_length

            assert len(input_ids) == max_seq_length
            assert len(input_mask) == max_seq_length
            assert len(segment_ids) == max_seq_length
            assert len(label_ids) == max_seq_length

            if ex_index < 5:
                logger.info("*** Example ***")
                logger.info("guid: %s", example.guid)
                logger.info("tokens: %s", " ".join([str(x) for x in tokens]))
                logger.info("input_ids: %s", " ".join([str(x) for x in input_ids]))
                logger.info("input_mask: %s", " ".join([str(x) for x in input_mask]))
                logger.info("segment_ids: %s", " ".join([str(x) for x in segment_ids]))
                logger.info("label_ids: %s", " ".join([str(x) for x in label_ids]))

            if "token_type_ids" not in tokenizer.model_input_names:
                segment_ids = None

            features.append(
                InputFeatures(
                    input_ids=input_ids, attention_mask=input_mask, token_type_ids=segment_ids, label_ids=label_ids
                )
            )
        return features


if is_torch_available():
    import torch
    from torch import nn
    from torch.utils.data import Dataset

    class TokenClassificationDataset(Dataset):
        

        features: List[InputFeatures]
        pad_token_label_id: int = nn.CrossEntropyLoss().ignore_index
        
        

        def __init__(
            self,
            token_classification_task: TokenClassificationTask,
            data_dir: str,
            tokenizer: PreTrainedTokenizer,
            labels: List[str],
            model_type: str,
            max_seq_length: Optional[int] = None,
            overwrite_cache=False,
            mode: Split = Split.train,
        ):
            
            cached_features_file = os.path.join(
                data_dir,
                "cached_{}_{}_{}".format(mode.value, tokenizer.__class__.__name__, str(max_seq_length)),
            )

            
            
            lock_path = cached_features_file + ".lock"
            with FileLock(lock_path):
                if os.path.exists(cached_features_file) and not overwrite_cache:
                    logger.info(f"Loading features from cached file {cached_features_file}")
                    self.features = torch.load(cached_features_file)
                else:
                    logger.info(f"Creating features from dataset file at {data_dir}")
                    examples = token_classification_task.read_examples_from_file(data_dir, mode)
                    
                    self.features = token_classification_task.convert_examples_to_features(
                        examples,
                        labels,
                        max_seq_length,
                        tokenizer,
                        cls_token_at_end=bool(model_type in ["xlnet"]),
                        
                        cls_token=tokenizer.cls_token,
                        cls_token_segment_id=2 if model_type in ["xlnet"] else 0,
                        sep_token=tokenizer.sep_token,
                        sep_token_extra=False,
                        
                        pad_on_left=bool(tokenizer.padding_side == "left"),
                        pad_token=tokenizer.pad_token_id,
                        pad_token_segment_id=tokenizer.pad_token_type_id,
                        pad_token_label_id=self.pad_token_label_id,
                    )
                    logger.info(f"Saving features into cached file {cached_features_file}")
                    torch.save(self.features, cached_features_file)

        def __len__(self):
            return len(self.features)

        def __getitem__(self, i) -> InputFeatures:
            return self.features[i]


if is_tf_available():
    import tensorflow as tf

    class TFTokenClassificationDataset:
        

        features: List[InputFeatures]
        pad_token_label_id: int = -100
        
        

        def __init__(
            self,
            token_classification_task: TokenClassificationTask,
            data_dir: str,
            tokenizer: PreTrainedTokenizer,
            labels: List[str],
            model_type: str,
            max_seq_length: Optional[int] = None,
            overwrite_cache=False,
            mode: Split = Split.train,
        ):
            examples = token_classification_task.read_examples_from_file(data_dir, mode)
            
            self.features = token_classification_task.convert_examples_to_features(
                examples,
                labels,
                max_seq_length,
                tokenizer,
                cls_token_at_end=bool(model_type in ["xlnet"]),
                
                cls_token=tokenizer.cls_token,
                cls_token_segment_id=2 if model_type in ["xlnet"] else 0,
                sep_token=tokenizer.sep_token,
                sep_token_extra=False,
                
                pad_on_left=bool(tokenizer.padding_side == "left"),
                pad_token=tokenizer.pad_token_id,
                pad_token_segment_id=tokenizer.pad_token_type_id,
                pad_token_label_id=self.pad_token_label_id,
            )

            def gen():
                for ex in self.features:
                    if ex.token_type_ids is None:
                        yield (
                            {"input_ids": ex.input_ids, "attention_mask": ex.attention_mask},
                            ex.label_ids,
                        )
                    else:
                        yield (
                            {
                                "input_ids": ex.input_ids,
                                "attention_mask": ex.attention_mask,
                                "token_type_ids": ex.token_type_ids,
                            },
                            ex.label_ids,
                        )

            if "token_type_ids" not in tokenizer.model_input_names:
                self.dataset = tf.data.Dataset.from_generator(
                    gen,
                    ({"input_ids": tf.int32, "attention_mask": tf.int32}, tf.int64),
                    (
                        {"input_ids": tf.TensorShape([None]), "attention_mask": tf.TensorShape([None])},
                        tf.TensorShape([None]),
                    ),
                )
            else:
                self.dataset = tf.data.Dataset.from_generator(
                    gen,
                    ({"input_ids": tf.int32, "attention_mask": tf.int32, "token_type_ids": tf.int32}, tf.int64),
                    (
                        {
                            "input_ids": tf.TensorShape([None]),
                            "attention_mask": tf.TensorShape([None]),
                            "token_type_ids": tf.TensorShape([None]),
                        },
                        tf.TensorShape([None]),
                    ),
                )

        def get_dataset(self):
            self.dataset = self.dataset.apply(tf.data.experimental.assert_cardinality(len(self.features)))

            return self.dataset

        def __len__(self):
            return len(self.features)

        def __getitem__(self, i) -> InputFeatures:
            return self.features[i]