prepare dataset

.gitignore

@@ -161,4 +161,5 @@ cython_debug/
 .idea/
 
 .DS_Store
-.ruff_cache
+.ruff_cache
+venv*/

scripts/prepare_pretrain_dataset.py

@@ -0,0 +1,179 @@
+import gc
+
+from datasets import load_dataset
+from litdata import optimize, TokensLoader
+from litgpt.tokenizer import Tokenizer
+from functools import partial
+
+
+def batch_iterator(name=None):
+    # code
+    if name in (None, 'bigcode/programming-languages-keywords'):
+        dataset = load_dataset('bigcode/programming-languages-keywords', split='train')
+
+        for row in dataset:
+            for n in row['keywords']:
+                yield n
+
+        del dataset
+        gc.collect()
+
+    # code
+    if name in (None, 'bigcode/the-stack-smol-xs'):
+        dataset = (
+            load_dataset('bigcode/the-stack-smol-xs', lang, split='train', trust_remote_code=True)
+            for lang in [
+                'ada', 'agda', 'alloy', 'antlr', 'applescript', 'assembly', 'augeas', 'awk', 'batchfile', 'bison', 'bluespec', 'c',
+                'c++', 'c-sharp', 'clojure', 'cmake', 'coffeescript', 'common-lisp', 'css', 'cuda', 'dart', 'dockerfile', 'elixir',
+                'elm', 'emacs-lisp','erlang', 'f-sharp', 'fortran', 'glsl', 'go', 'groovy', 'haskell','html', 'idris', 'isabelle', 'java', 
+                'java-server-pages', 'javascript', 'julia', 'kotlin', 'lean', 'literate-agda', 'literate-coffeescript', 'literate-haskell',
+                'lua', 'makefile', 'maple', 'markdown', 'mathematica', 'matlab', 'ocaml', 'pascal', 'perl', 'php', 'powershell', 'prolog',
+                'protocol-buffer', 'python', 'r', 'racket', 'restructuredtext', 'rmarkdown', 'ruby', 'rust', 'sas', 'scala', 'scheme', 
+                'shell', 'smalltalk', 'solidity', 'sparql', 'sql', 'stan', 'standard-ml', 'stata', 'systemverilog', 'tcl', 'tcsh', 'tex', 
+                'thrift', 'typescript', 'verilog', 'vhdl', 'visual-basic', 'xslt', 'yacc', 'zig'
+            ]
+        )
+
+        for d in dataset:
+            for row in d:
+                yield row['content']
+
+        del dataset
+        gc.collect()
+
+    # text
+    if name in (None, 'nampdn-ai/tiny-textbooks'):
+        dataset = load_dataset('nampdn-ai/tiny-textbooks', split='train')
+
+        for row in dataset:
+            yield row['text']
+
+        del dataset
+        gc.collect()
+
+    # text
+    if name in (None, 'xu-song/cc100-samples'):
+        dataset = (
+            load_dataset('xu-song/cc100-samples', lang, split='train')
+            for lang in ['am', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'bn_rom', 'br', 'bs', 'ca', 'cs', 'cy', 'da', 'de', 'el', 'en', 'eo', 'es', 'et', 'eu', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl', 'gn', 'gu', 'ha', 'he', 'hi', 'hi_rom', 'hr', 'ht', 'hu', 'hy', 'id', 'ig', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km', 'kn', 'ko', 'ku', 'ky', 'la', 'lg', 'li', 'ln', 'lo', 'lt', 'lv', 'mg', 'mk', 'ml', 'mn', 'mr', 'ms', 'my', 'my_zaw', 'ne', 'nl', 'no', 'ns', 'om', 'or', 'pa', 'pl', 'ps', 'pt', 'qu', 'rm', 'ro', 'ru', 'sa', 'si', 'sc', 'sd', 'sk', 'sl', 'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'ta_rom', 'te', 'te_rom', 'th', 'tl', 'tn', 'tr', 'ug', 'uk', 'ur', 'ur_rom', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo', 'zh-Hans', 'zh-Hant', 'zu']
+        )
+
+        for d in dataset:
+            for row in d['text']:
+                yield row
+
+        del dataset
+        gc.collect()
+
+    # code
+    if name in (None, 'm-a-p/CodeFeedback-Filtered-Instruction'):
+        dataset = load_dataset('m-a-p/CodeFeedback-Filtered-Instruction', split='train')
+
+        for row in dataset:
+            yield row['query'] + '\n' + row['answer']
+
+        del dataset
+        gc.collect()
+
+    # code
+    if name in (None, 'nampdn-ai/tiny-codes'):
+        dataset = load_dataset('nampdn-ai/tiny-codes', split='train')
+        
+        for row in dataset:
+            yield row['prompt'] + '\n' + row['response']
+        
+        del dataset
+        gc.collect()
+
+    # math
+    if name in (None, 'ajibawa-2023/Maths-College'):
+        dataset = load_dataset('ajibawa-2023/Maths-College', split='train')
+        
+        for row in dataset:
+            yield row['instruction'] + '\n' + row['output']
+        
+        del dataset
+        gc.collect()
+
+    # math
+    if name in (None, 'microsoft/orca-math-word-problems-200k'):
+        dataset = load_dataset('microsoft/orca-math-word-problems-200k', split='train')
+
+        for row in dataset:
+            yield row['question'] + '\n' + row['answer']
+
+        del dataset
+        gc.collect()
+
+    # text
+    if name in (None, 'mlabonne/FineTome-100k'):
+        dataset = load_dataset('mlabonne/FineTome-100k', split='train')
+
+        for row in dataset['conversations']:
+            yield '\n'.join(n['value'] for n in row)
+
+        del dataset
+        gc.collect()
+
+    # instruction
+    if name in (None, 'arcee-ai/agent-data'):
+        dataset = load_dataset('arcee-ai/agent-data', split='train')
+        
+        for row in dataset['conversations']:
+            yield '\n'.join(n['value'] for n in row)
+        
+        del dataset
+        gc.collect()
+
+    # instruction
+    if name in (None, 'cognitivecomputations/SystemChat-2.0'):
+        dataset = (
+            load_dataset('cognitivecomputations/SystemChat-2.0', data_files='SystemChat_filtered.jsonl', split='train'),
+            load_dataset('cognitivecomputations/SystemChat-2.0', data_files='SystemChat_multilingual.jsonl', split='train'),
+        )
+        
+        for d in dataset:
+            for row in d['messages']:
+                yield '\n'.join(n['content'] for n in row)
+        
+        del dataset
+        gc.collect()
+
+    # emoji
+    if name in (None, 'badrex/llm-emoji-dataset'):
+        dataset = load_dataset('badrex/llm-emoji-dataset', split='train')
+        
+        for row in dataset:
+            yield f'{row["character"]}\n{row["unicode"]}\n{row["short description"]}\n{row["tags"]}\n{row["LLM description"]}'
+        
+        del dataset
+        gc.collect()
+
+
+def tokenize_fn(dataset_name, tokenizer=None):
+    for text in batch_iterator(dataset_name):
+        text_ids = tokenizer.encode(text, bos=False, eos=True)
+        yield text_ids
+
+datasets_names = [
+    'bigcode/programming-languages-keywords',
+    'bigcode/the-stack-smol-xs',
+    'nampdn-ai/tiny-textbooks',
+    'xu-song/cc100-samples',
+    'm-a-p/CodeFeedback-Filtered-Instruction',
+    'nampdn-ai/tiny-codes',
+    'ajibawa-2023/Maths-College',
+    'microsoft/orca-math-word-problems-200k',
+    'mlabonne/FineTome-100k',
+    'arcee-ai/agent-data',
+    'cognitivecomputations/SystemChat-2.0',
+    'badrex/llm-emoji-dataset',
+]
+
+outputs = optimize(
+    fn=partial(tokenize_fn, tokenizer=Tokenizer('..')),
+    inputs=datasets_names,
+    output_dir='../data/',
+    # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
+    chunk_size=(2049 * 8012),
+)

scripts/requirements-lit.in

@@ -0,0 +1,9 @@
+# pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+tqdm
+datasets
+jinja2
+transformers
+bitsandbytes
+wandb
+litgpt[all]
+litdata

scripts/requirements.in

@@ -2,6 +2,5 @@ tqdm
 datasets
 jinja2
 transformers
-jsonlines
 bitsandbytes
 wandb