| """Corpus pipeline: fetch -> normalize (+filter, dedupe) -> tokenizer sample -> tokenize (+decontaminate, split, shard). |
| |
| Usage: python -m zenith.data.pipeline {fetch,normalize,sample,tokenize} [--sources a,b] [options] |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| from collections import defaultdict |
| from pathlib import Path |
|
|
| from zenith.data.clean import dedupe, keep, normalize_text |
| from zenith.data.decontam import build_ngram_index, is_contaminated |
| from zenith.data.shards import ShardWriter, split_of, write_manifest |
| from zenith.data.sources.base import Record |
| from zenith.tokenizer.export import encode_fast, load_fast |
|
|
| RAW_DIR = Path("data/raw") |
| INTERIM_DIR = Path("data/interim") |
| PROCESSED_DIR = Path("data/processed") |
| EVAL_DIR = Path("data/eval") |
| TOKENIZER_HF = Path("zenith/tokenizer/zenith-16k.hf.json") |
| SHOWCASE = Path("zenith/evals/showcase.txt") |
| DEFAULT_MIX = {"astro": 0.30, "arxiv": 0.15, "textbook": 0.05, "general": 0.50} |
| NGRAM = 8 |
|
|
|
|
| def _read_interim(path: Path): |
| with open(path) as f: |
| for line in f: |
| yield Record.from_json(line) |
|
|
|
|
| def normalize_stage(raw_dir: Path, interim_dir: Path, sources: dict) -> None: |
| interim_dir.mkdir(parents=True, exist_ok=True) |
|
|
| def cleaned(): |
| for mod in sources.values(): |
| for r in mod.iter_records(raw_dir): |
| r = Record(r.source, r.id, r.license, r.group, normalize_text(r.text)) |
| if keep(r): |
| yield r |
|
|
| handles = {name: open(interim_dir / f"{name}.jsonl", "w") for name in sources} |
| try: |
| for r in dedupe(cleaned()): |
| handles[r.source].write(r.to_json() + "\n") |
| finally: |
| for h in handles.values(): |
| h.close() |
|
|
|
|
| def sample_stage(interim_dir: Path, out_path: Path, total_bytes: int, mix: dict[str, float]) -> None: |
| """Write a tokenizer-training sample with the group mix, taking evenly spaced docs from each group.""" |
| by_group: dict[str, list[Path]] = defaultdict(list) |
| sizes: dict[str, int] = defaultdict(int) |
| for path in sorted(interim_dir.glob("*.jsonl")): |
| first = next(_read_interim(path), None) |
| if first is None: |
| continue |
| by_group[first.group].append(path) |
| sizes[first.group] += path.stat().st_size |
| with open(out_path, "w") as out: |
| for group, share in mix.items(): |
| budget = int(total_bytes * share) |
| stride = max(1, sizes[group] // max(budget, 1)) |
| written = 0 |
| for path in by_group.get(group, []): |
| for i, r in enumerate(_read_interim(path)): |
| if written >= budget: |
| break |
| if i % stride == 0: |
| chunk = r.text[: budget - written] |
| out.write(chunk + "\n\n") |
| written += len(chunk.encode()) + 2 |
|
|
|
|
| def eval_texts_from(eval_dir: Path, showcase_path: Path | None) -> tuple[list[str], str]: |
| """Every evaluation question (`question` or `text` field) plus showcase lines, and their combined hash.""" |
| texts, h = [], hashlib.sha256() |
| paths = sorted(Path(eval_dir).glob("*.jsonl")) + ([Path(showcase_path)] if showcase_path and Path(showcase_path).exists() else []) |
| for path in paths: |
| data = path.read_bytes() |
| h.update(data) |
| for line in data.decode().splitlines(): |
| if not line.strip(): |
| continue |
| if path.suffix == ".jsonl": |
| row = json.loads(line) |
| texts.append(row.get("question") or row["text"]) |
| else: |
| texts.append(line.strip()) |
| return texts, h.hexdigest() |
|
|
|
|
| def tokenize_stage(interim_dir: Path, out_dir: Path, tokenizer_path: Path, eval_dir: Path, sources: dict, |
| showcase_path: Path | None = SHOWCASE, shard_tokens: int = 100_000_000, batch: int = 512, |
| allow_no_decontam: bool = False) -> None: |
| fast = load_fast(tokenizer_path) |
| bos, eos = fast.token_to_id("<|bos|>"), fast.token_to_id("<|eos|>") |
| eval_texts, eval_sha = eval_texts_from(eval_dir, showcase_path) |
| if not eval_texts and not allow_no_decontam: |
| raise ValueError( |
| f"no evaluation texts in {eval_dir}: decontamination would silently do nothing. " |
| "Run `python -m zenith.sft.generate` first, or pass --no-decontam deliberately." |
| ) |
| letter_ids = {i for t, i in fast.get_vocab().items() if any(c.isalpha() for c in fast.decode([i]))} |
| index = build_ngram_index(encode_fast(fast, eval_texts), n=NGRAM, letter_ids=letter_ids) |
| decontamination = {"eval_texts": len(eval_texts), "ngrams": len(index), |
| "eval_sha256": eval_sha if eval_texts else None, "ngram": NGRAM} |
|
|
| writers: dict[tuple[str, str], ShardWriter] = {} |
| source_stats: dict[str, dict] = {} |
| for name in sources: |
| path = interim_dir / f"{name}.jsonl" |
| if not path.exists(): |
| continue |
| st = {"docs": 0, "tokens": 0, "bytes": 0, "contaminated": 0, "group": None, "licenses": set()} |
| pending: list[Record] = [] |
|
|
| def flush(): |
| for r, ids in zip(pending, encode_fast(fast, [r.text for r in pending])): |
| if is_contaminated(ids, index, n=NGRAM): |
| st["contaminated"] += 1 |
| continue |
| split = split_of(r.id) |
| key = (r.group, split) |
| if key not in writers: |
| writers[key] = ShardWriter(out_dir / r.group, split, bos, eos, shard_tokens) |
| n_bytes = len(r.text.encode("utf-8", errors="surrogatepass")) |
| writers[key].add(ids, n_bytes) |
| st["docs"] += 1 |
| st["tokens"] += len(ids) + 2 |
| st["bytes"] += n_bytes |
| st["group"] = r.group |
| st["licenses"].add(r.license) |
| pending.clear() |
|
|
| for r in _read_interim(path): |
| pending.append(r) |
| if len(pending) >= batch: |
| flush() |
| flush() |
| st["licenses"] = sorted(st["licenses"])[:20] |
| source_stats[name] = st |
|
|
| groups: dict[str, dict] = defaultdict(dict) |
| for (group, split), w in sorted(writers.items()): |
| groups[group][split] = {"tokens": w.tokens, "bytes": w.bytes, "shards": w.close()} |
| tok_sha = hashlib.sha256(Path(tokenizer_path).read_bytes()).hexdigest() |
| write_manifest(out_dir, {"tokenizer_sha256": tok_sha, "groups": dict(groups), "sources": source_stats, |
| "decontamination": decontamination}) |
|
|
|
|
| def main(argv=None) -> None: |
| from zenith.data.sources import SOURCES |
|
|
| ap = argparse.ArgumentParser(prog="python -m zenith.data.pipeline") |
| ap.add_argument("stage", choices=["fetch", "normalize", "sample", "tokenize"]) |
| ap.add_argument("--sources", help="comma-separated subset of sources") |
| ap.add_argument("--bytes", default="200MB", help="tokenizer sample size (sample stage)") |
| ap.add_argument("--tokenizer", type=Path, default=TOKENIZER_HF) |
| ap.add_argument("--no-decontam", action="store_true", help="tokenize without evaluation texts (not for real runs)") |
| args = ap.parse_args(argv) |
| sources = {k: v for k, v in SOURCES.items() if not args.sources or k in args.sources.split(",")} |
|
|
| if args.stage == "fetch": |
| for name, mod in sources.items(): |
| print(f"fetch {name}", flush=True) |
| mod.fetch(RAW_DIR) |
| elif args.stage == "normalize": |
| normalize_stage(RAW_DIR, INTERIM_DIR, sources) |
| elif args.stage == "sample": |
| units = {"KB": 10**3, "MB": 10**6, "GB": 10**9} |
| n = args.bytes.upper() |
| total = int(float(n[:-2]) * units[n[-2:]]) if n[-2:] in units else int(n) |
| sample_stage(INTERIM_DIR, INTERIM_DIR / "tokenizer_sample.txt", total, DEFAULT_MIX) |
| else: |
| tokenize_stage(INTERIM_DIR, PROCESSED_DIR, args.tokenizer, EVAL_DIR, sources, allow_no_decontam=args.no_decontam) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
