scripts/prepare_contrain_0_lang_math_dataset.py

from typing import Optional, Union
from functools import partial

import numpy as np
from datasets import load_dataset
from litdata import optimize, TokensLoader
from litgpt.tokenizer import Tokenizer


def batch_dict_iterator(path: str,
 name: Optional[str]=None,
 data_dir: Optional[str]=None,
 data_files: Optional[str]=None,
 keep_in_memory: bool=False,
 revision: Optional[str]=None,
 split: str='train',
 num_proc: Optional[int]=None,
 format: Optional[str]=None):
 assert isinstance(format, str) or callable(format)
 
 dataset = load_dataset(path=path,
 name=name,
 data_dir=data_dir,
 data_files=data_files,
 keep_in_memory=keep_in_memory,
 revision=revision,
 split=split,
 trust_remote_code=True,
 num_proc=num_proc)

 if callable(format):
 for row in dataset:
 text = format(row)
 yield text
 else:
 for row in dataset:
 text = format.format(**row)
 yield text


def batch_iterator(dataset_config: Union[list, dict]):
 if isinstance(dataset_config, dict):
 for text in batch_dict_iterator(**dataset_config):
 yield text
 elif isinstance(dataset_config, list):
 for dc in dataset_config:
 for text in batch_dict_iterator(**dc):
 yield text
 else:
 raise ValueError('')


def tokenize_fn(dataset_config: Union[dict, list], tokenizer: Optional[Tokenizer]=None):
 assert isinstance(dataset_config, (dict, list))

 for text in batch_iterator(dataset_config):
 text_ids = tokenizer.encode(text, bos=False, eos=True)
 yield text_ids


datasets_configs = [
 #
 # multilingual instruct
 #
 {'path': 'yahma/alpaca-cleaned', 'format': '{instruction} {input} {output}'}, # 44.3 MB, 51,760
 # saillab/taco-datasets 2.48 GB, 3,202,163
 [
 {'path': 'saillab/taco-datasets', 'data_dir': data_dir, 'split': 'train', 'format': '{instruction} {input} {output}'}
 for data_dir in [
 f'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4/{n}'
 for n in [
 'Afrikaans', 'Albanian', 'Amharic', 'Arabic', 'Armenian', 'Assamese',
 'Aymara', 'Azerbaijani', 'Bambara', 'Basque', 'Belarusian', 'Bengali',
 'Bhojpuri', 'Bosnian', 'Bulgarian', 'Catalan', 'Cebuano', 'Chichewa',
 'ChineseSimplified', 'ChineseTraditional', 'Corsican', 'Croatian',
 'Czech', 'Danish', 'Divehi', 'Dogri', 'Dutch', 'Esperanto', 'Estonian',
 'Ewe', 'Filipino', 'Finnish', 'French', 'Frisian', 'Galician',
 'Georgian', 'German', 'Greek', 'Guarani', 'Gujarati', 'Haitian_Creole',
 'Hausa', 'Hawaiian', 'Hebrew', 'Hindi', 'Hmong', 'Hungarian',
 'Icelandic', 'Igbo', 'Ilocano', 'Indonesian', 'Irish', 'Italian',
 'Japanese', 'Javanese', 'Kannada', 'Kazakh', 'Khmer', 'Kinyarwanda',
 'Konkani', 'Korean', 'Krio', 'Kurdish_Kurmanji', 'Kurdish_Sorani',
 'Kyrgyz', 'Lao', 'Latin', 'Latvian', 'Lingala', 'Lithuanian',
 'Luganda', 'Luxembourgish', 'Macedonian', 'Maithili', 'Malagasy',
 'Malay', 'Malayalam', 'Maltese', 'Maori', 'Marathi', 'Meiteilon_Manipuri',
 'Mizo', 'Mongolian', 'Myanmar_Burmese', 'Nepali', 'Norwegian',
 'Odia_Oriya', 'Oromo', 'Pashto', 'Persian', 'Polish', 'Portuguese',
 'Punjabi', 'Quechua', 'Romanian', 'Russian', 'Samoan', 'Sanskrit',
 'ScottishGaelic', 'Sepedi', 'Serbian', 'Sesotho', 'Shona', 'Sindhi',
 'Sinhala', 'Slovak', 'Slovenian', 'Somali', 'Spanish', 'Sundanese',
 'Swahili', 'Swedish', 'Tajik', 'Tamil', 'Tatar', 'Telugu', 'Thai',
 'Tigrinya', 'Tsonga', 'Turkish', 'Turkmen', 'Twi', 'Ukrainian',
 'Urdu', 'Uyghur', 'Uzbek', 'Vietnamese', 'Welsh', 'Xhosa',
 'Yiddish', 'Yoruba', 'Zulu',
 ]
 ]
 ],
 [
 {'path': 'saillab/taco-datasets', 'data_dir': 'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k/', 'data_files': n, 'split': 'train', 'format': '{instruction} {input} {output}'}
 for n in [
 'Afrikaans.json', 'Albanian.json', 'Amharic.json', 'Arabic.json', 'Armenian.json',
 'Assamese.json', 'Aymara.json', 'Azerbaijani.json', 'Bambara.json', 'Basque.json',
 'Belarusian.json', 'Bengali.json', 'Bhojpuri.json', 'Bosnian.json', 'Bulgarian.json',
 'Catalan.json', 'Cebuano.json', 'Chichewa.json', 'ChineseSimplified.json',
 'ChineseTraditional.json', 'Corsican.json', 'Croatian.json', 'Czech.json',
 'Danish.json', 'Dhivehi.json', 'Dogri.json', 'Dutch.json', 'English.json',
 'Esperanto.json', 'Estonian.json', 'Ewe.json', 'Filipino.json',
 'Finnish.json', 'French.json', 'Frisian.json', 'Galician.json',
 'Georgian.json', 'German.json', 'Greek.json', 'Guarani.json',
 'Gujarati.json', 'Haitian_Creole.json', 'Hausa.json', 'Hawaiian.json',
 'Hebrew.json', 'Hindi.json', 'Hmong.json', 'Hungarian.json',
 'Icelandic.json', 'Igbo.json', 'Ilocano.json', 'Indonesian.json',
 'Irish.json', 'Italian.json', 'Japanese.json', 'Javanese.json',
 'Kannada.json', 'Kazakh.json', 'Khmer.json', 'Kinyarwanda.json',
 'Konkani.json', 'Korean.json', 'Krio.json', 'Kurdish_Kurmanji.json',
 'Kurdish_Sorani.json', 'Kyrgyz.json', 'Lao.json', 'Latin.json',
 'Latvian.json', 'Lingala.json', 'Lithuanian.json', 'Luganda.json',
 'Luxembourgish.json', 'Macedonian.json', 'Maithili.json',
 'Malagasy.json', 'Malayalam.json', 'Malay.json', 'Maltese.json',
 'Maori.json', 'Marathi.json', 'Meiteilon_Manipuri.json',
 'Mizo.json', 'Mongolian.json', 'Myanmar_Burmese.json',
 'Nepali.json', 'Norwegian.json', 'Odia_Oriya.json', 'Oromo.json',
 'Pashto.json', 'Persian.json', 'Polish.json', 'Portuguese.json',
 'Punjabi.json', 'Quechua.json', 'Romanian.json', 'Russian.json',
 'Samoan.json', 'Sanskrit.json', 'ScottishGaelic.json', 'Sepedi.json',
 'Serbian.json', 'Sesotho.json', 'Shona.json', 'Sindhi.json',
 'Sinhala.json', 'Slovak.json', 'Slovenian.json', 'Somali.json',
 'Spanish.json', 'Sundanese.json', 'Swahili.json', 'Swedish.json',
 'Tajik.json', 'Tamil.json', 'Tatar.json', 'Telugu.json', 'Thai.json',
 'Tigrinya.json', 'Tsonga.json', 'Turkish.json', 'Turkmen.json',
 'Twi.json', 'Ukrainian.json', 'Urdu.json', 'Uyghur.json', 'Uzbek.json',
 'Vietnamese.json', 'Welsh.json', 'Xhosa.json', 'Yiddish.json',
 'Yoruba.json', 'Zulu.json',
 ]
 ],
 [
 # 193 MB, 1,141,967
 {'path': 'xu-song/cc100-samples', 'name': name, 'split': 'train', 'format': lambda n: n['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',
 ]
 ],

 #
 # general knowledge
 #
 # 2.89 GB, 430,000, English September of 2017
 *[
 {'path': 'jordiclive/wikipedia-summary-dataset', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['summary']}
 for i in range(0, 100, 5)
 ],

 #
 # math
 #
 # 9.05 GB, 2,583,257
 *[
 {'path': 'gair-prox/open-web-math-pro', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
 for i in range(0, 100, 5)
 ]
]

outputs = optimize(
 fn=partial(tokenize_fn, tokenizer=Tokenizer('..')),
 inputs=datasets_configs,
 output_dir='../contrain-lang-math-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,
)

#
# total number of chunks
#
from litdata import StreamingDataset, StreamingDataLoader, TokensLoader

dataset = StreamingDataset(
 input_dir='../contrain-lang-math-data/',
 item_loader=TokensLoader(block_size=8193),
)

print(len(dataset))