mirror of https://github.com/hpcaitech/ColossalAI
257 lines
8.2 KiB
Python
257 lines
8.2 KiB
Python
|
# coding=utf-8
|
||
|
|
# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
|
||
|
|
#
|
||
|
|
# 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.
|
||
|
|
|
||
|
|
"""Megatron tokenizers."""
|
||
|
|
|
||
|
|
from abc import ABC
|
||
|
|
from abc import abstractmethod
|
||
|
|
from colossalai.core import global_context as gpc
|
||
|
|
from colossalai.context import ParallelMode
|
||
|
|
|
||
|
|
from .bert_tokenization import FullTokenizer as FullBertTokenizer
|
||
|
|
|
||
|
|
|
||
|
|
def build_tokenizer(vocab_file, tokenizer_type, vocab_extra_ids=0):
|
||
|
|
"""Initialize tokenizer."""
|
||
|
|
if not gpc.is_initialized(ParallelMode.GLOBAL) or gpc.get_global_rank() == 0:
|
||
|
|
print('> building {} tokenizer ...'.format(tokenizer_type),
|
||
|
|
flush=True)
|
||
|
|
|
||
|
|
# Select and instantiate the tokenizer.
|
||
|
|
if tokenizer_type == 'BertWordPieceLowerCase':
|
||
|
|
tokenizer = _BertWordPieceTokenizer(vocab_file=vocab_file,
|
||
|
|
lower_case=True,
|
||
|
|
vocab_extra_ids=vocab_extra_ids)
|
||
|
|
elif tokenizer_type == 'BertWordPieceCase':
|
||
|
|
tokenizer = _BertWordPieceTokenizer(vocab_file=vocab_file,
|
||
|
|
lower_case=False,
|
||
|
|
vocab_extra_ids=vocab_extra_ids)
|
||
|
|
else:
|
||
|
|
raise NotImplementedError('{} tokenizer is not '
|
||
|
|
'implemented.'.format(tokenizer_type))
|
||
|
|
|
||
|
|
# Add vocab size.
|
||
|
|
padded_vocab_size = _vocab_size_with_padding(tokenizer.vocab_size)
|
||
|
|
|
||
|
|
return tokenizer, padded_vocab_size
|
||
|
|
|
||
|
|
|
||
|
|
def _vocab_size_with_padding(orig_vocab_size, make_vocab_size_divisible_by=128):
|
||
|
|
"""Pad vocab size so it is divisible by model parallel size and
|
||
|
|
still having GPU friendly size."""
|
||
|
|
|
||
|
|
after = orig_vocab_size
|
||
|
|
|
||
|
|
if gpc.is_initialized(ParallelMode.TENSOR):
|
||
|
|
multiple = make_vocab_size_divisible_by * gpc.get_world_size(ParallelMode.TENSOR)
|
||
|
|
else:
|
||
|
|
multiple = make_vocab_size_divisible_by
|
||
|
|
while (after % multiple) != 0:
|
||
|
|
after += 1
|
||
|
|
if not gpc.is_initialized(ParallelMode.GLOBAL) or gpc.get_global_rank() == 0:
|
||
|
|
print(' > padded vocab (size: {}) with {} dummy tokens '
|
||
|
|
'(new size: {})'.format(
|
||
|
|
orig_vocab_size, after - orig_vocab_size, after), flush=True)
|
||
|
|
return after
|
||
|
|
|
||
|
|
|
||
|
|
class AbstractTokenizer(ABC):
|
||
|
|
"""Abstract class for tokenizer."""
|
||
|
|
|
||
|
|
def __init__(self, name):
|
||
|
|
self.name = name
|
||
|
|
super().__init__()
|
||
|
|
|
||
|
|
@property
|
||
|
|
@abstractmethod
|
||
|
|
def vocab_size(self):
|
||
|
|
pass
|
||
|
|
|
||
|
|
@property
|
||
|
|
@abstractmethod
|
||
|
|
def vocab(self):
|
||
|
|
"""Dictionary from vocab text token to id token."""
|
||
|
|
pass
|
||
|
|
|
||
|
|
@property
|
||
|
|
@abstractmethod
|
||
|
|
def inv_vocab(self):
|
||
|
|
"""Dictionary from vocab id token to text token."""
|
||
|
|
pass
|
||
|
|
|
||
|
|
@abstractmethod
|
||
|
|
def tokenize(self, text):
|
||
|
|
pass
|
||
|
|
|
||
|
|
def detokenize(self, token_ids):
|
||
|
|
raise NotImplementedError('detokenizer is not implemented for {} '
|
||
|
|
'tokenizer'.format(self.name))
|
||
|
|
|
||
|
|
@property
|
||
|
|
def cls(self):
|
||
|
|
raise NotImplementedError('CLS is not provided for {} '
|
||
|
|
'tokenizer'.format(self.name))
|
||
|
|
|
||
|
|
@property
|
||
|
|
def sep(self):
|
||
|
|
raise NotImplementedError('SEP is not provided for {} '
|
||
|
|
'tokenizer'.format(self.name))
|
||
|
|
|
||
|
|
@property
|
||
|
|
def pad(self):
|
||
|
|
raise NotImplementedError('PAD is not provided for {} '
|
||
|
|
'tokenizer'.format(self.name))
|
||
|
|
|
||
|
|
@property
|
||
|
|
def eod(self):
|
||
|
|
raise NotImplementedError('EOD is not provided for {} '
|
||
|
|
'tokenizer'.format(self.name))
|
||
|
|
|
||
|
|
@property
|
||
|
|
def mask(self):
|
||
|
|
raise NotImplementedError('MASK is not provided for {} '
|
||
|
|
'tokenizer'.format(self.name))
|
||
|
|
|
||
|
|
|
||
|
|
class _BertWordPieceTokenizer(AbstractTokenizer):
|
||
|
|
"""Original BERT wordpiece tokenizer."""
|
||
|
|
|
||
|
|
def __init__(self, vocab_file, lower_case=True, vocab_extra_ids=0):
|
||
|
|
if lower_case:
|
||
|
|
name = 'BERT Lower Case'
|
||
|
|
else:
|
||
|
|
name = 'BERT Upper Case'
|
||
|
|
super().__init__(name)
|
||
|
|
self.tokenizer = FullBertTokenizer(vocab_file, do_lower_case=lower_case)
|
||
|
|
self.cls_id = self.tokenizer.vocab['[CLS]']
|
||
|
|
self.sep_id = self.tokenizer.vocab['[SEP]']
|
||
|
|
self.pad_id = self.tokenizer.vocab['[PAD]']
|
||
|
|
self.mask_id = self.tokenizer.vocab['[MASK]']
|
||
|
|
self._additional_special_tokens = []
|
||
|
|
|
||
|
|
# (dsachan) Add BOS and EOS tokens
|
||
|
|
SPECIAL_TOKENS = {'eos_token': '[EOS]',
|
||
|
|
'bos_token': '[BOS]'}
|
||
|
|
self._bos_token = '[BOS]'
|
||
|
|
self.add_token(self._bos_token)
|
||
|
|
self._bos_token_id = self.vocab.get(self._bos_token)
|
||
|
|
|
||
|
|
self._eos_token = '[EOS]'
|
||
|
|
self.add_token(self._eos_token)
|
||
|
|
self._eos_token_id = self.vocab.get(self._eos_token)
|
||
|
|
|
||
|
|
# (dsachan) Add additional special tokens
|
||
|
|
# These can be used as sentinel tokens in T5 model inputs
|
||
|
|
additional_special_tokens = []
|
||
|
|
additional_special_tokens.extend(
|
||
|
|
["<extra_id_{}>".format(i) for i in range(vocab_extra_ids)])
|
||
|
|
self.add_additional_special_tokens(additional_special_tokens)
|
||
|
|
|
||
|
|
def add_token(self, token):
|
||
|
|
if token not in self.vocab:
|
||
|
|
self.inv_vocab[self.vocab_size] = token
|
||
|
|
# self.vocab_size comes from len(vocab)
|
||
|
|
# and it will increase as we add elements
|
||
|
|
self.vocab[token] = self.vocab_size
|
||
|
|
|
||
|
|
def add_additional_special_tokens(self, tokens_list):
|
||
|
|
setattr(self, "additional_special_tokens", tokens_list)
|
||
|
|
for value in tokens_list:
|
||
|
|
self.add_token(value)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def vocab_size(self):
|
||
|
|
return self.tokenizer.vocab_size()
|
||
|
|
|
||
|
|
@property
|
||
|
|
def vocab(self):
|
||
|
|
return self.tokenizer.vocab
|
||
|
|
|
||
|
|
@property
|
||
|
|
def inv_vocab(self):
|
||
|
|
return self.tokenizer.inv_vocab
|
||
|
|
|
||
|
|
def tokenize(self, text):
|
||
|
|
text_tokens = self.tokenizer.tokenize(text)
|
||
|
|
return self.tokenizer.convert_tokens_to_ids(text_tokens)
|
||
|
|
|
||
|
|
def decode(self, ids):
|
||
|
|
tokens = self.tokenizer.convert_ids_to_tokens(ids)
|
||
|
|
return self.tokenizer.convert_tokens_to_string(tokens)
|
||
|
|
|
||
|
|
def decode_token_ids(self, token_ids):
|
||
|
|
tokens = self.tokenizer.convert_ids_to_tokens(token_ids)
|
||
|
|
exclude_list = ['[PAD]', '[CLS]']
|
||
|
|
non_pads = [t for t in tokens if t not in exclude_list]
|
||
|
|
|
||
|
|
result = ""
|
||
|
|
for s in non_pads:
|
||
|
|
if s.startswith("##"):
|
||
|
|
result += s[2:]
|
||
|
|
else:
|
||
|
|
result += " " + s
|
||
|
|
|
||
|
|
return result
|
||
|
|
|
||
|
|
@property
|
||
|
|
def cls(self):
|
||
|
|
return self.cls_id
|
||
|
|
|
||
|
|
@property
|
||
|
|
def sep(self):
|
||
|
|
return self.sep_id
|
||
|
|
|
||
|
|
@property
|
||
|
|
def pad(self):
|
||
|
|
return self.pad_id
|
||
|
|
|
||
|
|
@property
|
||
|
|
def mask(self):
|
||
|
|
return self.mask_id
|
||
|
|
|
||
|
|
@property
|
||
|
|
def bos_token(self):
|
||
|
|
""" Beginning of sentence token id """
|
||
|
|
return self._bos_token
|
||
|
|
|
||
|
|
@property
|
||
|
|
def eos_token(self):
|
||
|
|
""" End of sentence token id """
|
||
|
|
return self._eos_token
|
||
|
|
|
||
|
|
@property
|
||
|
|
def additional_special_tokens(self):
|
||
|
|
""" All the additional special tokens you may want to use (list of strings)."""
|
||
|
|
return self._additional_special_tokens
|
||
|
|
|
||
|
|
@property
|
||
|
|
def bos_token_id(self):
|
||
|
|
""" Id of the beginning of sentence token in the vocabulary."""
|
||
|
|
return self._bos_token_id
|
||
|
|
|
||
|
|
@property
|
||
|
|
def eos_token_id(self):
|
||
|
|
""" Id of the end of sentence token in the vocabulary."""
|
||
|
|
return self._eos_token_id
|
||
|
|
|
||
|
|
@property
|
||
|
|
def additional_special_tokens_ids(self):
|
||
|
|
""" Ids of all the additional special tokens in the vocabulary (list of integers)."""
|
||
|
|
return [self.vocab.get(token) for token in self._additional_special_tokens]
|
||
|
|
|
||
|
|
@additional_special_tokens.setter
|
||
|
|
def additional_special_tokens(self, value):
|
||
|
|
self._additional_special_tokens = value
|