project-llm-trainer 0.9.3__py3-none-any.whl → 0.9.4__py3-none-any.whl

llm_trainer/tokenizer.py

@@ -1,22 +1,15 @@
 import os
 import warnings
 from typing import List, Dict, Union
-from transformers import Qwen2TokenizerFast
-from transformers import AddedToken
-from transformers import LlamaTokenizerFast
+from transformers import AutoTokenizer
 import torch
-
-TOKEN_TYPE_QWEN = 'qwen'
-TOKEN_TYPE_ZH_LLAMA = "zh_llama"
-
-AVAILABLE_TOKEN_TYPES = [TOKEN_TYPE_QWEN, TOKEN_TYPE_ZH_LLAMA]
+from .log import log
 
 
 class Tokenizer:
-    def __init__(self, token_type: str = TOKEN_TYPE_ZH_LLAMA):
-        super().__init__()
-        assert token_type in AVAILABLE_TOKEN_TYPES, 'token type is unavailable'
-        self.token_type = token_type
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(os.environ['TOKEN_DIR'])
+        log(f'is fast tokenizer={self.tokenizer.is_fast}')
 
         self.text_end = '</s>'
 
@@ -36,31 +29,6 @@ class Tokenizer:
 
         self.text_image = '<image>'
 
-        if token_type == TOKEN_TYPE_QWEN:
-            self.tokenizer = Qwen2TokenizerFast(
-                vocab_file=f"{os.environ['TOKEN_DIR']}qwen_vocab.json",
-                merges_file=f"{os.environ['TOKEN_DIR']}qwen_merges.txt",
-                unk_token=self.text_unk,
-                eos_token=self.text_end,
-                pad_token=self.text_pad
-            )
-            additional_special_tokens = [
-                AddedToken(self.text_user, lstrip=False, rstrip=False),
-                AddedToken(self.text_assistant, lstrip=False, rstrip=False),
-                AddedToken(self.text_think_start, lstrip=False, rstrip=False),
-                AddedToken(self.text_think_end, lstrip=False, rstrip=False),
-                AddedToken(self.text_answer_start, lstrip=False, rstrip=False),
-                AddedToken(self.text_answer_end, lstrip=False, rstrip=False),
-                AddedToken(self.text_system, lstrip=False, rstrip=False),
-                AddedToken(self.text_image, lstrip=False, rstrip=False),
-            ]
-
-            self.tokenizer.add_special_tokens({"additional_special_tokens": additional_special_tokens})
-        else:
-            self.tokenizer = LlamaTokenizerFast.from_pretrained(os.environ['TOKEN_DIR'])
-            # self.tokenizer = AutoTokenizer.from_pretrained(os.environ['TOKEN_DIR'])
-            # self.tokenizer = PreTrainedTokenizerFast.from_pretrained(os.environ['TOKEN_DIR'], trust_remote_code=True)
-
         self.end = self.tokenizer.convert_tokens_to_ids(self.text_end)
 
         self.pad = self.tokenizer.convert_tokens_to_ids(self.text_pad)
@@ -105,8 +73,15 @@ class Tokenizer:
             padding = False,
             truncation = False,
             covert_tensor: bool = False,
+            return_attention_mask: bool = False
     ) -> Union[torch.Tensor, List[List[int]]]:
-        encoded = self.tokenizer(text, padding=padding, truncation=truncation)['input_ids']
+        encoded = self.tokenizer(
+            text,
+            padding=padding,
+            truncation=truncation,
+            return_attention_mask=return_attention_mask
+        )['input_ids']
+
         if covert_tensor:
             encoded = torch.tensor(encoded, dtype=torch.long)
 

llm_trainer/tools.py

@@ -28,7 +28,7 @@ class TrainerTools:
 
             self.parallel = self._new_parallel()
 
-            self.tokenizer = Tokenizer(os.environ.get('TOKENIZERS_TYPE', 'zh_llama'))
+            self.tokenizer = Tokenizer()
             self.use_amp = 'cuda' in self.parallel.device and not isinstance(self.parallel, DsParallel)
 
             log(f'word_size={self.parallel.world_size}, use_amp={self.use_amp}')

{project_llm_trainer-0.9.3.dist-info → project_llm_trainer-0.9.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: project_llm_trainer
-Version: 0.9.3
+Version: 0.9.4
 Summary: LLM and VLM trainer
 Author: qibin
 Author-email: qibin0506@gmail.com

{project_llm_trainer-0.9.3.dist-info → project_llm_trainer-0.9.4.dist-info}/RECORD

@@ -15,19 +15,19 @@ llm_trainer/parallel_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,
 llm_trainer/partition_utils.py,sha256=eEYNhfEIF4hGzZ3OLa6sEBIECz261drptEz_n7fZYtk,8396
 llm_trainer/scheduler.py,sha256=LAI_0VxClsIQkix0bRoduRD4vPfVuIZDhZgTAT_KK8k,4901
 llm_trainer/sft_trainer.py,sha256=rSOGZx53jMgOuJdztfxQASYJ62uD0dVaih4IAnSwGBc,1787
-llm_trainer/tokenizer.py,sha256=23MgueOq2FtIXR8KUDgmzK9jfSUXzx0v1VmolbvR-7U,7487
-llm_trainer/tools.py,sha256=5op5qrjjkK-Lr9oes5VxIVnOVYOYGoAdlIJq9mPUf64,2637
+llm_trainer/tokenizer.py,sha256=_jk4zb9JSHjwiWlWAW40zUqf85AmRrwDfQ_L2jkhRoM,5955
+llm_trainer/tools.py,sha256=Yca2OkXqFto37Jw13Feu1xzAP0s1proUcVKtLvdGxrk,2592
 llm_trainer/train_configs.py,sha256=afXUZ7M_Uoj0B3c2Nwf5xE-Lv7QAZZHTdW8LBw-QeWE,7704
 llm_trainer/trainer.py,sha256=bVghqvQY4bvYAZFPgyh2ywX8WanqAC525Lkg8bNv4FQ,29721
 llm_trainer/utils.py,sha256=xC5plG-8-_Al5yIF5xIU5lroOcBBk98TEhtUJrazZPE,12305
-project_llm_trainer-0.9.3.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.9.3.data/scripts/ddp_train,sha256=eZSud6KYQAoKLsYB5QB-FI2zq5AZm6Apq1azKdupV3o,477
-project_llm_trainer-0.9.3.data/scripts/ds_train,sha256=41q4rOxwbvZDUY0FDdAIpG13PEaUWBpthhvFvww8uOc,388
-project_llm_trainer-0.9.3.data/scripts/plot_loss,sha256=O9ooioAJ-79-X06LosgqF8XOqQe-beRxYm3LsLunmoU,908
-project_llm_trainer-0.9.3.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.9.3.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.9.3.data/scripts/smart_train,sha256=N8dp2n7k6bghGczedBVwOdtf1O66oM_cNPh9QmZt0bM,914
-project_llm_trainer-0.9.3.dist-info/METADATA,sha256=MXUhvp1gftCShHCKRPRxBgadUWmKmnC6trtU0Qf_Y1k,195
-project_llm_trainer-0.9.3.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.9.3.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.9.3.dist-info/RECORD,,
+project_llm_trainer-0.9.4.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.9.4.data/scripts/ddp_train,sha256=eZSud6KYQAoKLsYB5QB-FI2zq5AZm6Apq1azKdupV3o,477
+project_llm_trainer-0.9.4.data/scripts/ds_train,sha256=41q4rOxwbvZDUY0FDdAIpG13PEaUWBpthhvFvww8uOc,388
+project_llm_trainer-0.9.4.data/scripts/plot_loss,sha256=O9ooioAJ-79-X06LosgqF8XOqQe-beRxYm3LsLunmoU,908
+project_llm_trainer-0.9.4.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.9.4.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.9.4.data/scripts/smart_train,sha256=N8dp2n7k6bghGczedBVwOdtf1O66oM_cNPh9QmZt0bM,914
+project_llm_trainer-0.9.4.dist-info/METADATA,sha256=OAh8orL_7cliVbV9KK_2YxDheiD0NfmdJh-YU9_G92U,195
+project_llm_trainer-0.9.4.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.9.4.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.9.4.dist-info/RECORD,,