Update tokenization_chatglm.py

tokenization_chatglm.py

@@ -1,11 +1,13 @@
+import json
 import os
-import torch
+import re
 from typing import List, Optional, Union, Dict
 from sentencepiece import SentencePieceProcessor
 from transformers import PreTrainedTokenizer
 from transformers.utils import logging, PaddingStrategy
 from transformers.tokenization_utils_base import EncodedInput, BatchEncoding
 
+
 class SPTokenizer:
     def __init__(self, model_path: str):
         # reload tokenizer
@@ -30,6 +32,7 @@ class SPTokenizer:
     def tokenize(self, s: str):
         return self.sp_model.EncodeAsPieces(s)
 
+
     def encode(self, s: str, bos: bool = False, eos: bool = False) -> List[int]:
         assert type(s) is str
         t = self.sp_model.encode(s)
@@ -40,7 +43,18 @@ class SPTokenizer:
         return t
 
     def decode(self, t: List[int]) -> str:
-        return self.sp_model.decode(t)
+        text, buffer = "", []
+        for token in t:
+            if token in self.index_special_tokens:
+                if buffer:
+                    text += self.sp_model.decode(buffer)
+                    buffer = []
+                text += self.index_special_tokens[token]
+            else:
+                buffer.append(token)
+        if buffer:
+            text += self.sp_model.decode(buffer)
+        return text
 
     def decode_tokens(self, tokens: List[str]) -> str:
         text = self.sp_model.DecodePieces(tokens)
@@ -54,7 +68,9 @@ class SPTokenizer:
 
     def convert_id_to_token(self, index):
         """Converts an index (integer) in a token (str) using the vocab."""
-        if index in self.index_special_tokens or index in [self.eos_id, self.bos_id, self.pad_id] or index < 0:
+        if index in self.index_special_tokens:
+            return self.index_special_tokens[index]
+        if index in [self.eos_id, self.bos_id, self.pad_id] or index < 0 or index > self.sp_model.vocab_size():
             return ""
         return self.sp_model.IdToPiece(index)
 
@@ -64,8 +80,8 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
 
     model_input_names = ["input_ids", "attention_mask", "position_ids"]
 
-    def __init__(self, vocab_file, padding_side="left", **kwargs):
-        super().__init__(padding_side=padding_side, **kwargs)
+    def __init__(self, vocab_file, padding_side="left", clean_up_tokenization_spaces=False, encode_special_tokens=False,
+                 **kwargs):
         self.name = "GLMTokenizer"
 
         self.vocab_file = vocab_file
@@ -75,6 +91,10 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
             "<eos>": self.tokenizer.eos_id,
             "<pad>": self.tokenizer.pad_id
         }
+        self.encode_special_tokens = encode_special_tokens
+        super().__init__(padding_side=padding_side, clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+                         encode_special_tokens=encode_special_tokens,
+                         **kwargs)
 
     def get_command(self, token):
         if token in self.special_tokens:
@@ -82,6 +102,10 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         assert token in self.tokenizer.special_tokens, f"{token} is not a special token for {self.name}"
         return self.tokenizer.special_tokens[token]
 
+    @property
+    def unk_token(self) -> str:
+        return "<unk>"
+
     @property
     def pad_token(self) -> str:
         return "<unk>"
@@ -163,6 +187,7 @@ class ChatGLMTokenizer(PreTrainedTokenizer):
         prompt += "[Round {}]\n\n问：{}\n\n答：".format(len(history) + 1, query)
         return prompt
 
+
     def build_inputs_with_special_tokens(
             self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
     ) -> List[int]: