contrain dataset

scripts/prepare_contrain_dataset.py

@@ -1,11 +1,11 @@
-from typing import Optional, Union, Callable, Iterator, Any
+from typing import Optional, Union, Callable, Iterator
 from collections.abc import Collection
 from functools import partial
 
-import numpy as np
 from datasets import load_dataset
 from litdata import optimize, TokensLoader
 from litgpt.tokenizer import Tokenizer
+from litdata import StreamingDataset
 
 
 def batch_dict_iterator(path: Optional[str]=None,
@@ -65,8 +65,6 @@ def tokenize_fn(dataset_config: Union[dict, list], tokenizer: Optional[Tokenizer
  assert isinstance(dataset_config, (dict, list))
 
  for text in batch_iterator(dataset_config):
- # print(text)
- # break
  text_ids = tokenizer.encode(text, bos=False, eos=True)
  yield text_ids
 
@@ -717,7 +715,7 @@ datasets_configs = [
  {'role': 'user', 'content': r['prompt']},
  {'role': 'assistant', 'content': r['response']},
  ]}, # 12.4 MB, 3,000
- {'path': 'dvilasuero/dvilasuero/reflection-v1-final-dedup', 'transform': lambda r: [
+ {'path': 'dvilasuero/reflection-v1-final-dedup', 'transform': lambda r: [
  {'role': 'system', 'content': r['system']},
  {'role': 'user', 'content': r['prompt']},
  {'role': 'assistant', 'content': r['response']},
@@ -739,8 +737,17 @@ outputs = optimize(
  inputs=datasets_configs,
  output_dir='../contrain-data/',
  # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
- # chunk_size=(2049 * 8012),
  chunk_size=(8193 * 2003),
  num_workers=32,
  # compression='zstd',
 )
+
+#
+# total number of chunks
+#
+dataset = StreamingDataset(
+ input_dir='../contrain-data/',
+ item_loader=TokensLoader(block_size=8193),
+)
+
+print(len(dataset))

scripts/prepare_pretrain_dataset.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union
+from typing import Optional, Union, Iterator
 from functools import partial
 
 from datasets import load_dataset
@@ -15,7 +15,7 @@ def batch_dict_iterator(path: str,
  revision: Optional[str]=None,
  split: str='train',
  num_proc: Optional[int]=None,
- format: Optional[str]=None):
+ format: Optional[str]=None) -> Iterator[str]:
  assert isinstance(format, str) or callable(format)
 
  dataset = load_dataset(path=path,
@@ -86,12 +86,12 @@ datasets_configs = [
  # general knowledge
  #
  # 2.89 GB, 430,000, English September of 2017
- *[
+ [
  {'path': 'jordiclive/wikipedia-summary-dataset', 'split': f'train[{i}%:{i + 20}%]', 'format': lambda n: n['summary']}
  for i in range(0, 100, 20)
  ],
  # 3.18 GB, 1,010,500
- *[
+ [
  {'path': 'JeanKaddour/minipile', 'split': f'train[{i}%:{i + 20}%]', 'format': lambda n: n['text']}
  for i in range(0, 100, 20)
  ],