social-data-attribution-w2/src/dolma/materialize.py

"""Pipeline to stream, sample, and persist Dolma data."""

from __future__ import annotations

from pathlib import Path
from typing import Sequence

from .constants import DEFAULT_SPLIT
from .outputs import SampleManifest, write_jsonl, write_manifest, write_parquet
from .sample import sample_documents
from .shards import dolma_builder, stream_dolma


def _write_output(records, path: Path, output_format: str) -> None:
    suffix = output_format.lower()
    if suffix == "jsonl":
        write_jsonl(records, path)
    elif suffix == "parquet":
        write_parquet(records, path)
    else:
        raise ValueError(f"Unsupported output format: {output_format}")


def materialize_sample(
    output_path: Path,
    manifest_path: Path,
    *,
    subset: str | None = None,
    split: str = DEFAULT_SPLIT,
    cache_dir: Path | None = None,
    download_dir: Path | None = None,
    use_local_shards: bool = False,
    ensure_download: bool = False,
    allow_full_download: bool = False,
    shuffle_buffer: int = 10_000,
    seed: int = 0,
    num_documents: int | None = None,
    token_budget: int | None = None,
    stratify_by: Sequence[str] | None = None,
    output_format: str = "jsonl",
) -> SampleManifest:
    builder = dolma_builder(subset=subset, cache_dir=cache_dir or download_dir)
    dataset = stream_dolma(
        subset=subset,
        split=split,
        cache_dir=cache_dir,
        download_dir=download_dir,
        use_local_shards=use_local_shards,
        ensure_download=ensure_download,
        allow_full_download=allow_full_download,
        shuffle_buffer=shuffle_buffer,
        seed=seed,
    )

    selected, tokens = sample_documents(
        dataset,
        num_documents=num_documents,
        token_budget=token_budget,
        stratify_by=stratify_by,
    )

    _write_output(selected, output_path, output_format)

    manifest = SampleManifest(
        dataset=builder.info.dataset_name,
        subset=subset,
        split=split,
        version=str(builder.info.version) if builder.info.version else None,
        seed=seed,
        sampling={
            "num_documents": num_documents,
            "token_budget": token_budget,
            "stratify_by": list(stratify_by or []),
            "stratify_strategy": "approximate_streaming_cap" if stratify_by else None,
            "shuffle_buffer": shuffle_buffer,
            "cache_dir": str(cache_dir) if cache_dir else None,
            "download_dir": str(download_dir) if download_dir else None,
            "use_local_shards": use_local_shards,
            "ensure_download": ensure_download,
        },
        document_count=len(selected),
        token_estimate=tokens,
    )
    write_manifest(manifest, manifest_path)
    return manifest


__all__ = ["materialize_sample"]