| |
| """ |
| @author:cb |
| @contact:chenbo@bat100.net |
| @time:2023/5/30 14:21 |
| @filename:tokenization.py |
| @software:PyCharm |
| @description: |
| """ |
| import re |
| from transformers import FSMTTokenizer as fsmt |
|
|
|
|
| class FSMTTokenizer(fsmt): |
| space_re = re.compile('\s*(?=[^a-zA-Z0-9 ]+)\s*') |
|
|
| def moses_tokenize(self, text, lang): |
| if lang not in self.cache_moses_tokenizer: |
| moses_tokenizer = self.sm.MosesTokenizer(lang=lang) |
| self.cache_moses_tokenizer[lang] = moses_tokenizer |
| return self.cache_moses_tokenizer[lang].tokenize( |
| text, aggressive_dash_splits=True, return_str=False, escape=False |
| ) |
|
|
| def _switch_to_input_mode(self): |
| self.lang_prefix, self.lang_prefix_id = 'en', 64812 |
|
|
| def _switch_to_target_mode(self): |
| self.lang_prefix, self.lang_prefix_id = 'zh', 64870 |
|
|
| def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None): |
| """ |
| Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and |
| adding special tokens. A FAIRSEQ Transformer sequence has the following format: |
| |
| - single sequence: `<s> X </s>` |
| - pair of sequences: `<s> A </s> B </s>` |
| |
| Args: |
| token_ids_0 (`List[int]`): |
| List of IDs to which the special tokens will be added. |
| token_ids_1 (`List[int]`, *optional*): |
| Optional second list of IDs for sequence pairs. |
| |
| Returns: |
| `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. |
| """ |
| sep = [self.sep_token_id] |
| token_ids_0 = [self.lang_prefix_id] + token_ids_0 |
| |
| if token_ids_1 is None: |
| return token_ids_0 + sep |
| return token_ids_0 + sep + token_ids_1 + sep |
|
|
| def moses_pipeline(self, text, lang): |
| text = self.moses_punct_norm(text, lang) |
| return text |
|
|
| def _tokenize(self, text, lang="en", bypass_tokenizer=False): |
| """ |
| 原版FSMTTokenizer会把中文标点英文化，故重写 |
| :param text: |
| :param lang: |
| :param bypass_tokenizer: |
| :return: |
| """ |
| if self.do_lower_case: |
| text = text.lower() |
| if bypass_tokenizer: |
| text = text.split() |
| else: |
| text = self.moses_pipeline(text, lang=self.lang_prefix) |
| text = self.moses_tokenize(text, lang=self.lang_prefix) |
|
|
| split_tokens = [] |
| for token in text: |
| if token: |
| split_tokens.extend(list(self.bpe(token).split(" "))) |
|
|
| return split_tokens |
|
|
| def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs): |
| """ |
| |
| :param text: |
| :param is_split_into_words: |
| :param kwargs: |
| :return: |
| """ |
| if kwargs.get('src', True): |
| self._switch_to_input_mode() |
| else: |
| self._switch_to_target_mode() |
| return super(FSMTTokenizer, self).prepare_for_tokenization(text, is_split_into_words=False, **kwargs) |
| |
| def convert_tokens_to_string(self, tokens): |
| """ |
| 删除非英文字母前后的空格，业务上处理更合适 |
| :param tokens: |
| :return: |
| """ |
| tokens = super(FSMTTokenizer, self).convert_tokens_to_string(tokens) |
| tokens = FSMTTokenizer.space_re.sub('', tokens) |
| return tokens |
|
|
|
|
| if __name__ == '__main__': |
| tokenizer = FSMTTokenizer.from_pretrained(r'./') |
| r = tokenizer.tokenize(['hello', 'hi']) |
| print(r) |
