social-data-attribution-w2/src/dolma/quality/validation/run.py

"""Execution helpers for quality validation commands."""

from __future__ import annotations

import json

from dolma.quality.fasttext import QualityFastTextClassifier
from dolma.quality.r2 import (
    download_object_bytes,
    sidecar_keys_for_source,
    upload_object_bytes,
)
from dolma.quality.sidecar import encode_quality_rows, process_shard_bytes


def encode_json_bytes(payload: dict[str, object]) -> bytes:
    return (json.dumps(payload, indent=2) + "\n").encode("utf-8")


def run_local_smoke(
    client: object,
    *,
    bucket: str,
    output_prefix: str,
    source_keys: list[str],
    model_repo: str,
    batch_size: int,
) -> list[dict[str, object]]:
    classifier = QualityFastTextClassifier(model_repo=model_repo)
    summaries: list[dict[str, object]] = []
    for source_key in source_keys:
        payload = download_object_bytes(client, bucket=bucket, key=source_key)
        rows, stats = process_shard_bytes(
            payload,
            source_key=source_key,
            classifier=classifier,
            label_map=classifier.label_map,
            batch_size=batch_size,
            cpu_count=1,
            model_load_time_seconds=classifier.load_time_seconds,
        )
        keys = sidecar_keys_for_source(source_key, output_prefix)
        upload_object_bytes(
            client,
            bucket=bucket,
            key=keys["parquet"],
            payload=encode_quality_rows(rows),
            content_type="application/octet-stream",
        )
        upload_object_bytes(
            client,
            bucket=bucket,
            key=keys["stats"],
            payload=encode_json_bytes(stats),
            content_type="application/json",
        )
        upload_object_bytes(
            client,
            bucket=bucket,
            key=keys["done"],
            payload=b"done\n",
            content_type="text/plain",
        )
        summaries.append(stats)
    return summaries


__all__ = ["encode_json_bytes", "run_local_smoke"]