init

llm/corpus_fetch.py

@@ -0,0 +1,78 @@
+"""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)