Update README.md

README.md

@@ -35,7 +35,24 @@ import logging
 from typing import List, Tuple
 from transformers import AutoConfig
 from transformers.models.mt5.modeling_mt5 import MT5ForConditionalGeneration
-from utils import T5PegasusTokenizer
+
+import jieba
+from functools import partial
+from transformers import BertTokenizer
+
+class T5PegasusTokenizer(BertTokenizer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.pre_tokenizer = partial(jieba.cut, HMM=False)
+
+    def _tokenize(self, text, *arg, **kwargs):
+        split_tokens = []
+        for text in self.pre_tokenizer(text):
+            if text in self.vocab:
+                split_tokens.append(text)
+            else:
+                split_tokens.extend(super()._tokenize(text))
+        return split_tokens
 
 def load_model(model_path):
     config = AutoConfig.from_pretrained(model_path)