Upload bert_tokenizer.py

bert_tokenizer.py

@@ -2,20 +2,24 @@ import json
 import os
 import time
 from pathlib import Path
-from typing import List
+from types import NoneType
+from typing import List, Union, Optional
 
 import tokenizers
 import torch
 from huggingface_hub import hf_hub_download
 from huggingface_hub.file_download import http_user_agent
 from pypinyin import pinyin, Style
+from transformers.tokenization_utils_base import TruncationStrategy
+from transformers.utils import PaddingStrategy
+from transformers.utils.generic import TensorType
 
 try:
     from tokenizers import BertWordPieceTokenizer
 except:
     from tokenizers.implementations import BertWordPieceTokenizer
 
-from transformers import BertTokenizerFast
+from transformers import BertTokenizerFast, BatchEncoding
 
 cache_path = Path(os.path.abspath(__file__)).parent
 
@@ -60,6 +64,64 @@ class ChineseBertTokenizer(BertTokenizerFast):
         with open(config_path / 'pinyin2tensor.json', encoding='utf8') as fin:
             self.pinyin2tensor = json.load(fin)
 
+    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: Union[bool, str, PaddingStrategy] = False,
+                 truncation: Union[bool, str, TruncationStrategy] = None,
+                 max_length: Optional[int] = None,
+                 stride: int = 0,
+                 is_split_into_words: bool = False,
+                 pad_to_multiple_of: Optional[int] = None,
+                 return_tensors: Union[str, TensorType, NoneType] = 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) -> BatchEncoding:
+        encoding = super.__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
+
+        pinyin_ids = None
+        if type(text) == str:
+            pinyin_ids = self.convert_ids_to_pinyin_ids(input_ids)
+
+        if type(text) == list:
+            pinyin_ids = []
+            for ids in input_ids:
+                pinyin_ids.append(self.convert_ids_to_pinyin_ids(ids))
+
+        if torch.is_tensor(encoding.input_ids):
+            pinyin_ids = torch.LongTensor(pinyin_ids)
+
+        encoding['pinyin_ids'] = pinyin_ids
+
+        return encoding
+
     def tokenize_sentence(self, sentence):
         # convert sentence to ids
         tokenizer_output = self.tokenizer.encode(sentence)
@@ -73,6 +135,33 @@ class ChineseBertTokenizer(BertTokenizerFast):
         pinyin_ids = torch.LongTensor(pinyin_tokens).view(-1)
         return input_ids, pinyin_ids
 
+    def convert_ids_to_pinyin_ids(self, ids: List[str]):
+        pinyin_ids = []
+        tokens = self.convert_tokens_to_ids(ids)
+        for token in tokens:
+            if len(token) > 1:
+                pinyin_ids.append([0] * 8)
+                continue
+
+            pinyin_string = pinyin(token, style=Style.TONE3, errors=lambda x: [['not chinese'] for _ in x])[0]
+
+            if pinyin_string == "not chinese":
+                pinyin_ids.append([0] * 8)
+                continue
+
+            if pinyin_string in self.pinyin2tensor:
+                pinyin_ids.append(self.pinyin2tensor[pinyin_string])
+            else:
+                ids = [0] * 8
+                for i, p in enumerate(pinyin_string):
+                    if p not in self.pinyin_dict["char2idx"]:
+                        ids = [0] * 8
+                        break
+                    ids[i] = self.pinyin_dict["char2idx"][p]
+                pinyin_ids.append(pinyin_ids)
+
+        return pinyin_ids
+
     def convert_sentence_to_pinyin_ids(self, sentence: str, tokenizer_output: tokenizers.Encoding) -> List[List[int]]:
         # get pinyin of a sentence
         pinyin_list = pinyin(sentence, style=Style.TONE3, heteronym=True, errors=lambda x: [['not chinese'] for _ in x])