File size: 2,362 Bytes
90f058b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """MicroSupra-10k — preparação de dados (Fineweb-Edu + DCLM-Edu, streaming)."""
import os
import time
import numpy as np
from datasets import load_dataset, interleave_datasets
from transformers import PreTrainedTokenizerFast
from huggingface_hub import hf_hub_download
TOKENS_TRAIN = int(os.environ.get("TOKENS_TRAIN", 20_000_000))
TOKENS_VAL = int(os.environ.get("TOKENS_VAL", 1_000_000))
VOCAB = 1024
BOS, EOS = 0, 2 # <s>=0, </s>=2 (pad do original é <pad>=1; não usamos padding)
OUT = os.path.dirname(os.path.abspath(__file__)) + "/data"
os.makedirs(OUT, exist_ok=True)
print("[*] Baixando tokenizador do original (SupraLabs/MicroSupra-1k)...", flush=True)
tok_path = hf_hub_download("SupraLabs/MicroSupra-1k", "tokenizer.json")
tokenizer = PreTrainedTokenizerFast(
tokenizer_file=tok_path,
bos_token="<s>", eos_token="</s>", pad_token="<pad>", unk_token="<unk>",
)
print("[*] Abrindo streams: fineweb-edu sample-10BT + dclm-edu (edu_int_score>=3)...", flush=True)
fw = load_dataset("HuggingFaceFW/fineweb-edu", name="sample-10BT", split="train", streaming=True)
dclm = load_dataset("HuggingFaceTB/dclm-edu", split="train", streaming=True)
dclm = dclm.filter(lambda score: score >= 3, input_columns=["edu_int_score"])
mix = interleave_datasets([fw, dclm]) # 50/50
total_target = TOKENS_TRAIN + TOKENS_VAL
buf = np.empty(total_target, dtype=np.uint16)
pos, docs, t0 = 0, 0, time.time()
for ex in mix:
text = ex.get("text") or ""
if not text:
continue
ids = tokenizer.encode(text)
if not ids:
continue
ids = np.array([BOS] + ids + [EOS], dtype=np.int64)
ids = ids[ids < VOCAB] # segurança (BPE 1024 não deve estourar)
k = min(len(ids), total_target - pos)
buf[pos:pos + k] = ids[:k].astype(np.uint16)
pos += k
docs += 1
if pos % 2_000_000 < len(ids):
print(f" {pos/1e6:.1f}M tokens ({docs} docs, {time.time()-t0:.0f}s)", flush=True)
if pos >= total_target:
break
print(f"[*] Total tokenizado: {pos:,} tokens em {time.time()-t0:.0f}s", flush=True)
# separação: últimos TOKENS_VAL tokens viram validação
train = buf[:pos - TOKENS_VAL]
val = buf[pos - TOKENS_VAL:pos]
train.tofile(f"{OUT}/train.bin")
val.tofile(f"{OUT}/val.bin")
print(f"[*] train.bin: {len(train):,} tokens | val.bin: {len(val):,} tokens", flush=True)
print("[*] DONE", flush=True) |