Add MP-01 (9500) 100M base: weights, config, tokenizer, model code, model card and eval results
11e2c7c verified | """Streaming adapters for each pretraining source. Each yields cleaned text documents.""" | |
| from __future__ import annotations | |
| import io | |
| import json | |
| import random | |
| import re | |
| from pathlib import Path | |
| import pyarrow as pa | |
| import pyarrow.compute as pc | |
| import pyarrow.parquet as pq | |
| import yaml | |
| from huggingface_hub import HfApi, hf_hub_download | |
| ROOT = Path(__file__).resolve().parent.parent | |
| RAW = ROOT / "data" / "raw" | |
| _BLANKS = re.compile(r"\n{4,}") | |
| def load_cfg(path=None): | |
| return yaml.safe_load(open(path or ROOT / "configs" / "data.yaml")) | |
| def _clean(t: str) -> str: | |
| t = t.replace("\r\n", "\n").replace("\x00", "") | |
| return _BLANKS.sub("\n\n\n", t).strip() | |
| class Source: | |
| def __init__(self, name: str, cfg: dict, glob: dict): | |
| self.name, self.cfg, self.glob = name, cfg, glob | |
| self.repo = cfg["repo"] | |
| self.min_chars = cfg.get("min_chars", glob.get("min_chars", 200)) | |
| self.quota = int(glob["total_tokens"] * cfg["weight"]) | |
| self.per_shard_cap = cfg.get("per_shard_cap_tokens") | |
| self.api = HfApi() | |
| # -- file selection (deterministic shuffle so a partial run is still a uniform sample) -- | |
| def list_shards(self) -> list[str]: | |
| files = [f.path for f in self.api.list_repo_tree(self.repo, repo_type="dataset", path_in_repo=self.cfg["prefix"].rstrip("/"), recursive=True) | |
| if hasattr(f, "size") and f.path.endswith((".parquet", ".jsonl.zst"))] | |
| files.sort() | |
| random.Random(self.glob["seed"]).shuffle(files) | |
| return files | |
| def download(self, path: str) -> Path: | |
| return Path(hf_hub_download(self.repo, path, repo_type="dataset", local_dir=RAW / self.name)) | |
| def docs(self, local: Path): | |
| raise NotImplementedError | |
| def _emit(self, texts): | |
| for t in texts: | |
| if t is None: | |
| continue | |
| t = _clean(t) | |
| if len(t) >= self.min_chars: | |
| yield t | |
| class ParquetSource(Source): | |
| text_col = "text" | |
| extra_cols: list[str] = [] | |
| def mask(self, tbl: pa.Table): | |
| return None | |
| def format(self, tbl: pa.Table): | |
| return tbl[self.text_col].to_pylist() | |
| def docs(self, local: Path): | |
| pf = pq.ParquetFile(local) | |
| cols = [self.text_col] + self.extra_cols | |
| for batch in pf.iter_batches(batch_size=4096, columns=cols): | |
| tbl = pa.Table.from_batches([batch]) | |
| m = self.mask(tbl) | |
| if m is not None: | |
| tbl = tbl.filter(m) | |
| yield from self._emit(self.format(tbl)) | |
| class FineWebEdu(ParquetSource): | |
| extra_cols = ["metadata"] | |
| def mask(self, tbl): | |
| score = pc.struct_field(tbl["metadata"], "score") | |
| return pc.greater_equal(score, self.cfg["min_edu_score"]) | |
| class Cosmopedia(ParquetSource): | |
| pass | |
| class FineMath(ParquetSource): | |
| extra_cols = ["score", "language", "language_score"] | |
| def mask(self, tbl): | |
| return pc.and_(pc.equal(tbl["language"], "en"), pc.greater_equal(tbl["language_score"], 0.9)) | |
| class Wikipedia(ParquetSource): | |
| extra_cols = ["title"] | |
| def format(self, tbl): | |
| out = [] | |
| for title, text in zip(tbl["title"].to_pylist(), tbl["text"].to_pylist()): | |
| if title.startswith(("List of", "Lists of", "Index of")) or len(text) < self.cfg.get("min_chars", 2000): | |
| continue | |
| out.append(f"# {title}\n\n{text}") | |
| return out | |
| class DCLM(ParquetSource): | |
| extra_cols = ["fasttext_score", "language", "language_score"] | |
| def list_shards(self): | |
| base = "filtered/OH_eli5_vs_rw_v2_bigram_200k_train/fasttext_openhermes_reddit_eli5_vs_rw_v2_bigram_200k_train/processed_data" | |
| rng = random.Random(self.glob["seed"]) | |
| out = [f"{base}/global-shard_{g:02d}_of_10/local-shard_{l}_of_10/shard_{n:08d}_processed.parquet" | |
| for g in range(1, 11) for l in range(10) for n in range(0, 260)] | |
| rng.shuffle(out) | |
| return out[:self.cfg.get("max_shards", 200)] # far more than needed; builder stops at quota | |
| def mask(self, tbl): | |
| m = pc.and_(pc.greater_equal(tbl["fasttext_score"], self.cfg["min_fasttext"]), | |
| pc.greater_equal(tbl["language_score"], self.cfg["min_lang_score"])) | |
| return pc.and_(m, pc.equal(tbl["language"], "en")) | |
| class StackExchange(ParquetSource): | |
| text_col = "ThreadText" | |
| extra_cols = ["Score", "AnswerCount"] | |
| def list_shards(self): | |
| out = [] | |
| for site in self.cfg["sites"]: | |
| d = site if "." in site else f"{site}.stackexchange.com" | |
| if site == "stackoverflow": | |
| d = "stackoverflow.com" | |
| try: | |
| fs = [f.path for f in self.api.list_repo_tree(self.repo, repo_type="dataset", path_in_repo=d) if f.path.endswith(".parquet")] | |
| except Exception: | |
| continue | |
| off, n = self.cfg.get("shard_offset", 0), self.cfg.get("shards_per_site", 1) | |
| out.extend(sorted(fs)[off:off + n]) # few shards per site; per_shard_cap keeps the mix balanced | |
| random.Random(self.glob["seed"]).shuffle(out) | |
| return out | |
| def mask(self, tbl): | |
| return pc.and_(pc.greater_equal(tbl["Score"], self.cfg["min_score"]), pc.greater_equal(tbl["AnswerCount"], 1)) | |
| class CodePython(Source): | |
| def docs(self, local: Path): | |
| import zstandard | |
| with open(local, "rb") as fh, zstandard.ZstdDecompressor().stream_reader(fh) as r: | |
| for line in io.TextIOWrapper(r, encoding="utf-8"): | |
| o = json.loads(line) | |
| if o.get("score", 0) < self.cfg["min_score"] or o.get("length_bytes", 0) < self.cfg["min_bytes"]: | |
| continue | |
| yield from self._emit([o["text"]]) | |
| REGISTRY = {"fineweb_edu": FineWebEdu, "dclm": DCLM, "cosmopedia": Cosmopedia, "finemath": FineMath, | |
| "wikipedia": Wikipedia, "stackexchange": StackExchange, "code_python": CodePython} | |
| def get_source(name: str, cfg: dict | None = None) -> Source: | |
| cfg = cfg or load_cfg() | |
| return REGISTRY[name](name, cfg["sources"][name], cfg) | |