scripts/source_normalization_benchmark_pairs.py

#!/usr/bin/env python3
r"""Emit JSONL **pair** candidates for normalization / merge labeling (see ``backend/benchmark_pair_sources``).

Strategies (comma-separated ``--strategies``):

- ``merge_signature`` — same ``(cluster, merge_signature)`` bucket as ``propose_merge_clusters``.
- ``fingerprint_collision`` — same title fingerprint, different merge signature (alias / reject review).
- ``similarity_edge`` — ``similarityedge`` rows above ``--similarity-min-score``.

Example:
  python scripts/source_normalization_benchmark_pairs.py \
    --db data/store_slim.db --out /tmp/bench_pairs.jsonl \
    --strategies merge_signature,fingerprint_collision \
    --sample-random 500 --seed 42
"""
from __future__ import annotations

import argparse
import json
import random
import sqlite3
from pathlib import Path

from backend.benchmark_pair_sources import (
    collect_all_pair_records,
    dedupe_pair_records,
    sample_pair_records,
)


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--db", type=Path, required=True)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument(
        "--strategies",
        default="merge_signature,fingerprint_collision,similarity_edge",
        help="comma list: merge_signature, fingerprint_collision, similarity_edge",
    )
    ap.add_argument("--max-pairs-per-merge-bucket", type=int, default=40)
    ap.add_argument("--max-pairs-per-fingerprint", type=int, default=30)
    ap.add_argument("--similarity-min-score", type=float, default=0.82)
    ap.add_argument("--similarity-max-rows", type=int, default=50_000)
    ap.add_argument("--sample-random", type=int, default=0, help="after dedupe, shuffle and cap (0 = no cap)")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--no-dedupe", action="store_true", help="keep duplicate pair_key rows from multiple tiers")
    args = ap.parse_args()

    db = args.db.resolve()
    if not db.is_file():
        raise SystemExit(f"database not found: {db}")

    strategies = tuple(s.strip() for s in args.strategies.split(",") if s.strip())
    conn = sqlite3.connect(db)
    conn.row_factory = sqlite3.Row
    try:
        rows_sql = conn.execute(
            """
            SELECT k.id, k.name, k.cluster, COALESCE(m.popularity, 0.0) AS pop
            FROM kink k
            LEFT JOIN fetlifekinkmeta m ON m.kink_id = k.id
            """,
        ).fetchall()
        kink_rows = [(str(r["id"]), str(r["name"]), str(r["cluster"]), float(r["pop"] or 0.0)) for r in rows_sql]

        rng = random.Random(args.seed)  # noqa: S311 — reproducible export sampling
        pairs = collect_all_pair_records(
            kink_rows,
            conn if "similarity_edge" in strategies else None,
            strategies=strategies,
            max_pairs_per_merge_bucket=args.max_pairs_per_merge_bucket,
            max_pairs_per_fingerprint=args.max_pairs_per_fingerprint,
            similarity_min_score=args.similarity_min_score,
            similarity_max_pairs=args.similarity_max_rows,
            rng=rng,
        )
        if not args.no_dedupe:
            pairs = dedupe_pair_records(pairs)
        if args.sample_random > 0:
            pairs = sample_pair_records(pairs, args.sample_random, seed=args.seed)
    finally:
        conn.close()

    args.out.parent.mkdir(parents=True, exist_ok=True)
    with args.out.open("w", encoding="utf-8") as fh:
        for rec in pairs:
            fh.write(json.dumps(rec, ensure_ascii=False) + "\n")

    by_s: dict[str, int] = {}
    for rec in pairs:
        s = str(rec.get("strategy", "?"))
        by_s[s] = by_s.get(s, 0) + 1
    print(json.dumps({"wrote": str(args.out), "pair_lines": len(pairs), "by_strategy": by_s}, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())