|
from abc import ABC |
|
from abc import abstractmethod |
|
|
|
|
|
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 _SentencePieceTokenizer(AbstractTokenizer): |
|
"""SentencePieceTokenizer-Megatron wrapper""" |
|
|
|
def __init__(self, model_file, vocab_extra_ids=0): |
|
name = 'SentencePieceTokenizer' |
|
super().__init__(name) |
|
|
|
import sentencepiece |
|
self.tokenizer = sentencepiece.SentencePieceProcessor(model_file=model_file) |
|
self._initalize(vocab_extra_ids) |
|
|
|
def _populate_vocab(self): |
|
self._vocab = {} |
|
self._inv_vocab = {} |
|
|
|
for i in range(len(self.tokenizer)): |
|
t = self.tokenizer.id_to_piece(i) |
|
self._inv_vocab[i] = t |
|
self._vocab[t] = i |
|
|
|
def _initalize(self, vocab_extra_ids): |
|
self._populate_vocab() |
|
self._special_tokens = {} |
|
self._inv_special_tokens = {} |
|
|
|
self._t5_tokens = [] |
|
|
|
def _add_special_token(t): |
|
if t not in self._vocab: |
|
next_id = len(self._vocab) |
|
self._vocab[t] = next_id |
|
self._inv_vocab[next_id] = t |
|
self._special_tokens[t] = self._vocab[t] |
|
self._inv_special_tokens[self._vocab[t]] = t |
|
|
|
_add_special_token('<CLS>') |
|
self._cls_id = self._vocab['<CLS>'] |
|
_add_special_token('<SEP>') |
|
self._sep_id = self._vocab['<SEP>'] |
|
_add_special_token('<EOD>') |
|
self._eod_id = self._vocab['<EOD>'] |
|
_add_special_token('<MASK>') |
|
self._mask_id = self._vocab['<MASK>'] |
|
|
|
pad_id = self.tokenizer.pad_id() |
|
try: |
|
pad_token = self.tokenizer.id_to_piece(pad_id) |
|
except IndexError: |
|
pad_token = '<PAD>' |
|
_add_special_token(pad_token) |
|
self._pad_id = self._vocab[pad_token] |
|
|
|
bos_id = self.tokenizer.bos_id() |
|
try: |
|
bos_token = self.tokenizer.id_to_piece(bos_id) |
|
except IndexError: |
|
bos_token = '<BOS>' |
|
_add_special_token(bos_token) |
|
self._bos_id = self._vocab[bos_token] |
|
|
|
eos_id = self.tokenizer.eos_id() |
|
try: |
|
eos_token = self.tokenizer.id_to_piece(eos_id) |
|
except IndexError: |
|
eos_token = '<EOS>' |
|
_add_special_token(eos_token) |
|
self._eos_id = self._vocab[eos_token] |
|
|
|
for i in range(vocab_extra_ids): |
|
t = "<extra_id_{}>".format(i) |
|
_add_special_token(t) |
|
self._t5_tokens += [t] |
|
|
|
@property |
|
def vocab_size(self): |
|
return len(self._vocab) |
|
|
|
@property |
|
def vocab(self): |
|
return self._vocab |
|
|
|
@property |
|
def inv_vocab(self): |
|
return self._inv_vocab |
|
|
|
@property |
|
def decoder(self): |
|
return self._inv_vocab |
|
|
|
@property |
|
def encoder(self): |
|
return self._vocab |
|
|
|
|
|
|
|
def tokenize(self, text): |
|
ids = [] |
|
idx = 0 |
|
|
|
while 1: |
|
indices = {} |
|
for token in self._special_tokens: |
|
try: |
|
indices[token] = text[idx:].index(token) |
|
except ValueError: |
|
continue |
|
if len(indices) == 0: |
|
break |
|
|
|
next_token = min(indices, key=indices.get) |
|
next_idx = idx + indices[next_token] |
|
|
|
ids.extend(self.tokenizer.encode_as_ids(text[idx:next_idx])) |
|
ids.append(self._special_tokens[next_token]) |
|
idx = next_idx + len(next_token) |
|
|
|
ids.extend(self.tokenizer.encode_as_ids(text[idx:])) |
|
return ids |
|
|
|
|
|
|
|
def detokenize(self, ids): |
|
text = "" |
|
last_i = 0 |
|
|
|
for i, id in enumerate(ids): |
|
if id in self._inv_special_tokens: |
|
text += self.tokenizer.decode_ids(ids[last_i:i]) + " " |
|
text += self._inv_special_tokens[id] + " " |
|
last_i = i + 1 |
|
|
|
text += self.tokenizer.decode_ids(ids[last_i:]) |
|
return text |
|
|
|
@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 bos_token_id(self): |
|
return self._bos_id |
|
|
|
@property |
|
def bos(self): |
|
return self._bos_id |
|
|
|
@property |
|
def eod(self): |
|
return self._eod_id |
|
|
|
@property |
|
def eos_token_id(self): |
|
return self._eos_id |
|
|
|
@property |
|
def eos(self): |
|
return self._eos_id |
|
|
|
@property |
|
def mask(self): |
|
return self._mask_id |
|
|
|
@property |
|
def additional_special_tokens_ids(self): |
|
return [self.vocab[k] for k in self._t5_tokens] |
|
|
|
class _MMSentencePieceTokenizer(_SentencePieceTokenizer): |
|
"""SentencePieceTokenizer-Megatron wrapper""" |
|
|
|
def __init__(self, model_file, vocab_extra_ids=0): |
|
super().__init__(model_file, vocab_extra_ids) |
|
|
|
|
|
def _initalize(self, vocab_extra_ids): |
|
self._populate_vocab() |
|
self._special_tokens = {} |
|
self._inv_special_tokens = {} |
|
|
|
self._t5_tokens = [] |
|
|
|
def _add_special_token(t): |
|
if t not in self._vocab: |
|
next_id = len(self._vocab) |
|
self._vocab[t] = next_id |
|
self._inv_vocab[next_id] = t |
|
self._special_tokens[t] = self._vocab[t] |
|
self._inv_special_tokens[self._vocab[t]] = t |
|
|
|
_add_special_token('<CLS>') |
|
self._cls_id = self._vocab['<CLS>'] |
|
_add_special_token('<SEP>') |
|
self._sep_id = self._vocab['<SEP>'] |
|
_add_special_token('<EOD>') |
|
self._eod_id = self._vocab['<EOD>'] |
|
_add_special_token('<MASK>') |
|
self._mask_id = self._vocab['<MASK>'] |
|
|
|
_add_special_token('<SOA>') |
|
self._soa_id = self._vocab['<SOA>'] |
|
_add_special_token('<EOA>') |
|
self._eoa_id = self._vocab['<EOA>'] |
|
_add_special_token('<SOV>') |
|
self._sov_id = self._vocab['<SOV>'] |
|
_add_special_token('<EOV>') |
|
self._eov_id = self._vocab['<EOV>'] |
|
_add_special_token('<SOI>') |
|
self._soi_id = self._vocab['<SOI>'] |
|
_add_special_token('<EOI>') |
|
self._eoi_id = self._vocab['<EOI>'] |
|
_add_special_token('<s_local>') |
|
self._s_local_id = self._vocab['<s_local>'] |
|
_add_special_token('<e_local>') |
|
self._e_local_id = self._vocab['<e_local>'] |
|
_add_special_token('<s_global>') |
|
self._s_global_id = self._vocab['<s_global>'] |
|
_add_special_token('<e_global>') |
|
self._e_global_id = self._vocab['<e_global>'] |
|
_add_special_token('<stage_1>') |
|
self._stage_1_id = self._vocab['<stage_1>'] |
|
_add_special_token('<stage_2>') |
|
self._stage_2_id = self._vocab['<stage_2>'] |
|
pad_id = self.tokenizer.pad_id() |
|
try: |
|
pad_token = self.tokenizer.id_to_piece(pad_id) |
|
except IndexError: |
|
pad_token = '<PAD>' |
|
_add_special_token(pad_token) |
|
self._pad_id = self._vocab[pad_token] |
|
|
|
bos_id = self.tokenizer.bos_id() |
|
try: |
|
bos_token = self.tokenizer.id_to_piece(bos_id) |
|
except IndexError: |
|
bos_token = '<BOS>' |
|
_add_special_token(bos_token) |
|
self._bos_id = self._vocab[bos_token] |
|
|
|
eos_id = self.tokenizer.eos_id() |
|
try: |
|
eos_token = self.tokenizer.id_to_piece(eos_id) |
|
except IndexError: |
|
eos_token = '<EOS>' |
|
_add_special_token(eos_token) |
|
self._eos_id = self._vocab[eos_token] |
|
|
|
for i in range(vocab_extra_ids): |
|
t = "<extra_id_{}>".format(i) |
|
_add_special_token(t) |
|
self._t5_tokens += [t] |
|
|
|
@property |
|
def soa(self): |
|
return self._soa_id |
|
|
|
@property |
|
def eoa(self): |
|
return self._eoa_id |
|
|
|
@property |
|
def sov(self): |
|
return self._sov_id |
|
|
|
@property |
|
def eov(self): |
|
return self._eov_id |
|
|
|
@property |
|
def soi(self): |
|
return self._soi_id |
|
|
|
@property |
|
def eoi(self): |
|
return self._eoi_id |
|
|
|
@property |
|
def s_local(self): |
|
return self._s_local_id |
|
|
|
@property |
|
def e_local(self): |
|
return self._e_local_id |
|
|
|
@property |
|
def s_global(self): |
|
return self._s_global_id |
|
|
|
@property |
|
def e_global(self): |
|
return self._e_global_id |
|
|
|
@property |
|
def stage_1(self): |
|
return self._stage_1_id |
|
|
|
@property |
|
def stage_2(self): |
|
return self._stage_2_id |
|
|
