278 lines
11 KiB
Python
278 lines
11 KiB
Python
|
from __future__ import annotations
|
||
|
|
import random
|
||
|
|
import time
|
||
|
|
from decon import is_contaminated
|
||
|
|
STAGE_BOUNDARIES = [0.7, 0.9, 1.0]
|
||
|
|
STAGE_MIX = [{'web': 0.74, 'math': 0.09, 'code': 0.07, 'arithmetic': 0.05, 'dialogue': 0.05}, {'web': 0.59, 'math': 0.18, 'code': 0.12, 'arithmetic': 0.06, 'dialogue': 0.05}, {'web': 0.5, 'math': 0.24, 'code': 0.11, 'arithmetic': 0.09, 'dialogue': 0.06}]
|
||
|
|
WEB_MIX = {'fineweb_edu': 0.6, 'dclm': 0.4}
|
||
|
|
MIX_PRESETS = {'smollm2': {'fineweb_edu': 0.576, 'dclm': 0.384, 'cosmopedia': 0.04}, 'dclm_heavy': {'fineweb_edu': 0.384, 'dclm': 0.576, 'cosmopedia': 0.04}, 'smollm2_math': {'fineweb_edu': 0.55, 'dclm': 0.37, 'cosmopedia': 0.04, 'math': 0.04}, 'dclm_heavy_math': {'fineweb_edu': 0.37, 'dclm': 0.55, 'cosmopedia': 0.04, 'math': 0.04}, 'cooldown': {'fineweb_edu': 0.75, 'cosmopedia': 0.25}}
|
||
|
|
MUSEKO_MIX = {'web': 0.9, 'arithmetic': 0.1}
|
||
|
|
EOS = '<|endoftext|>'
|
||
|
|
|
||
|
|
def get_stage_mix(progress: float, mix_mode: str='default', arith_frac: float | None=None) -> dict:
|
||
|
|
if mix_mode == 'museko':
|
||
|
|
mix = dict(MUSEKO_MIX)
|
||
|
|
else:
|
||
|
|
mix = None
|
||
|
|
for boundary, stage in zip(STAGE_BOUNDARIES, STAGE_MIX):
|
||
|
|
if progress <= boundary:
|
||
|
|
mix = dict(stage)
|
||
|
|
break
|
||
|
|
if mix is None:
|
||
|
|
mix = dict(STAGE_MIX[-1])
|
||
|
|
if arith_frac is None:
|
||
|
|
return mix
|
||
|
|
other = {k: v for k, v in mix.items() if k != 'arithmetic'}
|
||
|
|
total_other = sum(other.values()) or 1.0
|
||
|
|
scale = (1.0 - arith_frac) / total_other
|
||
|
|
out = {k: v * scale for k, v in other.items()}
|
||
|
|
out['arithmetic'] = arith_frac
|
||
|
|
return out
|
||
|
|
|
||
|
|
def _cycling_stream(load_fn, name: str):
|
||
|
|
backoff = 5.0
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
it = iter(load_fn())
|
||
|
|
backoff = 5.0
|
||
|
|
except Exception as e:
|
||
|
|
msg = str(e)[:140]
|
||
|
|
print(f'[data] stream {name} failed to start, retry in {backoff:.0f}s: {msg}', flush=True)
|
||
|
|
time.sleep(backoff)
|
||
|
|
backoff = min(backoff * 2, 300.0)
|
||
|
|
continue
|
||
|
|
while True:
|
||
|
|
try:
|
||
|
|
row = next(it)
|
||
|
|
except StopIteration:
|
||
|
|
break
|
||
|
|
except Exception as e:
|
||
|
|
print(f'[data] stream {name} row error, skipping: {e}', flush=True)
|
||
|
|
continue
|
||
|
|
yield row
|
||
|
|
|
||
|
|
def _split(ds, rank: int, world_size: int, name: str=''):
|
||
|
|
if world_size <= 1:
|
||
|
|
return ds
|
||
|
|
n = None
|
||
|
|
try:
|
||
|
|
n = ds.n_shards
|
||
|
|
except Exception:
|
||
|
|
pass
|
||
|
|
if n is not None and n >= world_size:
|
||
|
|
return ds.shard(num_shards=world_size, index=rank)
|
||
|
|
print(f'[data] {name}: {n} files < {world_size} slices, striping rows instead', flush=True)
|
||
|
|
return _RowStripe(ds, rank, world_size)
|
||
|
|
|
||
|
|
class _RowStripe:
|
||
|
|
|
||
|
|
def __init__(self, ds, rank, world_size):
|
||
|
|
self.ds, self.rank, self.world_size = (ds, rank, world_size)
|
||
|
|
|
||
|
|
def __iter__(self):
|
||
|
|
for i, row in enumerate(self.ds):
|
||
|
|
if i % self.world_size == self.rank:
|
||
|
|
yield row
|
||
|
|
|
||
|
|
def _shard(ds, rank: int, world_size: int):
|
||
|
|
if world_size > 1:
|
||
|
|
return ds.shard(num_shards=world_size, index=rank)
|
||
|
|
return ds
|
||
|
|
|
||
|
|
def _rank_filter(gen, rank: int, world_size: int):
|
||
|
|
if world_size <= 1:
|
||
|
|
yield from gen
|
||
|
|
return
|
||
|
|
for i, item in enumerate(gen):
|
||
|
|
if i % world_size == rank:
|
||
|
|
yield item
|
||
|
|
|
||
|
|
def _fineweb_edu_texts(rank: int=0, world_size: int=1):
|
||
|
|
from datasets import load_dataset
|
||
|
|
|
||
|
|
def _load():
|
||
|
|
ds = load_dataset('HuggingFaceFW/fineweb-edu', 'sample-100BT', split='train', streaming=True)
|
||
|
|
return _split(ds, rank, world_size, 'fineweb-edu')
|
||
|
|
for row in _cycling_stream(_load, 'fineweb-edu'):
|
||
|
|
yield row['text']
|
||
|
|
|
||
|
|
def _dclm_texts(rank: int=0, world_size: int=1):
|
||
|
|
from datasets import load_dataset
|
||
|
|
|
||
|
|
def _load():
|
||
|
|
ds = load_dataset('mlfoundations/dclm-baseline-1.0', split='train', streaming=True)
|
||
|
|
return _split(ds, rank, world_size, 'dclm')
|
||
|
|
for row in _cycling_stream(_load, 'dclm'):
|
||
|
|
yield row['text']
|
||
|
|
|
||
|
|
def _cosmopedia_texts(rank: int=0, world_size: int=1):
|
||
|
|
from datasets import load_dataset
|
||
|
|
|
||
|
|
def _load():
|
||
|
|
ds = load_dataset('HuggingFaceTB/smollm-corpus', 'cosmopedia-v2', split='train', streaming=True)
|
||
|
|
return _split(ds, rank, world_size, 'cosmopedia-v2')
|
||
|
|
for row in _cycling_stream(_load, 'cosmopedia-v2'):
|
||
|
|
t = row.get('text') or ''
|
||
|
|
if t:
|
||
|
|
yield t
|
||
|
|
|
||
|
|
def _finemath_texts(rank: int=0, world_size: int=1):
|
||
|
|
from datasets import load_dataset
|
||
|
|
|
||
|
|
def _load():
|
||
|
|
return load_dataset('HuggingFaceTB/finemath', 'finemath-4plus', split='train', streaming=True)
|
||
|
|
gen = (row['text'] for row in _cycling_stream(_load, 'finemath'))
|
||
|
|
yield from _rank_filter(gen, rank, world_size)
|
||
|
|
|
||
|
|
def _code_texts(rank: int=0, world_size: int=1):
|
||
|
|
from datasets import load_dataset
|
||
|
|
|
||
|
|
def _load():
|
||
|
|
return load_dataset('semeru/Code-Text-Python', split='train', streaming=True)
|
||
|
|
|
||
|
|
def _texts():
|
||
|
|
for row in _cycling_stream(_load, 'code-text-python'):
|
||
|
|
docstring = row.get('docstring') or ''
|
||
|
|
code = row.get('code') or row.get('original_string') or ''
|
||
|
|
if not code:
|
||
|
|
continue
|
||
|
|
yield (f'"""{docstring}"""\n{code}' if docstring else code)
|
||
|
|
yield from _rank_filter(_texts(), rank, world_size)
|
||
|
|
|
||
|
|
def _smoltalk_texts(rank: int=0, world_size: int=1):
|
||
|
|
from datasets import load_dataset
|
||
|
|
|
||
|
|
def _load():
|
||
|
|
return load_dataset('HuggingFaceTB/smol-smoltalk', split='train', streaming=True)
|
||
|
|
|
||
|
|
def _texts():
|
||
|
|
for row in _cycling_stream(_load, 'smol-smoltalk'):
|
||
|
|
parts = []
|
||
|
|
for msg in row['messages']:
|
||
|
|
role = msg.get('role', 'user')
|
||
|
|
parts.append(f"<|{role}|>\n{msg.get('content', '')}")
|
||
|
|
yield '\n'.join(parts)
|
||
|
|
yield from _rank_filter(_texts(), rank, world_size)
|
||
|
|
|
||
|
|
def _realmath_texts(rank: int=0, world_size: int=1):
|
||
|
|
import re as _re
|
||
|
|
from datasets import load_dataset
|
||
|
|
|
||
|
|
def _orca():
|
||
|
|
return load_dataset('microsoft/orca-math-word-problems-200k', split='train', streaming=True)
|
||
|
|
|
||
|
|
def _gsm8k():
|
||
|
|
return load_dataset('openai/gsm8k', 'main', split='train', streaming=True)
|
||
|
|
calc = _re.compile('<<[^>]*>>')
|
||
|
|
|
||
|
|
def _texts():
|
||
|
|
orca = _cycling_stream(_orca, 'orca-math')
|
||
|
|
gsm = _cycling_stream(_gsm8k, 'gsm8k')
|
||
|
|
while True:
|
||
|
|
row = next(orca) if random.random() < 0.85 else next(gsm)
|
||
|
|
q = (row.get('question') or '').strip()
|
||
|
|
a = calc.sub('', row.get('answer') or '').strip()
|
||
|
|
if not q or not a:
|
||
|
|
continue
|
||
|
|
yield f'{q}\n{a}'
|
||
|
|
yield from _rank_filter(_texts(), rank, world_size)
|
||
|
|
|
||
|
|
def _arithmetic_texts(rank: int=0, world_size: int=1, source: str='wordmath'):
|
||
|
|
import random as _random
|
||
|
|
rng = _random.Random(1000 + rank)
|
||
|
|
if source == 'drill':
|
||
|
|
from gen_arithmetic import gen_example as _gen
|
||
|
|
while True:
|
||
|
|
yield _gen()
|
||
|
|
from gen_wordmath import gen_example as _gen
|
||
|
|
if source == 'wordmath':
|
||
|
|
while True:
|
||
|
|
yield _gen(rng)
|
||
|
|
real = _realmath_texts(rank, world_size)
|
||
|
|
while True:
|
||
|
|
yield (next(real) if rng.random() < 0.75 else _gen(rng))
|
||
|
|
|
||
|
|
def build_web_stream(reference_ngrams, decon_enabled=True, rank=0, world_size=1):
|
||
|
|
fw = _fineweb_edu_texts(rank, world_size)
|
||
|
|
dclm = _dclm_texts(rank, world_size)
|
||
|
|
while True:
|
||
|
|
src = fw if random.random() < WEB_MIX['fineweb_edu'] else dclm
|
||
|
|
text = next(src)
|
||
|
|
if decon_enabled and is_contaminated(text, reference_ngrams):
|
||
|
|
continue
|
||
|
|
yield text
|
||
|
|
|
||
|
|
def build_preset_streams(preset: str, reference_ngrams, decon_enabled=True, rank=0, world_size=1):
|
||
|
|
|
||
|
|
def _filtered(gen):
|
||
|
|
for text in gen:
|
||
|
|
if decon_enabled and is_contaminated(text, reference_ngrams):
|
||
|
|
continue
|
||
|
|
yield text
|
||
|
|
mix = MIX_PRESETS[preset]
|
||
|
|
builders = {'fineweb_edu': _fineweb_edu_texts, 'dclm': _dclm_texts, 'cosmopedia': _cosmopedia_texts, 'math': _finemath_texts}
|
||
|
|
return {k: _filtered(builders[k](rank, world_size)) for k in mix}
|
||
|
|
|
||
|
|
def build_source_streams(reference_ngrams, decon_enabled=True, rank=0, world_size=1, mix_mode='default', arith_source='wordmath'):
|
||
|
|
|
||
|
|
def _filtered(gen, apply_decon):
|
||
|
|
for text in gen:
|
||
|
|
if apply_decon and decon_enabled and is_contaminated(text, reference_ngrams):
|
||
|
|
continue
|
||
|
|
yield text
|
||
|
|
streams = {'web': build_web_stream(reference_ngrams, decon_enabled, rank, world_size), 'arithmetic': _filtered(_arithmetic_texts(rank, world_size, arith_source), True)}
|
||
|
|
if mix_mode != 'museko':
|
||
|
|
streams['math'] = _filtered(_finemath_texts(rank, world_size), True)
|
||
|
|
streams['code'] = _filtered(_code_texts(rank, world_size), True)
|
||
|
|
streams['dialogue'] = _filtered(_smoltalk_texts(rank, world_size), True)
|
||
|
|
return streams
|
||
|
|
|
||
|
|
class MixedTokenStream:
|
||
|
|
|
||
|
|
def __init__(self, tokenizer, streams: dict, seq_len: int, seed: int=0, mix_mode: str='default', arith_frac: float | None=None, preset: str | None=None):
|
||
|
|
self.preset = preset
|
||
|
|
self.tokenizer = tokenizer
|
||
|
|
self.streams = streams
|
||
|
|
self.seq_len = seq_len
|
||
|
|
self.eos_id = tokenizer.token_to_id(EOS)
|
||
|
|
assert self.eos_id is not None, 'tokenizer is missing <|endoftext|>'
|
||
|
|
self.buffer: list[int] = []
|
||
|
|
self.rng = random.Random(seed)
|
||
|
|
self.progress = 0.0
|
||
|
|
self.mix_mode = mix_mode
|
||
|
|
self.arith_frac = arith_frac
|
||
|
|
|
||
|
|
def _pull_one_doc_tokens(self) -> list[int]:
|
||
|
|
if self.preset is not None:
|
||
|
|
mix = MIX_PRESETS[self.preset]
|
||
|
|
else:
|
||
|
|
mix = get_stage_mix(self.progress, self.mix_mode, self.arith_frac)
|
||
|
|
names = list(mix.keys())
|
||
|
|
weights = list(mix.values())
|
||
|
|
source_name = self.rng.choices(names, weights=weights, k=1)[0]
|
||
|
|
text = next(self.streams[source_name])
|
||
|
|
ids = self.tokenizer.encode(text).ids
|
||
|
|
return ids + [self.eos_id]
|
||
|
|
|
||
|
|
def next_sequence(self):
|
||
|
|
while len(self.buffer) < self.seq_len + 1:
|
||
|
|
self.buffer.extend(self._pull_one_doc_tokens())
|
||
|
|
seq = self.buffer[:self.seq_len + 1]
|
||
|
|
self.buffer = self.buffer[self.seq_len:]
|
||
|
|
return seq
|
||
|
|
|
||
|
|
def next_batch(self, batch_size: int):
|
||
|
|
return [self.next_sequence() for _ in range(batch_size)]
|
||
|
|
if __name__ == '__main__':
|
||
|
|
import time
|
||
|
|
from tokenizers import Tokenizer
|
||
|
|
tok = Tokenizer.from_file('artifacts/tokenizer.json')
|
||
|
|
print('eos id:', tok.token_to_id(EOS), flush=True)
|
||
|
|
streams = build_source_streams(reference_ngrams=set(), decon_enabled=False)
|
||
|
|
mixer = MixedTokenStream(tok, streams, seq_len=64)
|
||
|
|
t0 = time.time()
|
||
|
|
for i in range(5):
|
||
|
|
seq = mixer.next_sequence()
|
||
|
|
print(f'seq {i}: len={len(seq)} sample={seq[:10]}', flush=True)
|
||
|
|
print(f'5 sequences in {time.time() - t0:.2f}s', flush=True)
|