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

"""Stratified sampling for the 6T working sample (SOC-134).

Supports two modes:
- Token floor: draw docs from each bin until a minimum token count is reached
- Docs per bin: draw a fixed number of docs from each bin
"""

from __future__ import annotations

import json
import logging
import random
from dataclasses import dataclass, field
from pathlib import Path

import pandas as pd

from dolma.constants import (
    FORMATS,
    TOPICS,
    WORKING_SAMPLE_SAMPLING_SEED,
    WORKING_SAMPLE_TOKEN_FLOOR_PER_BIN,
)

logger = logging.getLogger(__name__)

NUM_BINS = len(TOPICS) * len(FORMATS)


@dataclass
class BinResult:
    topic: str
    fmt: str
    bin_id: int
    requested_floor: int
    requested_docs: int
    realized_tokens: int
    realized_docs: int
    available_tokens: int
    available_docs: int
    underfilled: bool


@dataclass
class SampleResult:
    sample_df: pd.DataFrame
    bin_results: list[BinResult] = field(default_factory=list)
    token_floor_per_bin: int = 0
    docs_per_bin: int = 0
    seed: int = 0
    global_token_budget: int | None = None
    min_token_count: int | None = None
    max_token_count: int | None = None
    exclude_manifest_paths: list[str] = field(default_factory=list)
    excluded_doc_count: int = 0


def _strip_label_prefix(value: str) -> str:
    if value.startswith("__label__"):
        return value[len("__label__") :]
    return value


def _compute_bin_id(topic_idx: int, format_idx: int) -> int:
    return topic_idx * len(FORMATS) + format_idx + 1


def draw_working_sample(
    manifest: pd.DataFrame,
    token_floor_per_bin: int = WORKING_SAMPLE_TOKEN_FLOOR_PER_BIN,
    docs_per_bin: int | None = None,
    seed: int = WORKING_SAMPLE_SAMPLING_SEED,
    global_token_budget: int | None = None,
    min_token_count: int | None = None,
    max_token_count: int | None = None,
) -> SampleResult:
    if (
        docs_per_bin is not None
        and token_floor_per_bin != WORKING_SAMPLE_TOKEN_FLOOR_PER_BIN
    ):
        raise ValueError("Specify docs_per_bin or token_floor_per_bin, not both.")

    use_docs_mode = docs_per_bin is not None
    rng = random.Random(seed)

    topic_col = "topic" if "topic" in manifest.columns else "topic_label"
    format_col = "format" if "format" in manifest.columns else "format_label"
    token_col = "token_count"

    df = manifest.copy()
    df["_topic"] = df[topic_col].map(_strip_label_prefix)
    df["_format"] = df[format_col].map(_strip_label_prefix)

    sampled_parts: list[pd.DataFrame] = []
    bin_results: list[BinResult] = []
    total_sampled_tokens = 0

    for t_idx, topic in enumerate(TOPICS):
        for f_idx, fmt in enumerate(FORMATS):
            bin_id = _compute_bin_id(t_idx, f_idx)
            bin_df = df[(df["_topic"] == topic) & (df["_format"] == fmt)]
            available_docs = len(bin_df)
            available_tokens = int(bin_df[token_col].sum()) if available_docs > 0 else 0

            if available_docs == 0:
                bin_results.append(
                    BinResult(
                        topic=topic,
                        fmt=fmt,
                        bin_id=bin_id,
                        requested_floor=0 if use_docs_mode else token_floor_per_bin,
                        requested_docs=docs_per_bin or 0,
                        realized_tokens=0,
                        realized_docs=0,
                        available_tokens=0,
                        available_docs=0,
                        underfilled=True,
                    )
                )
                continue

            shuffled = bin_df.sample(frac=1, random_state=rng.randint(0, 2**31))

            if use_docs_mode:
                take = min(docs_per_bin, available_docs)
                selected = shuffled.iloc[:take]
                underfilled = available_docs < docs_per_bin
            else:
                cumulative = shuffled[token_col].cumsum()
                meets_floor = cumulative >= token_floor_per_bin
                if meets_floor.any():
                    cutoff_idx = meets_floor.idxmax()
                    cutoff_pos = shuffled.index.get_loc(cutoff_idx)
                    selected = shuffled.iloc[: cutoff_pos + 1]
                else:
                    selected = shuffled
                underfilled = int(selected[token_col].sum()) < token_floor_per_bin

            realized_tokens = int(selected[token_col].sum())

            selected = selected.copy()
            selected["bin_id"] = bin_id
            selected["bin_topic"] = topic
            selected["bin_format"] = fmt
            sampled_parts.append(selected)
            total_sampled_tokens += realized_tokens

            bin_results.append(
                BinResult(
                    topic=topic,
                    fmt=fmt,
                    bin_id=bin_id,
                    requested_floor=0 if use_docs_mode else token_floor_per_bin,
                    requested_docs=docs_per_bin or 0,
                    realized_tokens=realized_tokens,
                    realized_docs=len(selected),
                    available_tokens=available_tokens,
                    available_docs=available_docs,
                    underfilled=underfilled,
                )
            )

    if sampled_parts:
        sample_df = pd.concat(sampled_parts, ignore_index=True)
    else:
        sample_df = pd.DataFrame()

    sample_df.drop(columns=["_topic", "_format"], inplace=True, errors="ignore")

    underfilled_count = sum(1 for br in bin_results if br.underfilled)
    covered_count = sum(1 for br in bin_results if br.realized_docs > 0)

    mode_desc = (
        f"{docs_per_bin} docs/bin"
        if use_docs_mode
        else f"{token_floor_per_bin:,} token floor"
    )
    logger.info(
        "Working sample (%s): %s docs, %s tokens across %d/%d bins (%d underfilled)",
        mode_desc,
        f"{len(sample_df):,}",
        f"{total_sampled_tokens:,}",
        covered_count,
        NUM_BINS,
        underfilled_count,
    )

    return SampleResult(
        sample_df=sample_df,
        bin_results=bin_results,
        token_floor_per_bin=0 if use_docs_mode else token_floor_per_bin,
        docs_per_bin=docs_per_bin or 0,
        seed=seed,
        global_token_budget=global_token_budget,
        min_token_count=min_token_count,
        max_token_count=max_token_count,
    )


def bin_summary_dataframe(result: SampleResult) -> pd.DataFrame:
    rows = []
    for br in result.bin_results:
        rows.append(
            {
                "bin_id": br.bin_id,
                "topic": br.topic,
                "format": br.fmt,
                "requested_floor_tokens": br.requested_floor,
                "requested_docs_per_bin": br.requested_docs,
                "realized_tokens": br.realized_tokens,
                "realized_docs": br.realized_docs,
                "available_tokens": br.available_tokens,
                "available_docs": br.available_docs,
                "underfilled": br.underfilled,
                "shortfall_tokens": max(0, br.requested_floor - br.realized_tokens),
                "shortfall_docs": max(0, br.requested_docs - br.realized_docs),
            }
        )
    return pd.DataFrame(rows)


def sample_contract(result: SampleResult) -> dict:
    underfilled = sum(1 for br in result.bin_results if br.underfilled)
    covered = sum(1 for br in result.bin_results if br.realized_docs > 0)
    realized_tokens = sum(br.realized_tokens for br in result.bin_results)
    realized_docs = sum(br.realized_docs for br in result.bin_results)

    contract = {
        "WORKING_SAMPLE_TOKEN_FLOOR_PER_BIN": result.token_floor_per_bin,
        "WORKING_SAMPLE_DOCS_PER_BIN": result.docs_per_bin,
        "WORKING_SAMPLE_GLOBAL_TOKEN_BUDGET": result.global_token_budget,
        "WORKING_SAMPLE_MIN_TOKEN_COUNT": result.min_token_count,
        "WORKING_SAMPLE_MAX_TOKEN_COUNT": result.max_token_count,
        "WORKING_SAMPLE_REALIZED_TOKEN_TOTAL": realized_tokens,
        "WORKING_SAMPLE_REALIZED_DOC_COUNT": realized_docs,
        "WORKING_SAMPLE_UNDERFILLED_BIN_COUNT": underfilled,
        "WORKING_SAMPLE_COVERED_BIN_COUNT": covered,
        "WORKING_SAMPLE_TOTAL_BIN_COUNT": NUM_BINS,
        "WORKING_SAMPLE_SAMPLING_SEED": result.seed,
    }
    if result.exclude_manifest_paths:
        contract["WORKING_SAMPLE_EXCLUDE_MANIFEST_PATHS"] = (
            result.exclude_manifest_paths
        )
        contract["WORKING_SAMPLE_EXCLUDED_DOC_COUNT"] = result.excluded_doc_count
    return contract


def write_outputs(
    result: SampleResult,
    output_dir: Path,
) -> None:
    output_dir.mkdir(parents=True, exist_ok=True)

    manifest_path = output_dir / "working_sample_manifest.parquet"
    result.sample_df.to_parquet(manifest_path, index=False)
    logger.info("Wrote sample manifest: %s", manifest_path)

    summary_df = bin_summary_dataframe(result)
    summary_csv = output_dir / "bin_summary.csv"
    summary_df.to_csv(summary_csv, index=False)
    logger.info("Wrote bin summary: %s", summary_csv)

    contract = sample_contract(result)
    contract_path = output_dir / "sample_contract.json"
    contract_path.write_text(json.dumps(contract, indent=2) + "\n")
    logger.info("Wrote sample contract: %s", contract_path)
    for key, value in contract.items():
        logger.info("  %s = %s", key, f"{value:,}" if isinstance(value, int) else value)