tangled-llama-v-128k-base-v0.1 / scripts /prepare_pretrain_dataset.py
mtasic85's picture
new tokenizer 38400
c62a845
raw
history blame
6.9 kB
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):
# text
if name in (None, 'saillab/taco-datasets'):
dataset = (
load_dataset(name, 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'] + '\n' + row['input'] + '\n' + row['output']
del dataset
gc.collect()
# text
if name in (None, 'xu-song/cc100-samples'):
dataset = (
load_dataset(name, 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 name in (None, 'ontocord/fineweb-permissive-multilingual-2m'):
dataset = load_dataset(name, split='train')
for row in dataset:
yield row['text']
del dataset
gc.collect()
# text
if name in (None, 'nampdn-ai/tiny-textbooks'):
for split in ['train', 'test']:
dataset = load_dataset(name, split=split)
for row in dataset:
yield row['textbook']
del dataset
gc.collect()
# code
if name in (None, 'bigcode/the-stack-smol-xs'):
dataset = (
load_dataset(name, 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()
# code
if name in (None, 'nampdn-ai/tiny-codes'):
dataset = load_dataset(name, split='train')
for row in dataset:
yield row['prompt'] + '\n' + row['response']
del dataset
gc.collect()
# text + code
if name in (None, 'm-a-p/CodeFeedback-Filtered-Instruction'):
dataset = load_dataset(name, split='train')
for row in dataset:
yield row['query'] + '\n' + row['answer']
del dataset
gc.collect()
# math
if name in (None, 'gair-prox/open-web-math-pro'):
dataset = load_dataset(name, split='train')
for row in dataset:
yield row['text']
del dataset
gc.collect()
# math
if name in (None, 'ajibawa-2023/Maths-College'):
dataset = load_dataset(name, 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(name, split='train')
for row in dataset:
yield row['question'] + '\n' + row['answer']
del dataset
gc.collect()
# math serbian
if name in (None, 'datatab/orca_math_world_problem_200k_serbian'):
dataset = load_dataset(name, split='train')
for row in dataset:
yield row['question_translated_srb'] + '\n' + row['answer_translated_srb']
del dataset
gc.collect()
# emoji
if name in (None, 'badrex/llm-emoji-dataset'):
dataset = load_dataset(name, 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 = [
'saillab/taco-datasets',
'xu-song/cc100-samples',
'ontocord/fineweb-permissive-multilingual-2m',
'nampdn-ai/tiny-textbooks',
'bigcode/the-stack-smol-xs',
'nampdn-ai/tiny-codes',
'm-a-p/CodeFeedback-Filtered-Instruction',
'gair-prox/open-web-math-pro',
'ajibawa-2023/Maths-College',
'microsoft/orca-math-word-problems-200k',
'datatab/orca_math_world_problem_200k_serbian',
'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=(4097 * 4006),
num_workers=16,
)