social-data-attribution-w2/scripts/diagnostics/sample_truncated_docs.py

"""Phase 2: Sample truncated+high-quality docs from R2 and re-classify.

Pulls actual docs classified as 'truncated' format with quality_score >= 0.2,
validates doc_id alignment between SOC-91 and SOC-139 sidecars, and
re-classifies each doc locally to check for score drift.

Requires R2 credentials via with_r2_credentials.sh.
"""

from __future__ import annotations

import argparse
import json
import logging
import sys
from pathlib import Path

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

R2_BUCKET = "soc127-dedup"
R2_ENDPOINT_URL = "https://0934ab8e84ac8f4e81decaf3eb121337.r2.cloudflarestorage.com"
R2_INPUT_PREFIXES = [
    "soc127/phase1_pool_shared",
    "soc127/phase2_nonpool_final",
]
SOC91_PREFIX = "soc91-labels"
SOC139_PREFIX = "soc139-quality-sidecars"
QUALITY_SCORE_THRESHOLD = 0.2


def find_shards_with_both_sidecars(
    client, *, bucket: str, shard_count: int
) -> list[str]:
    from dolma.quality.r2 import list_keys

    soc91_done = set(
        list_keys(client, bucket=bucket, prefix=SOC91_PREFIX, suffix=".parquet")
    )
    soc139_done = set(
        list_keys(client, bucket=bucket, prefix=SOC139_PREFIX, suffix=".parquet")
    )

    soc91_basenames = {Path(k).stem for k in soc91_done}
    soc139_basenames = {Path(k).stem for k in soc139_done}
    common_basenames = soc91_basenames & soc139_basenames

    source_keys: list[str] = []
    for prefix in R2_INPUT_PREFIXES:
        source_keys.extend(
            list_keys(client, bucket=bucket, prefix=prefix, suffix=".jsonl.zst")
        )

    matched: list[str] = []
    for key in sorted(source_keys):
        basename = Path(key).name.removesuffix(".jsonl.zst")
        if basename in common_basenames:
            matched.append(key)
            if len(matched) >= shard_count:
                break

    log.info(
        "Found %d shards with both SOC-91 and SOC-139 sidecars (of %d requested)",
        len(matched),
        shard_count,
    )
    return matched


def process_shard(
    client, classifier, *, bucket: str, source_key: str
) -> tuple[list[dict], dict]:
    from dolma.quality.validation.io import (
        read_quality_rows,
        read_raw_doc_map,
        read_soc91_doc_map,
    )

    quality_rows = read_quality_rows(
        client,
        bucket=bucket,
        source_key=source_key,
        output_prefix=SOC139_PREFIX,
    )
    soc91_map = read_soc91_doc_map(
        client,
        bucket=bucket,
        source_key=source_key,
        soc91_prefix=SOC91_PREFIX,
    )
    raw_map = read_raw_doc_map(client, bucket=bucket, source_key=source_key)

    quality_ids = {str(r["doc_id"]) for r in quality_rows}
    soc91_ids = set(soc91_map.keys()) if soc91_map else set()

    join_stats = {
        "source_key": source_key,
        "quality_doc_count": len(quality_ids),
        "soc91_doc_count": len(soc91_ids),
        "intersection": len(quality_ids & soc91_ids),
        "quality_only": len(quality_ids - soc91_ids),
        "soc91_only": len(soc91_ids - quality_ids),
    }

    samples: list[dict] = []
    for row in quality_rows:
        doc_id = str(row["doc_id"])
        score = float(row["quality_score"])
        soc91_doc = soc91_map.get(doc_id) if soc91_map else None
        format_label = soc91_doc.get("format_url_label") if soc91_doc else None
        if format_label != "truncated" or score < QUALITY_SCORE_THRESHOLD:
            continue

        raw_doc = raw_map.get(doc_id, {})
        text = raw_doc.get("text_snippet", "")

        reclass_labels, reclass_probs = classifier.model.predict(
            text.replace("\n", " ").strip(), k=-1
        )
        reclass_score = 0.0
        for label, prob in zip(reclass_labels, reclass_probs):
            if label == "__label__1":
                reclass_score = float(prob)

        samples.append(
            {
                "doc_id": doc_id,
                "source_key": source_key,
                "quality_score": score,
                "reclassified_score": reclass_score,
                "score_delta": abs(reclass_score - score),
                "format_url_label": format_label,
                "topic_url_label": soc91_doc.get("topic_url_label")
                if soc91_doc
                else None,
                "text_first_2000": text[:2000],
                "url": raw_doc.get("url"),
            }
        )

    return samples, join_stats


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Phase 2: Sample truncated docs from R2"
    )
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--shard-count", type=int, default=10)
    parser.add_argument("--sample-size", type=int, default=50)
    args = parser.parse_args()

    import os
    from dolma.quality.fasttext import QualityFastTextClassifier
    from dolma.quality.r2 import create_r2_client, R2Config

    config = R2Config(
        endpoint_url=R2_ENDPOINT_URL,
        bucket=R2_BUCKET,
        access_key_id=os.environ["R2_ACCESS_KEY_ID"],
        secret_access_key=os.environ["R2_SECRET_ACCESS_KEY"],
        input_prefixes=tuple(R2_INPUT_PREFIXES),
        output_prefix=SOC139_PREFIX,
    )
    client = create_r2_client(config)

    log.info("Loading FastText quality model...")
    classifier = QualityFastTextClassifier()
    log.info("Model loaded in %.2fs", classifier.load_time_seconds)

    shards = find_shards_with_both_sidecars(
        client,
        bucket=R2_BUCKET,
        shard_count=args.shard_count,
    )
    if not shards:
        log.error("No shards found with both sidecars.")
        sys.exit(1)

    all_samples: list[dict] = []
    all_join_stats: list[dict] = []
    reclassification_mismatches = 0

    for i, source_key in enumerate(shards, 1):
        log.info("[%d/%d] Processing %s", i, len(shards), source_key)
        samples, join_stats = process_shard(
            client,
            classifier,
            bucket=R2_BUCKET,
            source_key=source_key,
        )
        all_join_stats.append(join_stats)
        log.info(
            "  join: %d quality, %d soc91, %d common, %d quality-only, %d soc91-only",
            join_stats["quality_doc_count"],
            join_stats["soc91_doc_count"],
            join_stats["intersection"],
            join_stats["quality_only"],
            join_stats["soc91_only"],
        )

        for s in samples:
            if s["score_delta"] > 1e-3:
                reclassification_mismatches += 1

        remaining = args.sample_size - len(all_samples)
        all_samples.extend(samples[:remaining])
        log.info(
            "  found %d truncated+high docs, total collected: %d/%d",
            len(samples),
            len(all_samples),
            args.sample_size,
        )

        if len(all_samples) >= args.sample_size:
            break

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

    samples_path = args.output_dir / "truncated_samples.jsonl"
    with open(samples_path, "w") as f:
        for s in all_samples:
            f.write(json.dumps(s) + "\n")
    log.info("Wrote %d samples to %s", len(all_samples), samples_path)

    join_path = args.output_dir / "join_stats.json"
    with open(join_path, "w") as f:
        json.dump(all_join_stats, f, indent=2)
    log.info("Wrote join stats to %s", join_path)

    log.info("\n=== Summary ===")
    log.info("Shards processed: %d", len(all_join_stats))
    log.info("Truncated+high samples: %d", len(all_samples))
    log.info(
        "Reclassification mismatches (delta > 1e-3): %d", reclassification_mismatches
    )

    total_quality = sum(j["quality_doc_count"] for j in all_join_stats)
    total_soc91 = sum(j["soc91_doc_count"] for j in all_join_stats)
    total_intersection = sum(j["intersection"] for j in all_join_stats)
    total_quality_only = sum(j["quality_only"] for j in all_join_stats)
    total_soc91_only = sum(j["soc91_only"] for j in all_join_stats)

    log.info(
        "Aggregate join: quality=%d soc91=%d intersection=%d quality_only=%d soc91_only=%d",
        total_quality,
        total_soc91,
        total_intersection,
        total_quality_only,
        total_soc91_only,
    )

    if total_quality > 0:
        join_rate = total_intersection / total_quality
        log.info("Join rate (quality docs matched in SOC-91): %.4f", join_rate)
        if join_rate < 0.95:
            log.info(
                "WARNING: Join rate below 95%% - hypothesis C (join mismatch) gains support"
            )


if __name__ == "__main__":
    main()