初始化项目，由ModelHub XC社区提供模型
Model: bench-labs/cagliostro-v1 Source: Original Platform
This commit is contained in:
ModelHub XC
277
data_pipeline.py
Normal file
277
data_pipeline.py
Normal file
@@ -0,0 +1,277 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user