| from typing import Dict, List, Optional, Tuple, Union
|
|
|
| from tokenizers import AddedToken, EncodeInput, Encoding, InputSequence, Tokenizer
|
| from tokenizers.decoders import Decoder
|
| from tokenizers.models import Model
|
| from tokenizers.normalizers import Normalizer
|
| from tokenizers.pre_tokenizers import PreTokenizer
|
| from tokenizers.processors import PostProcessor
|
|
|
|
|
| Offsets = Tuple[int, int]
|
|
|
|
|
| class BaseTokenizer:
|
| def __init__(self, tokenizer: Tokenizer, parameters=None):
|
| self._tokenizer = tokenizer
|
| self._parameters = parameters if parameters is not None else {}
|
|
|
| def __repr__(self):
|
| return "Tokenizer(vocabulary_size={}, {})".format(
|
| self._tokenizer.get_vocab_size(),
|
| ", ".join(k + "=" + str(v) for k, v in self._parameters.items()),
|
| )
|
|
|
| def num_special_tokens_to_add(self, is_pair: bool) -> int:
|
| """
|
| Return the number of special tokens that would be added for single/pair sentences.
|
| :param is_pair: Boolean indicating if the input would be a single sentence or a pair
|
| :return:
|
| """
|
| return self._tokenizer.num_special_tokens_to_add(is_pair)
|
|
|
| def get_vocab(self, with_added_tokens: bool = True) -> Dict[str, int]:
|
| """Returns the vocabulary
|
|
|
| Args:
|
| with_added_tokens: boolean:
|
| Whether to include the added tokens in the vocabulary
|
|
|
| Returns:
|
| The vocabulary
|
| """
|
| return self._tokenizer.get_vocab(with_added_tokens=with_added_tokens)
|
|
|
| def get_added_tokens_decoder(self) -> Dict[int, AddedToken]:
|
| """Returns the added reverse vocabulary
|
|
|
| Returns:
|
| The added vocabulary mapping ints to AddedTokens
|
| """
|
| return self._tokenizer.get_added_tokens_decoder()
|
|
|
| def get_vocab_size(self, with_added_tokens: bool = True) -> int:
|
| """Return the size of vocabulary, with or without added tokens.
|
|
|
| Args:
|
| with_added_tokens: (`optional`) bool:
|
| Whether to count in added special tokens or not
|
|
|
| Returns:
|
| Size of vocabulary
|
| """
|
| return self._tokenizer.get_vocab_size(with_added_tokens=with_added_tokens)
|
|
|
| def enable_padding(
|
| self,
|
| direction: Optional[str] = "right",
|
| pad_to_multiple_of: Optional[int] = None,
|
| pad_id: Optional[int] = 0,
|
| pad_type_id: Optional[int] = 0,
|
| pad_token: Optional[str] = "[PAD]",
|
| length: Optional[int] = None,
|
| ):
|
| """Change the padding strategy
|
|
|
| Args:
|
| direction: (`optional`) str:
|
| Can be one of: `right` or `left`
|
|
|
| pad_to_multiple_of: (`optional`) unsigned int:
|
| If specified, the padding length should always snap to the next multiple of
|
| the given value. For example if we were going to pad with a length of 250 but
|
| `pad_to_multiple_of=8` then we will pad to 256.
|
|
|
| pad_id: (`optional`) unsigned int:
|
| The indice to be used when padding
|
|
|
| pad_type_id: (`optional`) unsigned int:
|
| The type indice to be used when padding
|
|
|
| pad_token: (`optional`) str:
|
| The pad token to be used when padding
|
|
|
| length: (`optional`) unsigned int:
|
| If specified, the length at which to pad. If not specified
|
| we pad using the size of the longest sequence in a batch
|
| """
|
| return self._tokenizer.enable_padding(
|
| direction=direction,
|
| pad_to_multiple_of=pad_to_multiple_of,
|
| pad_id=pad_id,
|
| pad_type_id=pad_type_id,
|
| pad_token=pad_token,
|
| length=length,
|
| )
|
|
|
| def no_padding(self):
|
| """Disable padding"""
|
| return self._tokenizer.no_padding()
|
|
|
| @property
|
| def padding(self) -> Optional[dict]:
|
| """Get the current padding parameters
|
|
|
| Returns:
|
| None if padding is disabled, a dict with the currently set parameters
|
| if the padding is enabled.
|
| """
|
| return self._tokenizer.padding
|
|
|
| def enable_truncation(self, max_length: int, stride: Optional[int] = 0, strategy: Optional[str] = "longest_first"):
|
| """Change the truncation options
|
|
|
| Args:
|
| max_length: unsigned int:
|
| The maximum length at which to truncate
|
|
|
| stride: (`optional`) unsigned int:
|
| The length of the previous first sequence to be included
|
| in the overflowing sequence
|
|
|
| strategy: (`optional`) str:
|
| Can be one of `longest_first`, `only_first` or `only_second`
|
| """
|
| return self._tokenizer.enable_truncation(max_length, stride=stride, strategy=strategy)
|
|
|
| def no_truncation(self):
|
| """Disable truncation"""
|
| return self._tokenizer.no_truncation()
|
|
|
| @property
|
| def truncation(self) -> Optional[dict]:
|
| """Get the current truncation parameters
|
|
|
| Returns:
|
| None if truncation is disabled, a dict with the current truncation parameters if
|
| truncation is enabled
|
| """
|
| return self._tokenizer.truncation
|
|
|
| def add_tokens(self, tokens: List[Union[str, AddedToken]]) -> int:
|
| """Add the given tokens to the vocabulary
|
|
|
| Args:
|
| tokens: List[Union[str, AddedToken]]:
|
| A list of tokens to add to the vocabulary. Each token can either be
|
| a string, or an instance of AddedToken
|
|
|
| Returns:
|
| The number of tokens that were added to the vocabulary
|
| """
|
| return self._tokenizer.add_tokens(tokens)
|
|
|
| def add_special_tokens(self, special_tokens: List[Union[str, AddedToken]]) -> int:
|
| """Add the given special tokens to the vocabulary, and treat them as special tokens.
|
|
|
| The special tokens will never be processed by the model, and will be
|
| removed while decoding.
|
|
|
| Args:
|
| tokens: List[Union[str, AddedToken]]:
|
| A list of special tokens to add to the vocabulary. Each token can either be
|
| a string, or an instance of AddedToken
|
|
|
| Returns:
|
| The number of tokens that were added to the vocabulary
|
| """
|
| return self._tokenizer.add_special_tokens(special_tokens)
|
|
|
| def normalize(self, sequence: str) -> str:
|
| """Normalize the given sequence
|
|
|
| Args:
|
| sequence: str:
|
| The sequence to normalize
|
|
|
| Returns:
|
| The normalized string
|
| """
|
| return self._tokenizer.normalize(sequence)
|
|
|
| def encode(
|
| self,
|
| sequence: InputSequence,
|
| pair: Optional[InputSequence] = None,
|
| is_pretokenized: bool = False,
|
| add_special_tokens: bool = True,
|
| ) -> Encoding:
|
| """Encode the given sequence and pair. This method can process raw text sequences as well
|
| as already pre-tokenized sequences.
|
|
|
| Args:
|
| sequence: InputSequence:
|
| The sequence we want to encode. This sequence can be either raw text or
|
| pre-tokenized, according to the `is_pretokenized` argument:
|
|
|
| - If `is_pretokenized=False`: `InputSequence` is expected to be `str`
|
| - If `is_pretokenized=True`: `InputSequence` is expected to be
|
| `Union[List[str], Tuple[str]]`
|
|
|
| is_pretokenized: bool:
|
| Whether the input is already pre-tokenized.
|
|
|
| add_special_tokens: bool:
|
| Whether to add the special tokens while encoding.
|
|
|
| Returns:
|
| An Encoding
|
| """
|
| if sequence is None:
|
| raise ValueError("encode: `sequence` can't be `None`")
|
|
|
| return self._tokenizer.encode(sequence, pair, is_pretokenized, add_special_tokens)
|
|
|
| def encode_batch(
|
| self,
|
| inputs: List[EncodeInput],
|
| is_pretokenized: bool = False,
|
| add_special_tokens: bool = True,
|
| ) -> List[Encoding]:
|
| """Encode the given inputs. This method accept both raw text sequences as well as already
|
| pre-tokenized sequences.
|
|
|
| Args:
|
| inputs: List[EncodeInput]:
|
| A list of single sequences or pair sequences to encode. Each `EncodeInput` is
|
| expected to be of the following form:
|
| `Union[InputSequence, Tuple[InputSequence, InputSequence]]`
|
|
|
| Each `InputSequence` can either be raw text or pre-tokenized,
|
| according to the `is_pretokenized` argument:
|
|
|
| - If `is_pretokenized=False`: `InputSequence` is expected to be `str`
|
| - If `is_pretokenized=True`: `InputSequence` is expected to be
|
| `Union[List[str], Tuple[str]]`
|
|
|
| is_pretokenized: bool:
|
| Whether the input is already pre-tokenized.
|
|
|
| add_special_tokens: bool:
|
| Whether to add the special tokens while encoding.
|
|
|
| Returns:
|
| A list of Encoding
|
| """
|
|
|
| if inputs is None:
|
| raise ValueError("encode_batch: `inputs` can't be `None`")
|
|
|
| return self._tokenizer.encode_batch(inputs, is_pretokenized, add_special_tokens)
|
|
|
| def decode(self, ids: List[int], skip_special_tokens: Optional[bool] = True) -> str:
|
| """Decode the given list of ids to a string sequence
|
|
|
| Args:
|
| ids: List[unsigned int]:
|
| A list of ids to be decoded
|
|
|
| skip_special_tokens: (`optional`) boolean:
|
| Whether to remove all the special tokens from the output string
|
|
|
| Returns:
|
| The decoded string
|
| """
|
| if ids is None:
|
| raise ValueError("None input is not valid. Should be a list of integers.")
|
|
|
| return self._tokenizer.decode(ids, skip_special_tokens=skip_special_tokens)
|
|
|
| def decode_batch(self, sequences: List[List[int]], skip_special_tokens: Optional[bool] = True) -> str:
|
| """Decode the list of sequences to a list of string sequences
|
|
|
| Args:
|
| sequences: List[List[unsigned int]]:
|
| A list of sequence of ids to be decoded
|
|
|
| skip_special_tokens: (`optional`) boolean:
|
| Whether to remove all the special tokens from the output strings
|
|
|
| Returns:
|
| A list of decoded strings
|
| """
|
| if sequences is None:
|
| raise ValueError("None input is not valid. Should be list of list of integers.")
|
|
|
| return self._tokenizer.decode_batch(sequences, skip_special_tokens=skip_special_tokens)
|
|
|
| def token_to_id(self, token: str) -> Optional[int]:
|
| """Convert the given token to its corresponding id
|
|
|
| Args:
|
| token: str:
|
| The token to convert
|
|
|
| Returns:
|
| The corresponding id if it exists, None otherwise
|
| """
|
| return self._tokenizer.token_to_id(token)
|
|
|
| def id_to_token(self, id: int) -> Optional[str]:
|
| """Convert the given token id to its corresponding string
|
|
|
| Args:
|
| token: id:
|
| The token id to convert
|
|
|
| Returns:
|
| The corresponding string if it exists, None otherwise
|
| """
|
| return self._tokenizer.id_to_token(id)
|
|
|
| def save_model(self, directory: str, prefix: Optional[str] = None):
|
| """Save the current model to the given directory
|
|
|
| Args:
|
| directory: str:
|
| A path to the destination directory
|
|
|
| prefix: (Optional) str:
|
| An optional prefix, used to prefix each file name
|
| """
|
| return self._tokenizer.model.save(directory, prefix=prefix)
|
|
|
| def save(self, path: str, pretty: bool = True):
|
| """Save the current Tokenizer at the given path
|
|
|
| Args:
|
| path: str:
|
| A path to the destination Tokenizer file
|
| """
|
| return self._tokenizer.save(path, pretty)
|
|
|
| def to_str(self, pretty: bool = False):
|
| """Get a serialized JSON version of the Tokenizer as a str
|
|
|
| Args:
|
| pretty: bool:
|
| Whether the JSON string should be prettified
|
|
|
| Returns:
|
| str
|
| """
|
| return self._tokenizer.to_str(pretty)
|
|
|
| def post_process(
|
| self, encoding: Encoding, pair: Optional[Encoding] = None, add_special_tokens: bool = True
|
| ) -> Encoding:
|
| """Apply all the post-processing steps to the given encodings.
|
|
|
| The various steps are:
|
| 1. Truncate according to global params (provided to `enable_truncation`)
|
| 2. Apply the PostProcessor
|
| 3. Pad according to global params. (provided to `enable_padding`)
|
|
|
| Args:
|
| encoding: Encoding:
|
| The main Encoding to post process
|
|
|
| pair: Optional[Encoding]:
|
| An optional pair Encoding
|
|
|
| add_special_tokens: bool:
|
| Whether to add special tokens
|
|
|
| Returns:
|
| The resulting Encoding
|
| """
|
| return self._tokenizer.post_process(encoding, pair, add_special_tokens)
|
|
|
| @property
|
| def model(self) -> Model:
|
| return self._tokenizer.model
|
|
|
| @model.setter
|
| def model(self, model: Model):
|
| self._tokenizer.model = model
|
|
|
| @property
|
| def normalizer(self) -> Normalizer:
|
| return self._tokenizer.normalizer
|
|
|
| @normalizer.setter
|
| def normalizer(self, normalizer: Normalizer):
|
| self._tokenizer.normalizer = normalizer
|
|
|
| @property
|
| def pre_tokenizer(self) -> PreTokenizer:
|
| return self._tokenizer.pre_tokenizer
|
|
|
| @pre_tokenizer.setter
|
| def pre_tokenizer(self, pre_tokenizer: PreTokenizer):
|
| self._tokenizer.pre_tokenizer = pre_tokenizer
|
|
|
| @property
|
| def post_processor(self) -> PostProcessor:
|
| return self._tokenizer.post_processor
|
|
|
| @post_processor.setter
|
| def post_processor(self, post_processor: PostProcessor):
|
| self._tokenizer.post_processor = post_processor
|
|
|
| @property
|
| def decoder(self) -> Decoder:
|
| return self._tokenizer.decoder
|
|
|
| @decoder.setter
|
| def decoder(self, decoder: Decoder):
|
| self._tokenizer.decoder = decoder
|
|
|
