project-llm-trainer 0.7.7__py3-none-any.whl → 0.7.8__py3-none-any.whl

llm_trainer/dataset.py

@@ -1,28 +1,30 @@
-import os.path
-
 import torch
 from torch.utils.data import Dataset
 import pickle
 import csv
+import json
 
 from .tools import TrainerTools
 from .utils import repeat_image_tok
 
 
-def _try_load_pkl(file_path: str):
-    tokens = None
-    try:
-        with open(file_path, 'rb') as f:
-            tokens = pickle.load(f)
-    except Exception as e:
-        raise e
-    finally:
-        return tokens
+"""
+support jsonl and pkl
+"""
+def _get_file_type(file_path: str):
+    if file_path.endswith('.jsonl'):
+        return 'jsonl'
+    elif file_path.endswith('.pkl'):
+        return 'pkl'
+
+    return None
 
 
 class TextDataset(Dataset):
     """
-    适用于pretrain阶段
+    适用于pretrain阶段，数据格式支持jsonl和pkl，如果是jsonl会在init阶段全部encode成token
+    jsonl: {'text': 'text1'}\n{'text': 'text2'}
+    pkl: [0, 1, 2, 3 ...]
     """
     def __init__(
             self,
@@ -34,19 +36,17 @@ class TextDataset(Dataset):
 
         self.input_ids = []
 
-        tokens = _try_load_pkl(file_path)
-        if not tokens:
-            cache_file = f'{file_path}.cache'
-            if os.path.exists(cache_file):
-                tokens = _try_load_pkl(cache_file)
-            else:
-                tokens = []
-                with open(file_path, 'r') as f:
-                    for line in f:
-                        tokens.extend(TrainerTools().tokenizer.encode(line))
-
-                with open(cache_file, 'wb') as f:
-                    pickle.dump(tokens, f)
+        file_type = _get_file_type(file_path)
+        if file_type == 'jsonl':
+            tokens = []
+            with open(file_path, 'r') as f:
+                for line in f:
+                    tokens.extend(TrainerTools().tokenizer.encode(json.loads(line.strip())['text']))
+        elif file_type == 'pkl':
+            with open(file_path, 'rb') as f:
+                tokens = pickle.load(f)
+        else:
+            raise Exception(f'unsupported file type for {file_path}')
 
         for i in range(0, len(tokens) - block_size + 1, stride):
             self.input_ids.append(tokens[i:i+block_size])
@@ -60,7 +60,21 @@ class TextDataset(Dataset):
 
 class LineByLineTextDataset(Dataset):
     """
-    适用于sft阶段
+    适用于sft阶段，数据格式支持jsonl和pkl，如果是jsonl，则会在getitem阶段encode成token
+    jsonl: [
+            {'role': 'system', 'content': 'system_content'},
+            {'role': 'user', 'content': 'user_content'},
+            {'role': 'assistant', 'think': 'think_content', 'content': 'assistant_content'}
+           ]\n
+           [
+            {'role': 'system', 'content': 'system_content'},
+            {'role': 'user', 'content': 'user_content'},
+            {'role': 'assistant', 'think': 'think_content', 'content': 'assistant_content'}
+           ]
+    pkl: [
+            [0, 1, 2, 3],
+            [4, 5, 6, 7]
+         ]
     """
     def __init__(
             self,
@@ -75,22 +89,20 @@ class LineByLineTextDataset(Dataset):
         self.tokens_per_image = tokens_per_image
         self.input_ids = []
         self.image_tags = []
-
-        tokens = _try_load_pkl(file_path)
-        if not tokens:
-            cache_file = f'{file_path}.cache'
-            if os.path.exists(cache_file):
-                tokens = _try_load_pkl(cache_file)
-            else:
-                tokens = []
-                with open(file_path, 'r') as f:
-                    for line in f:
-                        tokens.append(TrainerTools().tokenizer.encode(line))
-
-                with open(cache_file, 'wb') as f:
-                    pickle.dump(tokens, f)
-
-        self.input_ids = tokens
+        self.plain_text = False
+
+        file_type = _get_file_type(file_path)
+        if file_type == 'jsonl':
+            self.plain_text = True
+
+            with open(file_path, 'r') as f:
+                for line in f:
+                    self.input_ids.append(json.loads(line.strip()))
+        elif file_type == 'pkl':
+            with open(file_path, 'rb') as f:
+                self.input_ids = pickle.load(f)
+        else:
+            raise Exception(f'unsupported file type for {file_path}')
 
         if image_tags_file_path:
             with open(image_tags_file_path, 'r') as f:
@@ -102,8 +114,14 @@ class LineByLineTextDataset(Dataset):
         return len(self.input_ids)
 
     def __getitem__(self, item):
-        inputs = torch.tensor(self.input_ids[item]).long()
+        if self.plain_text:
+            inputs = TrainerTools().tokenizer.apply_chat_template(self.input_ids[item])
+        else:
+            inputs = self.input_ids[item]
+
+        inputs = torch.tensor(inputs).long()
         image_tag = self.image_tags[item] if self.image_tags else None
+
         if self.tokens_per_image != -1:
             inputs = repeat_image_tok(inputs, self.tokens_per_image)
         else:
@@ -111,48 +129,134 @@ class LineByLineTextDataset(Dataset):
 
         inputs = inputs[:self.max_len]
 
-        return {'inputs': inputs, 'image_tag': image_tag}
+        return {
+            'inputs': inputs,
+            'image_tag': image_tag
+        }
 
 
 class DPODataset(Dataset):
+    """
+    适用于dpo阶段，数据格式支持jsonl和pkl，如果是jsonl，则会在getitem阶段encode成token
+    jsonl: {'chosen':
+                [{'role': 'system', 'content': 'system_content'},
+                {'role': 'user', 'content': 'user_content'},
+                {'role': 'assistant', 'think': 'think_content', 'content': 'assistant_content'}],
+            'rejected':
+                [{'role': 'system', 'content': 'system_content'},
+                {'role': 'user', 'content': 'user_content'},
+                {'role': 'assistant', 'think': 'think_content', 'content': 'assistant_content'}],
+            }\n
+           {'chosen':
+                [{'role': 'system', 'content': 'system_content'},
+                {'role': 'user', 'content': 'user_content'},
+                {'role': 'assistant', 'think': 'think_content', 'content': 'assistant_content'}],
+            'rejected':
+                [{'role': 'system', 'content': 'system_content'},
+                {'role': 'user', 'content': 'user_content'},
+                'role': 'assistant', 'think': 'think_content', 'content': 'assistant_content'}],
+            }
+    pkl: [
+            {'chosen': xxx, 'rejected': xxx},
+            {'chosen': xxx, 'rejected': xxx},
+         ]
+    """
     def __init__(self, file_path, max_len):
         self.max_len = max_len
         self.chosen_ids = []
         self.rejected_ids = []
-
-        # [{'chosen': xxx, 'rejected': xxx} ...]
-        tokens = _try_load_pkl(file_path)
-        for token in tokens:
-            self.chosen_ids.append(token['chosen'])
-            self.rejected_ids.append(token['rejected'])
+        self.plain_text = False
+
+        file_type = _get_file_type(file_path)
+        if file_type == 'jsonl':
+            self.plain_text = True
+
+            with open(file_path, 'r') as f:
+                for line in f:
+                    json_ = json.loads(line.strip())
+                    self.chosen_ids.append(json_['chosen'])
+                    self.rejected_ids.append(json_['rejected'])
+        elif file_type == 'pkl':
+            with open(file_path, 'rb') as f:
+                tokens = pickle.load(f)
+
+            for token in tokens:
+                self.chosen_ids.append(token['chosen'])
+                self.rejected_ids.append(token['rejected'])
+        else:
+            raise Exception(f'unsupported file type for {file_path}')
 
     def __len__(self):
         return len(self.chosen_ids)
 
     def __getitem__(self, item):
-        chosen_id = self.chosen_ids[item]
-        rejected_id = self.rejected_ids[item]
+        if self.plain_text:
+            chosen_id = TrainerTools().tokenizer.apply_chat_template(self.chosen_ids[item])
+            rejected_id = TrainerTools().tokenizer.apply_chat_template(self.rejected_ids[item])
+        else:
+            chosen_id = self.chosen_ids[item]
+            rejected_id = self.rejected_ids[item]
 
-        return {'chosen': chosen_id[:self.max_len], 'rejected': rejected_id[:self.max_len]}
+        return {
+            'chosen': chosen_id[:self.max_len],
+            'rejected': rejected_id[:self.max_len]
+        }
 
 
 class GRPORolloutDataset(Dataset):
+    """
+        适用于grpo(gspo)阶段，数据格式支持jsonl和pkl，如果是jsonl，则会在getitem阶段encode成token
+        jsonl: {'prompt':
+                    [{'role': 'system', 'content': 'system_content'},
+                    {'role': 'user', 'content': 'user_content'},
+                    {'role': 'assistant', 'think': 'think_content', 'content': 'assistant_content'}],
+                'answer': '10'
+               }\n
+               {'prompt':
+                    [{'role': 'system', 'content': 'system_content'},
+                    {'role': 'user', 'content': 'user_content'},
+                    {'role': 'assistant', 'think': 'think_content', 'content': 'assistant_content'}],
+                'answer': '10'
+               }
+        pkl: [
+                {'prompt': xxx, 'answer': xxx},
+                {'prompt': xxx, 'answer': xxx},
+             ]
+        """
     def __init__(self, file_path):
         self.questions = []
         self.answers = []
-
-        # [{'question': xxx, 'answer': ''}]
-        tokens = _try_load_pkl(file_path)
-        for token in tokens:
-            self.questions.append(token['prompt'])
-            self.answers.append(token['answer'])
+        self.plain_text = False
+
+        file_type = _get_file_type(file_path)
+        if file_type == 'jsonl':
+            self.plain_text = True
+
+            with open(file_path, 'r') as f:
+                for line in f:
+                    json_ = json.loads(line.strip())
+                    self.questions.append(json_['prompt'])
+                    self.answers.append(json_['answer'])
+        elif file_type == 'pkl':
+            with open(file_path, 'rb') as f:
+                tokens = pickle.load(f)
+
+            for token in tokens:
+                self.questions.append(token['prompt'])
+                self.answers.append(token['answer'])
+        else:
+            raise Exception(f'unsupported file type for {file_path}')
 
     def __len__(self):
         return len(self.questions)
 
     def __getitem__(self, item):
-        question = self.questions[item]
-        answer = self.answers[item]
+        if self.plain_text:
+            question = TrainerTools().tokenizer.apply_chat_template(self.questions[item])
+            answer = TrainerTools().tokenizer.encode(self.answers[item])
+        else:
+            question = self.questions[item]
+            answer = self.answers[item]
 
         return {
             'prompt': torch.tensor(question).long(),

{project_llm_trainer-0.7.7.dist-info → project_llm_trainer-0.7.8.dist-info}/METADATA

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

{project_llm_trainer-0.7.7.dist-info → project_llm_trainer-0.7.8.dist-info}/RECORD

@@ -1,6 +1,6 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
 llm_trainer/checkpoint.py,sha256=X5ZeUtJlxVz7pnWQLaS-y7UIZOaOAnZTt2L8rSAPzUs,4428
-llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
+llm_trainer/dataset.py,sha256=UL3fGeM4XSlyNQRZH-139u3LujqAQx3YyaxNRewk6LE,8935
 llm_trainer/dpo_trainer.py,sha256=RMfbTsl3eav4yTJ2PK59mi6a0ECVOg8WwYVsHvMbNUE,12353
 llm_trainer/ds_checkpoint.py,sha256=X2IWgpgi0yOtogph7n6DEwvK_0Ceb7juu1WMutv3HSk,2270
 llm_trainer/eval.py,sha256=ZyUfSo2Q8P-lrCdPEnGkoo5pGubd0AabREK5eMISRII,1109
@@ -20,14 +20,14 @@ llm_trainer/tools.py,sha256=5op5qrjjkK-Lr9oes5VxIVnOVYOYGoAdlIJq9mPUf64,2637
 llm_trainer/train_configs.py,sha256=N3ykM1uaLHcSNRC8ErYIxp9VYhSP7voJyAP-2D4ZJe0,7574
 llm_trainer/trainer.py,sha256=jS31zEXIIj9BoPTPlmaGYq61x72HGCjKfS2u3_gOkDk,27924
 llm_trainer/utils.py,sha256=xcdzpvPvXRKqsOK2yB7PZ9GmOvZMDFcglDPUZY2hJTY,11484
-project_llm_trainer-0.7.7.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.7.7.data/scripts/ddp_train,sha256=Z-309mM56CN0m3bxoeC5us4LUuwuNnoiOm3-fDdLMjQ,566
-project_llm_trainer-0.7.7.data/scripts/ds_train,sha256=3nXNNKmYI7miqyBdf-Ijl_rW1cGIKrAMZ1CSswN_gGo,665
-project_llm_trainer-0.7.7.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.7.7.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.7.7.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.7.7.data/scripts/smart_train,sha256=3oLIDuuqb4U4TU1lXy9V8lw_0gIf7i8tGsxlQ_s6bro,1220
-project_llm_trainer-0.7.7.dist-info/METADATA,sha256=1O3xW3QM5aJgk1EESixX5DjxQ5ReX_pikUjI5x2qOvk,195
-project_llm_trainer-0.7.7.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.7.7.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.7.7.dist-info/RECORD,,
+project_llm_trainer-0.7.8.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.7.8.data/scripts/ddp_train,sha256=Z-309mM56CN0m3bxoeC5us4LUuwuNnoiOm3-fDdLMjQ,566
+project_llm_trainer-0.7.8.data/scripts/ds_train,sha256=3nXNNKmYI7miqyBdf-Ijl_rW1cGIKrAMZ1CSswN_gGo,665
+project_llm_trainer-0.7.8.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.7.8.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.7.8.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.7.8.data/scripts/smart_train,sha256=3oLIDuuqb4U4TU1lXy9V8lw_0gIf7i8tGsxlQ_s6bro,1220
+project_llm_trainer-0.7.8.dist-info/METADATA,sha256=rSYUrEkdjPCyYUqT2SOw3-hzT40wU3AwEw-ouHh1rBY,195
+project_llm_trainer-0.7.8.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.7.8.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.7.8.dist-info/RECORD,,