Aluode/PerceptionLabPortable

1import os2from typing import Optional, Union3 4import tensorflow as tf5from tensorflow_text import BertTokenizer as BertTokenizerLayer6from tensorflow_text import FastBertTokenizer, ShrinkLongestTrimmer, case_fold_utf8, combine_segments, pad_model_inputs7 8from ...modeling_tf_utils import keras9from ...utils.import_utils import requires10from .tokenization_bert import BertTokenizer11 12 13@requires(backends=("tf", "tensorflow_text"))14class TFBertTokenizer(keras.layers.Layer):15    """16    This is an in-graph tokenizer for BERT. It should be initialized similarly to other tokenizers, using the17    `from_pretrained()` method. It can also be initialized with the `from_tokenizer()` method, which imports settings18    from an existing standard tokenizer object.19 20    In-graph tokenizers, unlike other Hugging Face tokenizers, are actually Keras layers and are designed to be run21    when the model is called, rather than during preprocessing. As a result, they have somewhat more limited options22    than standard tokenizer classes. They are most useful when you want to create an end-to-end model that goes23    straight from `tf.string` inputs to outputs.24 25    Args:26        vocab_list (`list`):27            List containing the vocabulary.28        do_lower_case (`bool`, *optional*, defaults to `True`):29            Whether or not to lowercase the input when tokenizing.30        cls_token_id (`str`, *optional*, defaults to `"[CLS]"`):31            The classifier token which is used when doing sequence classification (classification of the whole sequence32            instead of per-token classification). It is the first token of the sequence when built with special tokens.33        sep_token_id (`str`, *optional*, defaults to `"[SEP]"`):34            The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for35            sequence classification or for a text and a question for question answering. It is also used as the last36            token of a sequence built with special tokens.37        pad_token_id (`str`, *optional*, defaults to `"[PAD]"`):38            The token used for padding, for example when batching sequences of different lengths.39        padding (`str`, defaults to `"longest"`):40            The type of padding to use. Can be either `"longest"`, to pad only up to the longest sample in the batch,41            or `"max_length", to pad all inputs to the maximum length supported by the tokenizer.42        truncation (`bool`, *optional*, defaults to `True`):43            Whether to truncate the sequence to the maximum length.44        max_length (`int`, *optional*, defaults to `512`):45            The maximum length of the sequence, used for padding (if `padding` is "max_length") and/or truncation (if46            `truncation` is `True`).47        pad_to_multiple_of (`int`, *optional*, defaults to `None`):48            If set, the sequence will be padded to a multiple of this value.49        return_token_type_ids (`bool`, *optional*, defaults to `True`):50            Whether to return token_type_ids.51        return_attention_mask (`bool`, *optional*, defaults to `True`):52            Whether to return the attention_mask.53        use_fast_bert_tokenizer (`bool`, *optional*, defaults to `True`):54            If True, will use the FastBertTokenizer class from Tensorflow Text. If False, will use the BertTokenizer55            class instead. BertTokenizer supports some additional options, but is slower and cannot be exported to56            TFLite.57    """58 59    def __init__(60        self,61        vocab_list: list,62        do_lower_case: bool,63        cls_token_id: Optional[int] = None,64        sep_token_id: Optional[int] = None,65        pad_token_id: Optional[int] = None,66        padding: str = "longest",67        truncation: bool = True,68        max_length: int = 512,69        pad_to_multiple_of: Optional[int] = None,70        return_token_type_ids: bool = True,71        return_attention_mask: bool = True,72        use_fast_bert_tokenizer: bool = True,73        **tokenizer_kwargs,74    ):75        super().__init__()76        if use_fast_bert_tokenizer:77            self.tf_tokenizer = FastBertTokenizer(78                vocab_list, token_out_type=tf.int64, lower_case_nfd_strip_accents=do_lower_case, **tokenizer_kwargs79            )80        else:81            lookup_table = tf.lookup.StaticVocabularyTable(82                tf.lookup.KeyValueTensorInitializer(83                    keys=vocab_list,84                    key_dtype=tf.string,85                    values=tf.range(tf.size(vocab_list, out_type=tf.int64), dtype=tf.int64),86                    value_dtype=tf.int64,87                ),88                num_oov_buckets=1,89            )90            self.tf_tokenizer = BertTokenizerLayer(91                lookup_table, token_out_type=tf.int64, lower_case=do_lower_case, **tokenizer_kwargs92            )93 94        self.vocab_list = vocab_list95        self.do_lower_case = do_lower_case96        self.cls_token_id = vocab_list.index("[CLS]") if cls_token_id is None else cls_token_id97        self.sep_token_id = vocab_list.index("[SEP]") if sep_token_id is None else sep_token_id98        self.pad_token_id = vocab_list.index("[PAD]") if pad_token_id is None else pad_token_id99        self.paired_trimmer = ShrinkLongestTrimmer(max_length - 3, axis=1)  # Allow room for special tokens100        self.max_length = max_length101        self.padding = padding102        self.truncation = truncation103        self.pad_to_multiple_of = pad_to_multiple_of104        self.return_token_type_ids = return_token_type_ids105        self.return_attention_mask = return_attention_mask106 107    @classmethod108    def from_tokenizer(cls, tokenizer: "PreTrainedTokenizerBase", **kwargs):  # noqa: F821109        """110        Initialize a `TFBertTokenizer` from an existing `Tokenizer`.111 112        Args:113            tokenizer (`PreTrainedTokenizerBase`):114                The tokenizer to use to initialize the `TFBertTokenizer`.115 116        Examples:117 118        ```python119        from transformers import AutoTokenizer, TFBertTokenizer120 121        tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")122        tf_tokenizer = TFBertTokenizer.from_tokenizer(tokenizer)123        ```124        """125        do_lower_case = kwargs.pop("do_lower_case", None)126        do_lower_case = tokenizer.do_lower_case if do_lower_case is None else do_lower_case127        cls_token_id = kwargs.pop("cls_token_id", None)128        cls_token_id = tokenizer.cls_token_id if cls_token_id is None else cls_token_id129        sep_token_id = kwargs.pop("sep_token_id", None)130        sep_token_id = tokenizer.sep_token_id if sep_token_id is None else sep_token_id131        pad_token_id = kwargs.pop("pad_token_id", None)132        pad_token_id = tokenizer.pad_token_id if pad_token_id is None else pad_token_id133 134        vocab = tokenizer.get_vocab()135        vocab = sorted(vocab.items(), key=lambda x: x[1])136        vocab_list = [entry[0] for entry in vocab]137        return cls(138            vocab_list=vocab_list,139            do_lower_case=do_lower_case,140            cls_token_id=cls_token_id,141            sep_token_id=sep_token_id,142            pad_token_id=pad_token_id,143            **kwargs,144        )145 146    @classmethod147    def from_pretrained(cls, pretrained_model_name_or_path: Union[str, os.PathLike], *init_inputs, **kwargs):148        """149        Instantiate a `TFBertTokenizer` from a pre-trained tokenizer.150 151        Args:152            pretrained_model_name_or_path (`str` or `os.PathLike`):153                The name or path to the pre-trained tokenizer.154 155        Examples:156 157        ```python158        from transformers import TFBertTokenizer159 160        tf_tokenizer = TFBertTokenizer.from_pretrained("google-bert/bert-base-uncased")161        ```162        """163        try:164            tokenizer = BertTokenizer.from_pretrained(pretrained_model_name_or_path, *init_inputs, **kwargs)165        except:  # noqa: E722166            from .tokenization_bert_fast import BertTokenizerFast167 168            tokenizer = BertTokenizerFast.from_pretrained(pretrained_model_name_or_path, *init_inputs, **kwargs)169        return cls.from_tokenizer(tokenizer, **kwargs)170 171    def unpaired_tokenize(self, texts):172        if self.do_lower_case:173            texts = case_fold_utf8(texts)174        tokens = self.tf_tokenizer.tokenize(texts)175        return tokens.merge_dims(1, -1)176 177    def call(178        self,179        text,180        text_pair=None,181        padding=None,182        truncation=None,183        max_length=None,184        pad_to_multiple_of=None,185        return_token_type_ids=None,186        return_attention_mask=None,187    ):188        if padding is None:189            padding = self.padding190        if padding not in ("longest", "max_length"):191            raise ValueError("Padding must be either 'longest' or 'max_length'!")192        if max_length is not None and text_pair is not None:193            # Because we have to instantiate a Trimmer to do it properly194            raise ValueError("max_length cannot be overridden at call time when truncating paired texts!")195        if max_length is None:196            max_length = self.max_length197        if truncation is None:198            truncation = self.truncation199        if pad_to_multiple_of is None:200            pad_to_multiple_of = self.pad_to_multiple_of201        if return_token_type_ids is None:202            return_token_type_ids = self.return_token_type_ids203        if return_attention_mask is None:204            return_attention_mask = self.return_attention_mask205        if not isinstance(text, tf.Tensor):206            text = tf.convert_to_tensor(text)207        if text_pair is not None and not isinstance(text_pair, tf.Tensor):208            text_pair = tf.convert_to_tensor(text_pair)209        if text_pair is not None:210            if text.shape.rank > 1:211                raise ValueError("text argument should not be multidimensional when a text pair is supplied!")212            if text_pair.shape.rank > 1:213                raise ValueError("text_pair should not be multidimensional!")214        if text.shape.rank == 2:215            text, text_pair = text[:, 0], text[:, 1]216        text = self.unpaired_tokenize(text)217        if text_pair is None:  # Unpaired text218            if truncation:219                text = text[:, : max_length - 2]  # Allow room for special tokens220            input_ids, token_type_ids = combine_segments(221                (text,), start_of_sequence_id=self.cls_token_id, end_of_segment_id=self.sep_token_id222            )223        else:  # Paired text224            text_pair = self.unpaired_tokenize(text_pair)225            if truncation:226                text, text_pair = self.paired_trimmer.trim([text, text_pair])227            input_ids, token_type_ids = combine_segments(228                (text, text_pair), start_of_sequence_id=self.cls_token_id, end_of_segment_id=self.sep_token_id229            )230        if padding == "longest":231            pad_length = input_ids.bounding_shape(axis=1)232            if pad_to_multiple_of is not None:233                # No ceiling division in tensorflow, so we negate floordiv instead234                pad_length = pad_to_multiple_of * (-tf.math.floordiv(-pad_length, pad_to_multiple_of))235        else:236            pad_length = max_length237 238        input_ids, attention_mask = pad_model_inputs(input_ids, max_seq_length=pad_length, pad_value=self.pad_token_id)239        output = {"input_ids": input_ids}240        if return_attention_mask:241            output["attention_mask"] = attention_mask242        if return_token_type_ids:243            token_type_ids, _ = pad_model_inputs(244                token_type_ids, max_seq_length=pad_length, pad_value=self.pad_token_id245            )246            output["token_type_ids"] = token_type_ids247        return output248 249    def get_config(self):250        return {251            "vocab_list": self.vocab_list,252            "do_lower_case": self.do_lower_case,253            "cls_token_id": self.cls_token_id,254            "sep_token_id": self.sep_token_id,255            "pad_token_id": self.pad_token_id,256        }257 258 259__all__ = ["TFBertTokenizer"]260