|
import gc |
|
from typing import Optional |
|
|
|
from datasets import load_dataset |
|
from litdata import optimize, TokensLoader |
|
from litgpt.tokenizer import Tokenizer |
|
from functools import partial |
|
|
|
|
|
""" |
|
def batch_iterator_1(path=None): |
|
# text |
|
if path in (None, 'saillab/taco-datasets'): |
|
dataset = ( |
|
load_dataset(path, data_dir=data_dir, split='train') |
|
for data_dir in [ |
|
'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4', |
|
'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k', |
|
] |
|
) |
|
|
|
for d in dataset: |
|
for row in d: |
|
for n in row: |
|
yield ( |
|
row['instruction'] + |
|
' ' + |
|
row['input'] + |
|
' ' + |
|
row['output'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# text |
|
if path in (None, 'xu-song/cc100-samples'): |
|
dataset = ( |
|
load_dataset(path, 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() |
|
|
|
# text |
|
if path in (None, 'ontocord/fineweb-permissive-multilingual-2m'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset['text']: |
|
yield row |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# text |
|
if path in (None, 'MuskumPillerum/General-Knowledge'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset: |
|
if not row['Question'] or not row['Answer']: |
|
continue |
|
|
|
yield row['Question'] + ' ' + row['Answer'] |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# text |
|
if path in (None, 'yirenc/general_knowledge_boolean'): |
|
for split in ['train', 'validation']: |
|
dataset = load_dataset(path, split=split) |
|
|
|
for row in dataset: |
|
yield row['question'] + '? ' + str(row['answer']) + '. ' + row['passage'] |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# text |
|
if path in (None, 'nampdn-ai/tiny-textbooks'): |
|
for split in ['train', 'test']: |
|
dataset = load_dataset(path, split=split) |
|
|
|
for row in dataset['textbook']: |
|
yield row |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# code |
|
if path in (None, 'nampdn-ai/tiny-codes'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset: |
|
yield ( |
|
row['prompt'] + |
|
' ' + |
|
row['response'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# code |
|
if path in (None, 'bigcode/the-stack-smol-xs'): |
|
dataset = ( |
|
load_dataset(path, 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['content']: |
|
yield row |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# text + code |
|
if path in (None, 'm-a-p/CodeFeedback-Filtered-Instruction'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset: |
|
yield ( |
|
row['query'] + |
|
' ' + |
|
row['answer'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# code |
|
if path in (None, 'jtatman/python-code-dataset-500k'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset: |
|
yield ( |
|
row['instruction'] + |
|
' ' + |
|
row['output'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# code |
|
if path in (None, 'iamtarun/python_code_instructions_18k_alpaca'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset: |
|
yield ( |
|
row['instruction'] + |
|
' ' + |
|
row['input'] + |
|
' ' + |
|
row['output'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# code |
|
if path in (None, 'HuggingFaceH4/CodeAlpaca_20K'): |
|
for split in ['train', 'test']: |
|
dataset = load_dataset(path, split=split) |
|
|
|
for row in dataset: |
|
yield ( |
|
row['prompt'] + |
|
' ' + |
|
row['completion'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# math |
|
if path in (None, 'gair-prox/open-web-math-pro'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset['text']: |
|
yield row |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# math |
|
if path in (None, 'rvv-karma/Math-QA'): |
|
for split in ['train', 'val', 'test']: |
|
dataset = load_dataset(path, split=split) |
|
|
|
for row in dataset: |
|
yield ( |
|
row['question'] + |
|
' ' + |
|
row['answer'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# math |
|
if path in (None, 'ajibawa-2023/Maths-College'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset: |
|
yield ( |
|
row['instruction'] + |
|
' ' + |
|
row['output'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# math |
|
if path in (None, 'microsoft/orca-math-word-problems-200k'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset: |
|
yield ( |
|
row['question'] + |
|
' ' + |
|
row['answer'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# math |
|
if path in (None, 'fblgit/simple-math'): |
|
for split in ['train', 'test']: |
|
dataset = load_dataset(path, revision='refs/convert/parquet', split=split) |
|
|
|
for row in dataset: |
|
yield ( |
|
str(row['instruction']) + |
|
' = ' + |
|
str(row['output']) |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# reasoning |
|
if path in (None, 'SkunkworksAI/reasoning-0.01'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset: |
|
yield ( |
|
row['instruction'] + |
|
' ' + |
|
row['reasoning'] + |
|
' ' + |
|
row['output'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
|
|
# emoji |
|
if path in (None, 'badrex/llm-emoji-dataset'): |
|
dataset = load_dataset(path, split='train') |
|
|
|
for row in dataset: |
|
yield ( |
|
row['character'] + |
|
' ' + |
|
row['unicode'] + |
|
' ' + |
|
row['short description'] + |
|
' ' + |
|
str(row['tags']) + |
|
' ' + |
|
row['LLM description'] |
|
) |
|
|
|
del dataset |
|
gc.collect() |
|
""" |
|
|
|
def batch_iterator(path: str, |
|
name: Optional[str]=None, |
|
data_dir: Optional[str]=None, |
|
data_files: Optional[str]=None, |
|
revision: Optional[str]=None, |
|
split: str='train', |
|
format: Optional[str]=None): |
|
assert format is not None |
|
|
|
dataset = load_dataset(path=path, |
|
name=name, |
|
data_dir=data_dir, |
|
data_files=data_files, |
|
revision=revision, |
|
split=split, |
|
trust_remote_code=True) |
|
|
|
for row in dataset: |
|
text = format.format(**row) |
|
yield text |
|
|
|
|
|
def tokenize_fn(datasets_config, tokenizer=None): |
|
for text in batch_iterator(**datasets_config): |
|
text_ids = tokenizer.encode(text, bos=False, eos=True) |
|
yield text_ids |
|
|
|
""" |
|
datasets_names = [ |
|
'saillab/taco-datasets', |
|
# 'xu-song/cc100-samples', |
|
# 'ontocord/fineweb-permissive-multilingual-2m', |
|
'MuskumPillerum/General-Knowledge', |
|
'yirenc/general_knowledge_boolean', |
|
'nampdn-ai/tiny-textbooks', |
|
# 'nampdn-ai/tiny-codes', |
|
'bigcode/the-stack-smol-xs', |
|
'm-a-p/CodeFeedback-Filtered-Instruction', |
|
# 'jtatman/python-code-dataset-500k', |
|
'iamtarun/python_code_instructions_18k_alpaca', |
|
'HuggingFaceH4/CodeAlpaca_20K', |
|
# 'gair-prox/open-web-math-pro', |
|
'rvv-karma/Math-QA', |
|
# 'ajibawa-2023/Maths-College', |
|
'microsoft/orca-math-word-problems-200k', |
|
'fblgit/simple-math', |
|
# 'SkunkworksAI/reasoning-0.01', |
|
'badrex/llm-emoji-dataset', |
|
] |
|
""" |
|
|
|
datasets_configs = [ |
|
{'path': 'yahma/alpaca-cleaned', 'format': '{instruction} {input} {output}'}, |
|
{'path': 'gbharti/wealth-alpaca_lora', 'format': '{instruction} {input} {output}'}, |
|
*[ |
|
{'path': 'saillab/taco-datasets', 'data_dir': data_dir, 'split': 'train[:10%]', 'format': '{instruction} {input} {output}'} |
|
for data_dir in [ |
|
'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4', |
|
'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k', |
|
] |
|
], |
|
*[ |
|
{'path': 'xu-song/cc100-samples', 'name': name, 'split': 'train[:10%]', 'format': '{text}'} |
|
for name 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', |
|
] |
|
], |
|
{'path': 'ontocord/fineweb-permissive-multilingual-2m', 'split': 'train[:5%]', 'format': '{text}'}, |
|
{'path': 'MuskumPillerum/General-Knowledge', 'format': '{Question} {Answer}'}, |
|
{'path': 'yirenc/general_knowledge_boolean', 'split': 'train+validation', 'format': '{question}? {answer}. {passage}'}, |
|
{'path': 'nampdn-ai/tiny-textbooks', 'split': 'train+test', 'format': '{textbook}'}, |
|
{'path': 'nampdn-ai/tiny-codes', 'split': 'train[:5%]', 'format': '{prompt} {response}'}, |
|
*[ |
|
{'path': 'bigcode/the-stack-smol-xs', 'name': name, 'format': '{content}'} |
|
for name 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', |
|
] |
|
], |
|
{'path': 'm-a-p/CodeFeedback-Filtered-Instruction', 'split': 'train', 'format': '{query} {answer}'}, |
|
{'path': 'jtatman/python-code-dataset-500k', 'format': '{instruction} {output}'}, |
|
{'path': 'iamtarun/python_code_instructions_18k_alpaca', 'format': '{instruction} {input} {output}'}, |
|
{'path': 'HuggingFaceH4/CodeAlpaca_20K', 'split': 'train+test', 'format': '{prompt} {completion}'}, |
|
{'path': 'gair-prox/open-web-math-pro', 'split': 'train[:5%]', 'format': '{text}'}, |
|
{'path': 'rvv-karma/Math-QA', 'split': 'train+val+test', 'format': '{question} {answer}'}, |
|
{'path': 'ajibawa-2023/Maths-College', 'split': 'train[:10%]', 'format': '{instruction} {output}'}, |
|
{'path': 'microsoft/orca-math-word-problems-200k', 'format': '{question} {answer}'}, |
|
{'path': 'fblgit/simple-math', 'revision': 'refs/convert/parquet', 'split': 'train+test', 'format': '{instruction} = {output}'}, |
|
{'path': 'SkunkworksAI/reasoning-0.01', 'format': '{instruction} {reasoning} {output}'}, |
|
{'path': 'badrex/llm-emoji-dataset', 'format': '{character} {unicode} {short description} {tags} {LLM description}'}, |
|
] |
|
|
|
outputs = optimize( |
|
fn=partial(tokenize_fn, tokenizer=Tokenizer('..')), |
|
inputs=datasets_configs, |
|
output_dir='../pretrain-data/', |
|
|
|
chunk_size=(2049 * 8012), |
|
num_workers=32, |
|
) |
|
|
