Files
model-training/llm/corpus_fetch.py
T
2026-07-19 23:52:25 +04:00

79 lines
2.5 KiB
Python

"""Phase 2.2 — pull open RU datasets → data/cpt_raw/<bucket>/*.jsonl
Streams from HuggingFace so we never materialize full dumps. Each bucket
over-pulls ~1.5x its token target (cleaning drops 20-40%). Writes newline JSONL
{"text": "..."} shards; corpus_clean.py consumes them.
Usage:
huggingface-cli login # accept CulturaX/OSCAR licenses first
python corpus_fetch.py culturax # one bucket
python corpus_fetch.py all # every open bucket (not synth/logs)
Buckets synth (corpus_synth.py) and logs (corpus_logs.py) are handled separately.
"""
import sys
import json
from datasets import load_dataset
from corpus_common import RAW, TARGET_TOKENS, BUCKET_SHARE
# rough chars-per-token for RU under the Qwen tokenizer; only used to size pulls.
CHARS_PER_TOK = 3.5
OVERPULL = 1.5
# bucket -> (loader kwargs, function mapping a row to its text string)
SOURCES = {
"culturax": (
dict(path="allenai/c4", name="ru", split="train", streaming=True),
lambda r: r.get("text", ""),
),
"wiki": (
dict(path="wikimedia/wikipedia", name="20231101.ru", split="train", streaming=True),
lambda r: r.get("text", ""),
),
"books": (
dict(path="IlyaGusev/gazeta", split="train", streaming=True),
lambda r: r.get("text", ""),
),
}
def target_chars(bucket: str) -> int:
return int(TARGET_TOKENS * BUCKET_SHARE[bucket] * CHARS_PER_TOK * OVERPULL)
def fetch(bucket: str):
if bucket not in SOURCES:
raise SystemExit(f"unknown bucket {bucket}; known: {list(SOURCES)}")
budget = target_chars(bucket)
out_dir = RAW / bucket
out_dir.mkdir(parents=True, exist_ok=True)
out_path = out_dir / "shard-000.jsonl"
kw, getter = SOURCES[bucket]
ds = load_dataset(**kw)
written = 0
n = 0
with open(out_path, "w", encoding="utf-8") as f:
for row in ds:
text = (getter(row) or "").strip()
if len(text) < 200: # skip fragments early
continue
f.write(json.dumps({"text": text}, ensure_ascii=False) + "\n")
written += len(text)
n += 1
if n % 5000 == 0:
print(f"[{bucket}] {n} docs, {written/1e6:.1f}M chars / {budget/1e6:.1f}M target")
if written >= budget:
break
print(f"[{bucket}] DONE {n} docs, {written/1e6:.1f}M chars -> {out_path}")
if __name__ == "__main__":
args = sys.argv[1:] or ["all"]
buckets = list(SOURCES) if args == ["all"] else args
for b in buckets:
fetch(b)