Refactor Qwen2VLTokenizer to use AutoTokenizer

Author: heavyrain-lzy

flagscale/train/megatron/training/tokenizer/tokenizer.py

@@ -246,39 +246,55 @@ def vocab_size(self):
 
 
 class _Qwen2VLTokenizer(_FlagScaleTokenizerBase):
-    """Qwen2-VL tokenizer with AutoProcessor and multimodal support."""
-
     def __init__(self, tokenizer_path, extra_vocab_size):
-        super().__init__(path=tokenizer_path)
+        super().__init__(tokenizer_path)
+        from transformers import AutoTokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_path,
+            padding_side="right",
+            use_fast=True,
+            split_special_tokens=False,
+            trust_remote_code=True,
+            revision = "main",
+            token = None,
+        )
+        self.extra_vocab_size = extra_vocab_size
+        self.special_tokens_map = {k:v for k, v in zip(self.tokenizer.all_special_tokens, self.tokenizer.all_special_ids)}
+        self.image_token = '<|image_pad|>'
+        self.video_token = '<|video_pad|>'
+        self.vision_start_token = '<|vision_start|>'
+        self.vision_end_token = '<|vision_end|>'
+
         from transformers import AutoProcessor
-        self.processor = AutoProcessor.from_pretrained(tokenizer_path, trust_remote_code=True)
-        self.tokenizer = self.processor.tokenizer
-        self.image_token = self.processor.image_processor.image_token
-        self.video_token = self.processor.image_processor.video_token
-        self.vision_start_token = self.processor.image_processor.vision_start_token
-        self.vision_end_token = self.processor.image_processor.vision_end_token
-        self.special_tokens_map = {
-            self.image_token: self.tokenizer.convert_tokens_to_ids(self.image_token),
-            self.video_token: self.tokenizer.convert_tokens_to_ids(self.video_token),
-            self.vision_start_token: self.tokenizer.convert_tokens_to_ids(self.vision_start_token),
-            self.vision_end_token: self.tokenizer.convert_tokens_to_ids(self.vision_end_token),
-        }
-        self._vocab_size = len(self.tokenizer) + extra_vocab_size
-        self._inv_vocab = None
+        self.processor = AutoProcessor.from_pretrained(
+            tokenizer_path,
+            revision = "main",
+            token = None,
+        )
+        # NOTE: In Qwen2-VL, template in chat_template.json is same within tokenizer_config.json and both can be used.
+        # However, in Qwen 2.5-VL, the two templates are different and only the one in chat_template.json is OK.
+        self.chat_template = self.processor.chat_template
+
+    def __call__(self, text, return_tensors=None,
+                    padding=None, max_length=None, truncation=None, add_special_tokens=None):
+
+        return self.tokenizer(text, return_tensors=return_tensors, padding=padding,
+                max_length=max_length, truncation=truncation, add_special_tokens=add_special_tokens)
 
+    def apply_chat_template(self, conversations, tokenize:bool=True, **kwargs):
+        return self.tokenizer.apply_chat_template(conversations, tokenize=tokenize, chat_template=self.chat_template, **kwargs)
+    
     @property
     def vocab_size(self):
-        return self._vocab_size
+        return self.tokenizer.vocab_size + self.extra_vocab_size
 
     @property
     def vocab(self):
-        return self.tokenizer.get_vocab()
+        return self.tokenizer.vocab
 
     @property
     def inv_vocab(self):
-        if self._inv_vocab is None:
-            self._inv_vocab = {v: k for k, v in self.vocab.items()}
-        return self._inv_vocab
+        return self.tokenizer.decoder
 
     def tokenize(self, text):
         return self.tokenizer.tokenize(text)
@@ -301,27 +317,26 @@ def pad_token_id(self):
     @property
     def eos_token_id(self):
         return self.tokenizer.eos_token_id
-
+    
     @property
     def image_token_id(self):
         return self.special_tokens_map[self.image_token]
-
+    
     @property
     def video_token_id(self):
         return self.special_tokens_map[self.video_token]
-
+    
     @property
     def vision_start_token_id(self):
         return self.special_tokens_map[self.vision_start_token]
-
+    
     @property
     def vision_end_token_id(self):
         return self.special_tokens_map[self.vision_end_token]
-
+    
     def encode(self, x):
         return self.tokenizer.encode(x)
 
-
 class _RWKVTokenizerFS(_FlagScaleTokenizerBase):
     """RWKV Trie-based tokenizer."""