Update README.md

README.md

@@ -25,7 +25,39 @@ Task: Summarization
 import jieba_fast as jieba
 jieba.initialize()
 from transformers import PegasusForConditionalGeneration,BertTokenizer
-                                            
+class PegasusTokenizer(BertTokenizer):
+    model_input_names = ["input_ids", "attention_mask"]
+    def __init__(self, **kwargs):
+        super().__init__(pre_tokenizer=lambda x: jieba.cut(x, HMM=False), **kwargs)
+        self.add_special_tokens({'additional_special_tokens':["<mask_1>"]})
+
+    def build_inputs_with_special_tokens(
+            self,
+            token_ids_0: List[int],
+            token_ids_1: Optional[List[int]] = None) -> List[int]:
+
+        if token_ids_1 is None:
+            return token_ids_0 + [self.eos_token_id]
+        return token_ids_0 + token_ids_1 + [self.eos_token_id]
+
+    def _special_token_mask(self, seq):
+        all_special_ids = set(
+            self.all_special_ids)  # call it once instead of inside list comp
+        # all_special_ids.remove(self.unk_token_id)  # <unk> is only sometimes special
+        return [1 if x in all_special_ids else 0 for x in seq]
+
+    def get_special_tokens_mask(
+            self,
+            token_ids_0: List[int],
+            token_ids_1: Optional[List[int]] = None,
+            already_has_special_tokens: bool = False) -> List[int]:
+        if already_has_special_tokens:
+            return self._special_token_mask(token_ids_0)
+        elif token_ids_1 is None:
+            return self._special_token_mask(token_ids_0) + [self.eos_token_id]
+        else:
+            return self._special_token_mask(token_ids_0 +
+                                            token_ids_1) + [self.eos_token_id]                              
 model = PegasusForConditionalGeneration.from_pretrained('dongxq/test_model')
 tokenizer = BertTokenizer.from_pretrained('dongxq/test_model')