490 lines
19 KiB
Python
490 lines
19 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import gzip
|
|
import json
|
|
import os
|
|
import re
|
|
import time
|
|
from collections import Counter, defaultdict
|
|
from pathlib import Path
|
|
|
|
from datasets import load_dataset
|
|
from transformers import AutoTokenizer
|
|
|
|
|
|
BUDGETS = {
|
|
"english_web": 250_000_000,
|
|
"english_edu": 200_000_000,
|
|
"chinese_clean": 250_000_000,
|
|
"code": 150_000_000,
|
|
"math": 100_000_000,
|
|
"science": 30_000_000,
|
|
"qa_as_text": 20_000_000,
|
|
}
|
|
|
|
|
|
STREAM_SOURCES = {
|
|
"english_web": [
|
|
{"kind": "hf", "name": "HuggingFaceFW/fineweb", "config": "CC-MAIN-2025-26", "split": "train", "max_rows": 0},
|
|
{"kind": "hf", "name": "HuggingFaceFW/fineweb", "config": "CC-MAIN-2025-21", "split": "train", "max_rows": 0},
|
|
],
|
|
"english_edu": [
|
|
{"kind": "hf", "name": "HuggingFaceFW/fineweb-edu", "config": None, "split": "train", "max_rows": 0},
|
|
],
|
|
"chinese_clean": [
|
|
{"kind": "hf", "name": "BAAI/CCI3-HQ", "config": None, "split": "train", "max_rows": 0},
|
|
{"kind": "hf", "name": "Skywork/SkyPile-150B", "config": None, "split": "train", "max_rows": 0},
|
|
],
|
|
"code": [
|
|
{"kind": "hf", "name": "bigcode/starcoderdata", "config": None, "split": "train", "max_rows": 0},
|
|
{"kind": "hf", "name": "codeparrot/github-code", "config": None, "split": "train", "max_rows": 0},
|
|
],
|
|
"math": [
|
|
{"kind": "hf", "name": "open-web-math/open-web-math", "config": None, "split": "train", "max_rows": 0},
|
|
{"kind": "hf", "name": "GAIR/MathPile", "config": None, "split": "train", "max_rows": 0},
|
|
],
|
|
}
|
|
|
|
|
|
LOCAL_GLOBS = {
|
|
"science": [
|
|
"data/offline_text_only_reasoning_sources_20260611/science_reasoning__*.jsonl",
|
|
"data/offline_text_only_reasoning_sources_20260611/logic__*.jsonl",
|
|
],
|
|
"qa_as_text": [
|
|
"data/training_mix_v4_train1m_test2p8k_noupsample_nobbh_20260611/train_1m.jsonl",
|
|
"data/open_recovery_sft_mix_alt_sources_1m_parquet_20260607/normalized.jsonl",
|
|
],
|
|
}
|
|
|
|
LOCAL_PARQUET_GLOBS = {
|
|
"english_web": [
|
|
"data/raw_parquets/fineweb_2025/*.parquet",
|
|
],
|
|
"english_edu": [
|
|
"data/raw_parquets/fineweb_edu/*.parquet",
|
|
],
|
|
"code": [
|
|
"data/raw_parquets/starcoder/python/*.parquet",
|
|
"data/raw_parquets/starcoder/javascript/*.parquet",
|
|
"data/raw_parquets/starcoder/typescript/*.parquet",
|
|
"data/raw_parquets/starcoder/java/*.parquet",
|
|
"data/raw_parquets/starcoder/cpp/*.parquet",
|
|
"data/raw_parquets/starcoder/go/*.parquet",
|
|
"data/raw_parquets/starcoder/rust/*.parquet",
|
|
"data/raw_parquets/starcoder/shell/*.parquet",
|
|
],
|
|
}
|
|
|
|
|
|
def clean_text(text):
|
|
if text is None:
|
|
return ""
|
|
text = str(text).replace("\x00", " ")
|
|
text = re.sub(r"[ \t\r\f\v]+", " ", text)
|
|
text = re.sub(r"\n{4,}", "\n\n\n", text)
|
|
return text.strip()
|
|
|
|
|
|
def first_user_assistant(messages):
|
|
user = None
|
|
assistant = None
|
|
for msg in messages or []:
|
|
role = msg.get("role")
|
|
content = clean_text(msg.get("content"))
|
|
if not content:
|
|
continue
|
|
if role == "user" and user is None:
|
|
user = content
|
|
elif role == "assistant" and user is not None:
|
|
assistant = content
|
|
break
|
|
if user and assistant:
|
|
return user, assistant
|
|
return None, None
|
|
|
|
|
|
def qa_text(user, assistant):
|
|
if not user or not assistant:
|
|
return ""
|
|
user_label = "" if user.lstrip().lower().startswith(("question:", "problem:", "context:", "support:", "fact1:")) else "Question:\n"
|
|
assistant_label = "" if assistant.lstrip().lower().startswith(("answer:", "solution:", "final answer:", "explanation:")) else "Answer:\n"
|
|
return f"{user_label}{user}\n\n{assistant_label}{assistant}"
|
|
|
|
|
|
def extract_text(row, category, source_name):
|
|
if category == "science":
|
|
user, assistant = first_user_assistant(row.get("messages"))
|
|
if user and assistant:
|
|
return qa_text(user, assistant)
|
|
|
|
if category == "qa_as_text" or row.get("messages"):
|
|
user, assistant = first_user_assistant(row.get("messages"))
|
|
if user and assistant:
|
|
return qa_text(user, assistant)
|
|
|
|
if category == "code":
|
|
content = clean_text(row.get("content") or row.get("text") or row.get("code"))
|
|
if not content:
|
|
return ""
|
|
path = clean_text(row.get("path") or row.get("max_stars_repo_path") or row.get("repo_name") or "")
|
|
lang = clean_text(row.get("lang") or row.get("language") or "")
|
|
if path or lang:
|
|
attrs = []
|
|
if path:
|
|
attrs.append(f'path="{path[:300]}"')
|
|
if lang:
|
|
attrs.append(f'language="{lang[:80]}"')
|
|
return f"<file {' '.join(attrs)}>\n{content}\n</file>"
|
|
return content
|
|
|
|
for key in ("text", "content", "markdown", "raw_content"):
|
|
text = clean_text(row.get(key))
|
|
if text:
|
|
return text
|
|
return ""
|
|
|
|
|
|
def token_count(tok, text):
|
|
return len(tok.encode(text, add_special_tokens=False))
|
|
|
|
|
|
def open_gz_writer(path):
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
return gzip.open(path, "at", encoding="utf-8")
|
|
|
|
|
|
def write_doc(writer, doc_id, category, source, text, ntok, metadata=None):
|
|
writer.write(
|
|
json.dumps(
|
|
{
|
|
"id": doc_id,
|
|
"category": category,
|
|
"source": source,
|
|
"text": text,
|
|
"token_count": ntok,
|
|
"metadata": metadata or {},
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
+ "\n"
|
|
)
|
|
|
|
|
|
def should_keep(row, category, text, ntok, args):
|
|
if ntok < args.min_tokens:
|
|
return False, "too_short"
|
|
if ntok > args.max_doc_tokens:
|
|
return False, "too_long"
|
|
if category == "english_web":
|
|
lang = row.get("language")
|
|
score = row.get("language_score")
|
|
if lang and str(lang).lower() != "en":
|
|
return False, "non_en"
|
|
if score is not None:
|
|
try:
|
|
if float(score) < args.min_fineweb_language_score:
|
|
return False, "low_language_score"
|
|
except Exception:
|
|
pass
|
|
return True, ""
|
|
|
|
|
|
def iter_local_jsonl(paths):
|
|
for path in paths:
|
|
with open(path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
if line.strip():
|
|
yield path, json.loads(line)
|
|
|
|
|
|
def flush_pack(writer, category, source, parts, ntok, stats, args):
|
|
if not parts:
|
|
return [], 0
|
|
if ntok < args.min_tokens:
|
|
stats["rejected"][category]["pack_too_short"] += 1
|
|
return [], 0
|
|
idx = stats["docs_by_category"][category]
|
|
text = "\n\n---\n\n".join(parts)
|
|
write_doc(
|
|
writer,
|
|
f"{category}_packed_{idx:09d}",
|
|
category,
|
|
source,
|
|
text,
|
|
ntok,
|
|
{"local_path": source, "packed_items": len(parts)},
|
|
)
|
|
stats["docs_by_category"][category] += 1
|
|
stats["tokens_by_category"][category] += ntok
|
|
stats["tokens_by_source"][source] += ntok
|
|
if stats["docs_by_category"][category] % args.log_every == 0:
|
|
print_progress(stats, category)
|
|
return [], 0
|
|
|
|
|
|
def collect_local(base, tok, category, out_dir, args, stats):
|
|
paths = []
|
|
for pattern in LOCAL_GLOBS.get(category, []):
|
|
paths.extend(str(p) for p in sorted((base).glob(pattern)))
|
|
if not paths:
|
|
stats["sources"].append({"category": category, "kind": "local", "error": "no local files"})
|
|
return
|
|
|
|
out_path = out_dir / "documents" / f"{category}.jsonl.gz"
|
|
with open_gz_writer(out_path) as writer:
|
|
pack_parts = []
|
|
pack_tokens = 0
|
|
pack_source = ""
|
|
for path, row in iter_local_jsonl(paths):
|
|
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
|
|
break
|
|
text = extract_text(row, category, path)
|
|
if not text:
|
|
stats["rejected"][category]["empty"] += 1
|
|
continue
|
|
ntok = token_count(tok, text)
|
|
if category == "science" and ntok < args.min_tokens:
|
|
if ntok > args.max_doc_tokens:
|
|
stats["rejected"][category]["too_long"] += 1
|
|
continue
|
|
if pack_source and pack_source != path:
|
|
pack_parts, pack_tokens = flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
|
|
pack_source = path
|
|
if pack_tokens + ntok > args.science_pack_tokens and pack_parts:
|
|
pack_parts, pack_tokens = flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
|
|
pack_parts.append(text)
|
|
pack_tokens += ntok
|
|
continue
|
|
keep, reason = should_keep(row, category, text, ntok, args)
|
|
if not keep:
|
|
stats["rejected"][category][reason] += 1
|
|
continue
|
|
if category == "science" and pack_parts:
|
|
pack_parts, pack_tokens = flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
|
|
idx = stats["docs_by_category"][category]
|
|
write_doc(writer, f"{category}_local_{idx:09d}", category, path, text, ntok, {"local_path": path})
|
|
stats["docs_by_category"][category] += 1
|
|
stats["tokens_by_category"][category] += ntok
|
|
stats["tokens_by_source"][path] += ntok
|
|
if stats["docs_by_category"][category] % args.log_every == 0:
|
|
print_progress(stats, category)
|
|
if category == "science" and pack_parts and stats["tokens_by_category"][category] < BUDGETS[category] * args.scale:
|
|
flush_pack(writer, category, pack_source, pack_parts, pack_tokens, stats, args)
|
|
|
|
|
|
def collect_local_parquet(base, tok, category, out_dir, args, stats):
|
|
try:
|
|
import pyarrow.parquet as pq
|
|
except Exception as exc:
|
|
stats["sources"].append({"category": category, "kind": "local_parquet", "error": repr(exc)})
|
|
return
|
|
|
|
paths = []
|
|
for pattern in LOCAL_PARQUET_GLOBS.get(category, []):
|
|
paths.extend(sorted(base.glob(pattern)))
|
|
if not paths:
|
|
stats["sources"].append({"category": category, "kind": "local_parquet", "error": "no parquet files"})
|
|
return
|
|
|
|
out_path = out_dir / "documents" / f"{category}.jsonl.gz"
|
|
with open_gz_writer(out_path) as writer:
|
|
for path in paths:
|
|
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
|
|
break
|
|
source_label = str(path.relative_to(base))
|
|
source_rec = {"category": category, "kind": "local_parquet", "source": source_label, "rows_seen": 0, "docs_written": 0, "tokens": 0, "error": ""}
|
|
try:
|
|
pf = pq.ParquetFile(path)
|
|
for batch in pf.iter_batches(batch_size=args.parquet_batch_size):
|
|
rows = batch.to_pylist()
|
|
for row in rows:
|
|
source_rec["rows_seen"] += 1
|
|
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
|
|
break
|
|
text = extract_text(row, category, source_label)
|
|
if not text:
|
|
stats["rejected"][category]["empty"] += 1
|
|
continue
|
|
ntok = token_count(tok, text)
|
|
keep, reason = should_keep(row, category, text, ntok, args)
|
|
if not keep:
|
|
stats["rejected"][category][reason] += 1
|
|
continue
|
|
idx = stats["docs_by_category"][category]
|
|
write_doc(
|
|
writer,
|
|
f"{category}_{safe_source(source_label)}_{idx:09d}",
|
|
category,
|
|
source_label,
|
|
text,
|
|
ntok,
|
|
{k: row.get(k) for k in ("url", "dump", "date", "language", "language_score", "path", "lang", "license") if k in row},
|
|
)
|
|
stats["docs_by_category"][category] += 1
|
|
stats["tokens_by_category"][category] += ntok
|
|
stats["tokens_by_source"][source_label] += ntok
|
|
source_rec["docs_written"] += 1
|
|
source_rec["tokens"] += ntok
|
|
if stats["docs_by_category"][category] % args.log_every == 0:
|
|
print_progress(stats, category)
|
|
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
|
|
break
|
|
except Exception as exc:
|
|
source_rec["error"] = repr(exc)
|
|
print(json.dumps({"event": "local_parquet_error", **source_rec}, ensure_ascii=False), flush=True)
|
|
stats["sources"].append(source_rec)
|
|
|
|
|
|
def collect_stream(tok, category, out_dir, args, stats):
|
|
out_path = out_dir / "documents" / f"{category}.jsonl.gz"
|
|
with open_gz_writer(out_path) as writer:
|
|
for spec in STREAM_SOURCES.get(category, []):
|
|
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
|
|
break
|
|
source_label = spec["name"] if not spec.get("config") else f"{spec['name']}:{spec['config']}"
|
|
source_rec = {"category": category, "kind": "hf_stream", "source": source_label, "rows_seen": 0, "docs_written": 0, "tokens": 0, "error": ""}
|
|
try:
|
|
ds = load_dataset(
|
|
spec["name"],
|
|
spec.get("config"),
|
|
split=spec.get("split", "train"),
|
|
streaming=True,
|
|
trust_remote_code=True,
|
|
)
|
|
for row in ds:
|
|
source_rec["rows_seen"] += 1
|
|
if spec.get("max_rows") and source_rec["rows_seen"] > spec["max_rows"]:
|
|
break
|
|
if stats["tokens_by_category"][category] >= BUDGETS[category] * args.scale:
|
|
break
|
|
text = extract_text(row, category, source_label)
|
|
if not text:
|
|
stats["rejected"][category]["empty"] += 1
|
|
continue
|
|
ntok = token_count(tok, text)
|
|
keep, reason = should_keep(row, category, text, ntok, args)
|
|
if not keep:
|
|
stats["rejected"][category][reason] += 1
|
|
continue
|
|
idx = stats["docs_by_category"][category]
|
|
write_doc(
|
|
writer,
|
|
f"{category}_{safe_source(source_label)}_{idx:09d}",
|
|
category,
|
|
source_label,
|
|
text,
|
|
ntok,
|
|
{k: row.get(k) for k in ("url", "dump", "date", "language", "language_score", "path", "lang", "license") if k in row},
|
|
)
|
|
stats["docs_by_category"][category] += 1
|
|
stats["tokens_by_category"][category] += ntok
|
|
stats["tokens_by_source"][source_label] += ntok
|
|
source_rec["docs_written"] += 1
|
|
source_rec["tokens"] += ntok
|
|
if stats["docs_by_category"][category] % args.log_every == 0:
|
|
print_progress(stats, category)
|
|
except Exception as exc:
|
|
source_rec["error"] = repr(exc)
|
|
print(json.dumps({"event": "source_error", **source_rec}, ensure_ascii=False), flush=True)
|
|
stats["sources"].append(source_rec)
|
|
|
|
|
|
def safe_source(text):
|
|
return re.sub(r"[^A-Za-z0-9._-]+", "_", text)[:120]
|
|
|
|
|
|
def print_progress(stats, category):
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"event": "progress",
|
|
"category": category,
|
|
"docs": stats["docs_by_category"][category],
|
|
"tokens": stats["tokens_by_category"][category],
|
|
"target": int(BUDGETS[category] * stats["scale"]),
|
|
"elapsed_sec": time.time() - stats["start_time"],
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
flush=True,
|
|
)
|
|
|
|
|
|
def dump_stats(out_dir, stats):
|
|
serializable = {
|
|
"scale": stats["scale"],
|
|
"budgets": {k: int(v * stats["scale"]) for k, v in BUDGETS.items()},
|
|
"tokens_by_category": dict(stats["tokens_by_category"]),
|
|
"docs_by_category": dict(stats["docs_by_category"]),
|
|
"tokens_by_source": dict(stats["tokens_by_source"].most_common()),
|
|
"rejected": {k: dict(v) for k, v in stats["rejected"].items()},
|
|
"sources": stats["sources"],
|
|
"elapsed_sec": time.time() - stats["start_time"],
|
|
}
|
|
(out_dir / "stats.json").write_text(json.dumps(serializable, ensure_ascii=False, indent=2), encoding="utf-8")
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--base-dir", default="/ssd/yi/Tokenizer_Swap")
|
|
parser.add_argument("--tokenizer", default="model_building/generated_models/Qwen3-0.6B-DSV4-tokenizer-remap-v2")
|
|
parser.add_argument("--out-dir", default="data/cpt_docmix_1b_8192_20260613")
|
|
parser.add_argument("--scale", type=float, default=1.0, help="1.0 = 1B token budget; 0.001 = 1M token smoke")
|
|
parser.add_argument("--categories", default="english_web,english_edu,chinese_clean,code,math,science,qa_as_text")
|
|
parser.add_argument("--min-tokens", type=int, default=128)
|
|
parser.add_argument("--max-doc-tokens", type=int, default=32768)
|
|
parser.add_argument("--min-fineweb-language-score", type=float, default=0.65)
|
|
parser.add_argument("--log-every", type=int, default=1000)
|
|
parser.add_argument("--parquet-batch-size", type=int, default=1000)
|
|
parser.add_argument("--science-pack-tokens", type=int, default=2048)
|
|
args = parser.parse_args()
|
|
|
|
base = Path(args.base_dir)
|
|
out_dir = base / args.out_dir
|
|
out_dir.mkdir(parents=True, exist_ok=True)
|
|
(out_dir / "manifest.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"tokenizer": str(base / args.tokenizer),
|
|
"budgets": {k: int(v * args.scale) for k, v in BUDGETS.items()},
|
|
"seq_len_for_later_packing": 8192,
|
|
"stream_sources": STREAM_SOURCES,
|
|
"local_globs": LOCAL_GLOBS,
|
|
"local_parquet_globs": LOCAL_PARQUET_GLOBS,
|
|
"min_tokens": args.min_tokens,
|
|
"max_doc_tokens": args.max_doc_tokens,
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
tok = AutoTokenizer.from_pretrained(base / args.tokenizer, trust_remote_code=True)
|
|
stats = {
|
|
"scale": args.scale,
|
|
"start_time": time.time(),
|
|
"tokens_by_category": Counter(),
|
|
"docs_by_category": Counter(),
|
|
"tokens_by_source": Counter(),
|
|
"rejected": defaultdict(Counter),
|
|
"sources": [],
|
|
}
|
|
|
|
for category in [x.strip() for x in args.categories.split(",") if x.strip()]:
|
|
if category in LOCAL_PARQUET_GLOBS:
|
|
collect_local_parquet(base, tok, category, out_dir, args, stats)
|
|
if category in STREAM_SOURCES:
|
|
collect_stream(tok, category, out_dir, args, stats)
|
|
if category in LOCAL_GLOBS:
|
|
collect_local(base, tok, category, out_dir, args, stats)
|
|
dump_stats(out_dir, stats)
|
|
|
|
dump_stats(out_dir, stats)
|
|
print(out_dir / "stats.json")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|