social-data-attribution-w2/scripts/sampling/extract_stratified_docs.py

"""Extract stratified document text from archive-dolma3-pool-150b shards.

Reads a single shard from HuggingFace, extracts the text field for
doc_ids present in a pre-split rows file, merges that text with the
original stratified metadata, and writes merged documents to a
compressed jsonl.zst output file plus a per-worker parquet manifest.

Designed for a 5-job / 20-worker-per-job SLURM layout where each
worker handles one shard (shard_index = job_id * workers_per_job + worker_id).

Requires pre-split row files generated by pre_split_stratified_ids.py.

Usage:
    python scripts/sampling/extract_stratified_docs.py \
        --job-id 0 \
        --worker-id 0 \
        --workers-per-job 20 \
        --id-dir stratified_data/shard_rows \
        --shard-list scripts/slurm/data/stratified_shard_list.txt \
        --output-dir stratified_data/merged_docs \
        --cache-dir /tmp/hf_cache
"""

import argparse
import io
import json
import logging
from pathlib import Path

import pyarrow as pa
import pyarrow.parquet as pq
import zstandard

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
)
log = logging.getLogger(__name__)


def shard_path_to_hf_path(shard_path: str) -> str:
    stem = shard_path.replace(".enriched.jsonl.zst", ".jsonl.zst")
    parts = stem.split("_")
    job_part = f"{parts[0]}_{parts[1]}"
    worker_part = f"{parts[2]}_{parts[3]}"
    train_part = "_".join(parts[4:])
    return f"{job_part}/{worker_part}/{train_part}"


def load_doc_rows(rows_file: Path) -> dict[str, dict]:
    rows: dict[str, dict] = {}
    for line in rows_file.read_text().splitlines():
        if not line.strip():
            continue
        rec = json.loads(line)
        rows[rec["doc_id"]] = rec
    return rows


def extract_docs(
    shard_file: Path,
    doc_rows: dict[str, dict],
    output_path: Path,
) -> list[str]:
    dctx = zstandard.ZstdDecompressor()
    cctx = zstandard.ZstdCompressor(level=3)
    matched_ids: list[str] = []

    with (
        open(shard_file, "rb") as fin,
        open(output_path, "wb") as fout,
    ):
        reader = dctx.stream_reader(fin)
        text_in = io.TextIOWrapper(reader, encoding="utf-8")
        writer = cctx.stream_writer(fout)

        for line in text_in:
            hf_rec = json.loads(line)
            doc_id = hf_rec.get("id")
            if doc_id in doc_rows:
                merged = dict(doc_rows[doc_id])
                merged["text"] = hf_rec["text"]
                writer.write((json.dumps(merged) + "\n").encode("utf-8"))
                matched_ids.append(doc_id)
                if len(matched_ids) == len(doc_rows):
                    break

        writer.close()

    return matched_ids


def write_manifest(
    manifest_path: Path,
    matched_ids: list[str],
    shard_path: str,
    job_id: int,
    worker_id: int,
) -> None:
    table = pa.table(
        {
            "doc_id": matched_ids,
            "shard_path": [shard_path] * len(matched_ids),
            "job_id": [job_id] * len(matched_ids),
            "worker_id": [worker_id] * len(matched_ids),
        }
    )
    pq.write_table(table, manifest_path)
    log.info("Wrote manifest: %d rows -> %s", len(matched_ids), manifest_path)


def main() -> None:
    parser = argparse.ArgumentParser(description="Extract stratified docs from a shard")
    parser.add_argument("--job-id", type=int, required=True)
    parser.add_argument("--worker-id", type=int, required=True)
    parser.add_argument("--workers-per-job", type=int, default=20)
    parser.add_argument("--id-dir", type=Path, required=True)
    parser.add_argument("--shard-list", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--cache-dir", type=Path, default=Path("/tmp/hf_cache"))
    parser.add_argument("--dataset", type=str, default="HCAI-Lab/archive-dolma3-pool-150b")
    args = parser.parse_args()

    shard_index = args.job_id * args.workers_per_job + args.worker_id

    shard_lines = args.shard_list.read_text().strip().splitlines()
    if shard_index >= len(shard_lines):
        log.error(
            "Shard index %d out of range (max %d)", shard_index, len(shard_lines) - 1
        )
        raise SystemExit(1)

    shard_path = shard_lines[shard_index].strip()
    hf_path = shard_path_to_hf_path(shard_path)

    log.info(
        "Job: %d, Worker: %d, Shard index: %d", args.job_id, args.worker_id, shard_index
    )
    log.info("Manifest shard_path: %s", shard_path)
    log.info("HF file path: %s", hf_path)

    rows_file = args.id_dir / shard_path.replace(".jsonl.zst", ".rows.jsonl")
    if not rows_file.exists():
        log.info("No rows file for this shard (%s), nothing to extract.", rows_file)
        return

    doc_rows = load_doc_rows(rows_file)
    log.info("Loaded %d doc rows from %s", len(doc_rows), rows_file)

    if not doc_rows:
        log.info("Empty rows file, nothing to extract.")
        return

    output_name = shard_path.replace(".enriched.jsonl.zst", ".docs.jsonl.zst")
    output_path = args.output_dir / output_name
    manifest_name = shard_path.replace(".enriched.jsonl.zst", ".manifest.parquet")
    manifest_path = args.output_dir / manifest_name

    if output_path.exists() and manifest_path.exists():
        log.info("Output already exists: %s. Skipping.", output_path)
        return

    log.info("Downloading shard from HuggingFace...")
    from huggingface_hub import hf_hub_download

    local_path = hf_hub_download(
        args.dataset,
        hf_path,
        repo_type="dataset",
        cache_dir=str(args.cache_dir),
    )
    log.info("Downloaded to: %s", local_path)

    args.output_dir.mkdir(parents=True, exist_ok=True)

    log.info("Extracting matching documents...")
    matched_ids = extract_docs(Path(local_path), doc_rows, output_path)
    log.info(
        "Extracted %d / %d docs -> %s", len(matched_ids), len(doc_rows), output_path
    )

    write_manifest(manifest_path, matched_ids, shard_path, args.job_id, args.worker_id)

    if len(matched_ids) < len(doc_rows):
        log.warning(
            "Missing %d doc_ids (not found in shard)",
            len(doc_rows) - len(matched_ids),
        )


if __name__ == "__main__":
    main()