prepare pretrain datasets

scripts/pretrain_datasets.py

@@ -4,13 +4,13 @@ pretrain_datasets = [
     #
     # 3.17 GB, 2,226,907
     *[
-        {'path': 'ontocord/fineweb-permissive-multilingual-2m', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
-        for i in range(0, 100, 5)
+        {'path': 'ontocord/fineweb-permissive-multilingual-2m', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 10)
     ],
     # 1.64 GB, 1,001,000
     *[
-        {'path': 'distily/c4_multilingual_1M', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
-        for i in range(0, 100, 5)
+        {'path': 'distily/c4_multilingual_1M', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 10)
     ],
 
     #
@@ -32,9 +32,14 @@ pretrain_datasets = [
     #
     # math
     #
-    # 12.6 GB, 21,972,791 - we use 1M subset - 639 MB, 1,000,000
+    # # 12.6 GB, 21,972,791 - we use 1M subset - 639 MB, 1,000,000
+    # *[
+    #     {'path': 'nvidia/OpenMathInstruct-2', 'split': f'train_1M[{i}%:{i + 10}%]', 'format': '{problem} {generated_solution} {expected_answer}'}
+    #     for i in range(0, 100, 10)
+    # ],
+    # 12.6 GB, 14M rows
     *[
-        {'path': 'nvidia/OpenMathInstruct-2', 'split': f'train_1M[{i}%:{i + 5}%]', 'format': '{problem} {generated_solution} {expected_answer}'}
+        {'path': 'nvidia/OpenMathInstruct-2', 'split': f'train[{i}%:{i + 5}%]', 'format': '{problem} {generated_solution} {expected_answer}'}
         for i in range(0, 100, 5)
     ],
 
@@ -43,12 +48,12 @@ pretrain_datasets = [
     #
     # 1.44 GB, 63,357
     *[
-        {'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['abstract']}
-        for i in range(0, 100, 5)
+        {'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['abstract']}
+        for i in range(0, 100, 10)
     ],
     *[
-        {'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['markdown']}
-        for i in range(0, 100, 5)
+        {'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['markdown']}
+        for i in range(0, 100, 10)
     ],
 
     #
@@ -56,8 +61,8 @@ pretrain_datasets = [
     #
     # 7.81 GB, ~2,804,025
     *[
-        {'path': 'rombodawg/code_bagel_hermes-2.5', 'split': f'train[{i}%:{i + 5}%]', 'format': '{input} {output}'}
-        for i in range(0, 100, 5)
+        {'path': 'rombodawg/code_bagel_hermes-2.5', 'split': f'train[{i}%:{i + 10}%]', 'format': '{input} {output}'}
+        for i in range(0, 100, 10)
     ],
 
     #
@@ -65,8 +70,8 @@ pretrain_datasets = [
     #
     # 3.18 GB, 1,010,500 - paper says that extracted is 6GB
     *[
-        {'path': 'JeanKaddour/minipile', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
-        for i in range(0, 100, 5)
+        {'path': 'JeanKaddour/minipile', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 10)
     ],
     {'path': 'JeanKaddour/minipile', 'split': 'validation', 'format': lambda n: n['text']},
     {'path': 'JeanKaddour/minipile', 'split': 'test', 'format': lambda n: n['text']},

scripts/utils.py

@@ -6,7 +6,7 @@ from datasets import load_dataset
 from litgpt.tokenizer import Tokenizer
 from transformers import AutoTokenizer
 
-# def _batch_text_iterator(path: str,
+
 def batch_text_iterator(path: str,
                         name: Optional[str]=None,
                         data_dir: Optional[str]=None,
@@ -41,7 +41,6 @@ def batch_text_iterator(path: str,
     gc.collect()
 
 
-# def _batch_chat_iterator(path: str,
 def batch_chat_iterator(path: str,
                         name: Optional[str]=None,
                         data_dir: Optional[str]=None,
@@ -84,66 +83,14 @@ def batch_chat_iterator(path: str,
     gc.collect()
 
 
-# def batch_text_iterator(dataset_config: Union[list, dict]) -> Iterator[str]:
-#     assert isinstance(dataset_config, (dict, list)), dataset_config
-#
-#     if isinstance(dataset_config, dict):
-#         for text in _batch_text_iterator(**dataset_config):
-#             yield text
-#     elif isinstance(dataset_config, list):
-#         for dc in dataset_config:
-#             for text in _batch_text_iterator(**dc):
-#                 yield text
-
-
-# def batch_chat_iterator(dataset_config: Union[list, dict]) -> Iterator[list[dict[str, str]]]:
-#     assert isinstance(dataset_config, (dict, list)), dataset_config
-#
-#     if isinstance(dataset_config, dict):
-#         for messages in _batch_chat_iterator(**dataset_config):
-#             yield messages
-#     elif isinstance(dataset_config, list):
-#         for dc in dataset_config:
-#             for messages in _batch_chat_iterator(**dc):
-#                 yield messages
-
-
-# def tokenize_text_fn(dataset_config: list, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer, min_len: Optional[int]=None, max_len: Optional[int]=None) -> Iterator[torch.Tensor]:
 def tokenize_text_fn(dataset_config: dict, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer) -> Iterator[torch.Tensor]:
     for text in batch_text_iterator(**dataset_config):
         text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=True)
-
-        # if min_len is None and max_len is None:
-        #     yield text_ids
-
-        # if min_len is None:
-        #     min_len = 0
-
-        # if max_len is None:
-        #     max_len = len(text_ids)
-
-        # if min_len <= len(text_ids) <= max_len:
-        #     yield text_ids
-
         yield text_ids
 
 
-# def tokenize_chat_fn(dataset_config: list, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer, min_len: Optional[int]=None, max_len: Optional[int]=None) -> Iterator[torch.Tensor]:
 def tokenize_chat_fn(dataset_config: dict, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer) -> Iterator[torch.Tensor]:
     for messages in batch_chat_iterator(**dataset_config):
         text: str = hf_tokenizer.apply_chat_template(messages, tokenize=False)
         text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=False)
-
-        # if min_len is None and max_len is None:
-        #     yield text_ids
-
-        # if min_len is None:
-        #     min_len = 0
-
-        # if max_len is None:
-        #     max_len = len(text_ids)
-
-        # if min_len <= len(text_ids) <= max_len:
-        #     yield text_ids
-
         yield text_ids