|
from typing import List, Union, Optional |
|
|
|
import pypinyin |
|
import torch |
|
from torch import NoneType |
|
|
|
from transformers import BertTokenizerFast |
|
|
|
|
|
class Pinyin2(object): |
|
def __init__(self): |
|
super(Pinyin2, self).__init__() |
|
pho_vocab = ['P'] |
|
pho_vocab += [chr(x) for x in range(ord('1'), ord('5') + 1)] |
|
pho_vocab += [chr(x) for x in range(ord('a'), ord('z') + 1)] |
|
pho_vocab += ['U'] |
|
assert len(pho_vocab) == 33 |
|
self.pho_vocab_size = len(pho_vocab) |
|
self.pho_vocab = {c: idx for idx, c in enumerate(pho_vocab)} |
|
|
|
def get_pho_size(self): |
|
return self.pho_vocab_size |
|
|
|
@staticmethod |
|
def get_pinyin(c): |
|
if len(c) > 1: |
|
return 'U' |
|
s = pypinyin.pinyin( |
|
c, |
|
style=pypinyin.Style.TONE3, |
|
neutral_tone_with_five=True, |
|
errors=lambda x: ['U' for _ in x], |
|
)[0][0] |
|
if s == 'U': |
|
return s |
|
assert isinstance(s, str) |
|
assert s[-1] in '12345' |
|
s = s[-1] + s[:-1] |
|
return s |
|
|
|
def convert(self, chars): |
|
pinyins = list(map(self.get_pinyin, chars)) |
|
pinyin_ids = [list(map(self.pho_vocab.get, pinyin)) for pinyin in pinyins] |
|
pinyin_lens = [len(pinyin) for pinyin in pinyins] |
|
pinyin_ids = torch.nn.utils.rnn.pad_sequence( |
|
[torch.tensor(x) for x in pinyin_ids], |
|
batch_first=True, |
|
padding_value=0, |
|
) |
|
return pinyin_ids, pinyin_lens |
|
|
|
|
|
class ReaLiSeTokenizer(BertTokenizerFast): |
|
|
|
def __init__(self, **kwargs): |
|
super(ReaLiSeTokenizer, self).__init__(**kwargs) |
|
|
|
self.pho2_convertor = Pinyin2() |
|
|
|
def __call__(self, |
|
text: Union[str, List[str], List[List[str]]] = None, |
|
text_pair: Union[str, List[str], List[List[str]], NoneType] = None, |
|
text_target: Union[str, List[str], List[List[str]]] = None, |
|
text_pair_target: Union[str, List[str], List[List[str]], NoneType] = None, |
|
add_special_tokens: bool = True, |
|
padding=False, |
|
truncation=None, |
|
max_length: Optional[int] = None, |
|
stride: int = 0, |
|
is_split_into_words: bool = False, |
|
pad_to_multiple_of: Optional[int] = None, |
|
return_tensors=None, |
|
return_token_type_ids: Optional[bool] = None, |
|
return_attention_mask: Optional[bool] = None, |
|
return_overflowing_tokens: bool = False, return_special_tokens_mask: bool = False, |
|
return_offsets_mapping: bool = False, |
|
return_length: bool = False, |
|
verbose: bool = True, **kwargs): |
|
encoding = super(ReaLiSeTokenizer, self).__call__( |
|
text=text, |
|
text_pair=text_pair, |
|
text_target=text_target, |
|
text_pair_target=text_pair_target, |
|
add_special_tokens=add_special_tokens, |
|
padding=padding, |
|
truncation=truncation, |
|
max_length=max_length, |
|
stride=stride, |
|
is_split_into_words=is_split_into_words, |
|
pad_to_multiple_of=pad_to_multiple_of, |
|
return_tensors=return_tensors, |
|
return_token_type_ids=return_token_type_ids, |
|
return_attention_mask=return_attention_mask, |
|
return_overflowing_tokens=return_overflowing_tokens, |
|
return_offsets_mapping=return_offsets_mapping, |
|
return_length=return_length, |
|
verbose=verbose, |
|
) |
|
|
|
input_ids = encoding['input_ids'] |
|
if type(text) == str and return_tensors is None: |
|
input_ids = [input_ids] |
|
|
|
pho_idx_list = [] |
|
pho_lens_list = [] |
|
for ids in input_ids: |
|
chars = self.convert_ids_to_tokens(ids) |
|
pho_idx, pho_lens = self.pho2_convertor.convert(chars) |
|
if return_tensors is None: |
|
pho_idx = pho_idx.tolist() |
|
pho_idx_list.append(pho_idx) |
|
pho_lens_list += pho_lens |
|
|
|
pho_idx = pho_idx_list |
|
pho_lens = pho_lens_list |
|
if return_tensors == 'pt': |
|
pho_idx = torch.vstack(pho_idx) |
|
pho_lens = torch.LongTensor(pho_lens) |
|
|
|
if type(text) == str and return_tensors is None: |
|
pho_idx = pho_idx[0] |
|
|
|
encoding['pho_idx'] = pho_idx |
|
encoding['pho_lens'] = pho_lens |
|
|
|
return encoding |
|
|
