#!/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()