scripts/score_parquet.py

#!/usr/bin/env python3
"""Score parquet shards with one language's C / A / Q scorers.

One row is one source file; a parquet shard holds many rows. Four columns are
appended to every input table and the result is written under ``--output-root``,
mirroring the input directory structure:

    category         argmax of the file-role classifier C
    cls_confidence   softmax max probability of C
    algo_rel_score   sigmoid output of the algorithmic-relevance model A
    quality_score    fused 0..10 output of the quality model Q

With ``--apply-policy`` a boolean ``selected`` column is added as well, using
the retention rule ``algo_rel_score >= t AND category not in excluded AND
quality_score >= q``. It marks rows rather than dropping them.

Inputs must carry an ``embedding`` column.  A ``relative_path`` column (or a
``meta`` struct with ``file_path``) is required only when the scorers
consume path features — the script checks and refuses to run silently without
them.

This is a single-process scorer.  It walks a directory recursively, but has no
worker pool and no resume ledger, so for corpus-scale runs drive it from your own
scheduler (one invocation per shard or per subtree) rather than pointing it at
millions of rows in one go.
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq

sys.path.insert(0, str(Path(__file__).resolve().parent))

from path_features import extract_relative_paths  # noqa: E402
from scorers import LanguageScorers  # noqa: E402


def collect_inputs(input_path: Path) -> list[Path]:
    if input_path.is_file():
        return [input_path]
    files = sorted(input_path.rglob("*.parquet"))
    if not files:
        raise FileNotFoundError(f"no parquet files under {input_path}")
    return files


def score_table(table, scorers: LanguageScorers, batch_size: int):
    embeddings = table.column("embedding").to_pylist()
    if scorers.needs_paths():
        paths = extract_relative_paths(table, scorers.path_column)
        if not any(paths):
            raise ValueError(
                "these scorers use path features but no 'relative_path' column "
                "(or meta.file_path) was found in the input"
            )
    else:
        paths = [""] * len(embeddings)

    columns = {"category": [], "cls_confidence": [], "algo_rel_score": [], "quality_score": []}
    for start in range(0, len(embeddings), batch_size):
        stop = start + batch_size
        chunk = np.asarray(embeddings[start:stop], dtype=np.float32)
        scored = scorers.score(chunk, paths[start:stop])
        columns["category"].extend(scored["category"])
        for key in ("cls_confidence", "algo_rel_score", "quality_score"):
            columns[key].extend(np.asarray(scored[key]).reshape(-1).tolist())
    return columns


def load_policy_defaults(language: str) -> dict:
    """Read the retention policy defaults for one language from config.json.

    Kept out of the argparse defaults on purpose: the cuts are a policy choice,
    not a property of the models, so they live in one editable place rather than
    baked into this file. ``scorers.py`` never reads this — only the CLI does.

    ``selection_policy`` holds the cuts that apply to every language. An optional
    ``per_language`` block overrides any of them for a single language, so each
    language can carry its own operating point without touching the others::

        "selection_policy": {
          "algo_rel_min": 0.8,
          "per_language": {"r": {"algo_rel_min": 0.9}}
        }

    Precedence is CLI flag > per_language > shared default.
    """
    path = Path(__file__).resolve().parent.parent / "config.json"
    try:
        policy = json.loads(path.read_text())["selection_policy"]
        override = (policy.get("per_language") or {}).get(language, {})
        merged = {**policy, **override}
        return {
            "algo_rel_min": float(merged["algo_rel_min"]),
            "quality_min": float(merged["quality_min"]),
            "exclude_categories": ",".join(merged["exclude_categories"]),
            "source": "per_language" if override else "shared",
        }
    except (OSError, KeyError, ValueError, TypeError) as exc:
        raise SystemExit(
            f"could not read selection_policy from {path}: {exc}. "
            "Pass --algo-rel-min / --quality-min / --exclude-categories explicitly, "
            "or restore the config file."
        )


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--scorers", required=True, help="scorer-set dir for one language, e.g. classifiers/cpp")
    parser.add_argument("--input", help="parquet file or directory of parquet files")
    parser.add_argument("--output-root", help="destination directory for scored parquet")
    parser.add_argument("--batch-size", type=int, default=1024)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--limit", type=int, default=0, help="score at most N input parquet files (0 = all)")
    parser.add_argument("--describe", action="store_true", help="print the scorer-set configuration and exit")
    parser.add_argument("--apply-policy", action="store_true", help="also write a boolean 'selected' column")
    parser.add_argument("--algo-rel-min", type=float, help="default: config.json selection_policy")
    parser.add_argument("--quality-min", type=float, help="default: config.json selection_policy")
    parser.add_argument("--exclude-categories", help="comma separated categories to veto; default: config.json")
    args = parser.parse_args()

    scorers = LanguageScorers.load(args.scorers, device=args.device)

    if args.describe:
        print(json.dumps(scorers.describe(), indent=2, ensure_ascii=False))
        return

    if not args.input or not args.output_root:
        parser.error("--input and --output-root are required unless --describe is given")

    input_path = Path(args.input)
    output_root = Path(args.output_root)
    files = collect_inputs(input_path)
    if args.limit:
        files = files[: args.limit]
    root = input_path if input_path.is_dir() else input_path.parent

    # Resolve the retention cuts only when they are actually needed, so a plain
    # scoring run does not fail on a missing/edited config.json.
    algo_rel_min = quality_min = None
    excluded = set()
    if args.apply_policy:
        defaults = load_policy_defaults(scorers.language)
        algo_rel_min = (
            args.algo_rel_min if args.algo_rel_min is not None
            else defaults["algo_rel_min"]
        )
        quality_min = (
            args.quality_min if args.quality_min is not None
            else defaults["quality_min"]
        )
        categories = (
            args.exclude_categories if args.exclude_categories is not None
            else defaults["exclude_categories"]
        )
        excluded = {c.strip() for c in categories.split(",") if c.strip()}
        print(
            f"policy [{scorers.language}, {defaults['source']} defaults]: "
            f"algo_rel_score >= {algo_rel_min} AND category not in "
            f"{sorted(excluded)} AND quality_score >= {quality_min}",
            flush=True,
        )

    total_rows = 0
    for index, in_path in enumerate(files, 1):
        started = time.time()
        table = pq.read_table(in_path)
        columns = score_table(table, scorers, args.batch_size)

        out = table
        for name, values, arrow_type in (
            ("category", columns["category"], pa.string()),
            ("cls_confidence", columns["cls_confidence"], pa.float32()),
            ("algo_rel_score", columns["algo_rel_score"], pa.float32()),
            ("quality_score", columns["quality_score"], pa.float32()),
        ):
            if name in out.column_names:
                out = out.drop([name])
            out = out.append_column(name, pa.array(values, type=arrow_type))

        if args.apply_policy:
            selected = [
                bool(
                    relevance >= algo_rel_min
                    and category not in excluded
                    and quality >= quality_min
                )
                for relevance, category, quality in zip(
                    columns["algo_rel_score"], columns["category"], columns["quality_score"]
                )
            ]
            if "selected" in out.column_names:
                out = out.drop(["selected"])
            out = out.append_column("selected", pa.array(selected, type=pa.bool_()))

        out_path = output_root / in_path.relative_to(root)
        out_path.parent.mkdir(parents=True, exist_ok=True)
        pq.write_table(out, out_path)
        total_rows += table.num_rows
        print(
            f"[{index}/{len(files)}] {in_path.name}: {table.num_rows} rows "
            f"in {time.time() - started:.1f}s -> {out_path}",
            flush=True,
        )

    print(f"done: {len(files)} parquet files, {total_rows} rows -> {output_root}")


if __name__ == "__main__":
    main()