[tab] /home/franz.srambical/slurm/jobs/franz/berlin/crowd-pilot/preference_utils.py (python)
#!/usr/bin/env python3
"""
Utility functions for analyzing preference data from .crowd-pilot-preferences.jsonl.

Two modes:
1. Calculate minAvgLogprob threshold for target acceptance rate
2. Calculate overall acceptance rate

Usage:
  # Calculate threshold for 25% acceptance rate
  python preference_utils.py --mode threshold --path .crowd-pilot-preferences.jsonl --target 0.25
  
  # Calculate overall acceptance rate
  python preference_utils.py --mode acceptance-rate --path .crowd-pilot-preferences.jsonl
"""

import argparse
import json
import math
from dataclasses import dataclass
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np


@dataclass
class Row:
    avglogprob: float
    accepted: int  # 1 if accepted else 0
    outcome: str


def load_rows(path: Path) -> List[Row]:
    rows: List[Row] = []
    bad = 0
    total = 0
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            total += 1
            line = line.strip()
            if not line:
                continue
            try:
                obj = json.loads(line)
            except Exception:
                bad += 1
                continue

            outcome = obj.get("outcome")
            if outcome not in ("accepted", "rejected", "ignored"):
                # skip null / unknown outcomes
                continue

            # Exclude "ignored" samples from acceptance rate calculation
            # (they represent cases where user didn't provide explicit feedback)
            if outcome == "ignored":
                continue

            avg = obj.get("completion", {}).get("avgLogprob")
            if not isinstance(avg, (int, float)):
                continue
            avg = float(avg)
            if math.isnan(avg) or math.isinf(avg):
                continue

            accepted = 1 if outcome == "accepted" else 0
            rows.append(Row(avglogprob=avg, accepted=accepted, outcome=outcome))

    if len(rows) == 0:
        raise SystemExit(
            f"No usable rows found in {path}. "
            f"Check that the file exists and contains outcome + completion.avgLogprob."
        )
    return rows

def calculate_mean_acceptance_rate(rows: List[Row]) -> float:
    """
    Calculate the mean acceptance rate across all rows.
    Returns: accepted / (accepted + rejected)
    Note: rows should already have "ignored" samples excluded.
    """
    acc = np.array([r.accepted for r in rows], dtype=np.float64)
    return acc.mean()


def find_threshold_for_target(rows: List[Row], target: float) -> Tuple[Optional[float], dict]:
    """
    Returns (tau, stats). tau is the lowest-confidence avglogprob threshold (smallest avglogprob)
    such that cumulative acceptance rate among rows with avglogprob >= tau is >= target.

    If even the best (highest-confidence only) can't reach target, tau=None.
    """
    rows_sorted = sorted(rows, key=lambda r: r.avglogprob, reverse=True)  # high->low confidence
    avg = np.array([r.avglogprob for r in rows_sorted], dtype=np.float64)
    acc = np.array([r.accepted for r in rows_sorted], dtype=np.int32)

    cum_acc = np.cumsum(acc)
    cum_n = np.arange(1, len(acc) + 1, dtype=np.int64)
    cum_rate = cum_acc / cum_n

    ok = np.where(cum_rate >= target)[0]
    if ok.size == 0:
        # cannot achieve target even with only highest-confidence samples
        stats = {
            "max_rate": float(cum_rate.max()),
            "max_rate_at_n": int(cum_n[cum_rate.argmax()]),
            "total_rows": int(len(rows_sorted)),
        }
        return None, stats

    # Choose the largest n such that cum_rate[n-1] >= target => last index in ok.
    idx = int(ok[-1])
    tau = float(avg[idx])

    stats = {
        "tau": tau,
        "shown_n": int(cum_n[idx]),
        "accepted_n": int(cum_acc[idx]),
        "achieved_rate": float(cum_rate[idx]),
        "total_rows": int(len(rows_sorted)),
        "overall_accept_rate": float(acc.mean()),
    }
    return tau, stats


def nearby_table(rows: List[Row], tau: float, target: float, k: int = 5) -> List[dict]:
    """
    Produce a small table around the chosen threshold: a few stricter (higher tau) and
    looser (lower tau) thresholds with resulting accept rates.
    """
    rows_sorted = sorted(rows, key=lambda r: r.avglogprob, reverse=True)
    avg = np.array([r.avglogprob for r in rows_sorted], dtype=np.float64)
    acc = np.array([r.accepted for r in rows_sorted], dtype=np.int32)

    mask = avg >= tau
    if mask.any():
        idx = int(np.where(mask)[0][-1])
    else:
        # tau is higher than all values, use index 0 (strictest threshold)
        idx = 0

    # pick indices around it
    indices = [i for i in range(max(0, idx - k), min(len(rows_sorted), idx + k + 1))]
    table = []
    for i in indices:
        thr = float(avg[i])
        mask = avg >= thr
        n = int(mask.sum())
        a = int(acc[mask].sum())
        rate = a / n if n else 0.0
        table.append(
            {
                "threshold_avgLogprob": thr,
                "shown_n": n,
                "accepted_n": a,
                "accept_rate": rate,
                "meets_target": rate >= target,
            }
        )
    return table


def mode_threshold(rows: List[Row], target: float, nearby: int = 5):
    """Calculate minAvgLogprob threshold for target acceptance rate."""
    tau, stats = find_threshold_for_target(rows, target)

    print(f"Loaded {len(rows)} usable samples")
    print(f"Overall accept rate (on shown samples in file): {stats.get('overall_accept_rate', float('nan')):.4f}")

    if tau is None:
        print("
No threshold can achieve the target acceptance rate.")
        print(f"Max achievable cumulative rate: {stats['max_rate']:.4f} at n={stats['max_rate_at_n']}")
        return

    print("
=== Recommended minAvgLogprob ===")
    print(f"Target acceptance rate: {target:.2f}")
    print(f"minAvgLogprob (tau):   {stats['tau']:.6f}")
    print(f"Shown samples (>=tau): {stats['shown_n']} / {stats['total_rows']}")
    print(f"Accepted among shown:  {stats['accepted_n']}")
    print(f"Achieved accept rate:  {stats['achieved_rate']:.4f}")

    print("
=== Nearby thresholds (sanity check) ===")
    table = nearby_table(rows, tau=stats["tau"], target=target, k=nearby)
    # print nicely
    header = f"{'thr(avgLogprob)':>16}  {'shown':>6}  {'acc':>6}  {'rate':>7}  {'>=target':>9}"
    print(header)
    print("-" * len(header))
    for r in table:
        print(
            f"{r['threshold_avgLogprob']:16.6f}  "
            f"{r['shown_n']:6d}  "
            f"{r['accepted_n']:6d}  "
            f"{r['accept_rate']:7.4f}  "
            f"{str(r['meets_target']):>9}"
        )


def mode_acceptance_rate(rows: List[Row]):
    """Calculate and display overall acceptance rate."""
    acceptance_rate = calculate_mean_acceptance_rate(rows)
    accepted_count = sum(1 for r in rows if r.accepted == 1)
    rejected_count = sum(1 for r in rows if r.accepted == 0)
    total_count = len(rows)

    print(f"Loaded {total_count} usable samples (accepted + rejected, ignored excluded)")
    print(f"
=== Acceptance Rate ===")
    print(f"Accepted: {accepted_count}")
    print(f"Rejected: {rejected_count}")
    print(f"Total:    {total_count}")
    print(f"Acceptance rate: {acceptance_rate:.4f} ({acceptance_rate*100:.2f}%)")


def main():
    ap = argparse.ArgumentParser(
        description="Analyze preference data from crowd-pilot preference logs."
    )
    ap.add_argument("--path", type=str, default=".crowd-pilot-preferences.jsonl",
                    help="Path to preference log JSONL.")
    ap.add_argument("--mode", type=str, choices=["threshold", "acceptance-rate"],
                    default="threshold",
                    help="Mode: 'threshold' to find minAvgLogprob, 'acceptance-rate' to calculate overall rate.")
    ap.add_argument("--target", type=float, default=0.25,
                    help="Target acceptance rate among shown suggestions (for threshold mode).")
    ap.add_argument("--nearby", type=int, default=5,
                    help="How many thresholds above/below to show in the nearby table (for threshold mode).")
    args = ap.parse_args()

    path = Path(args.path)
    if not path.exists():
        raise SystemExit(f"File not found: {path}")

    rows = load_rows(path)

    if args.mode == "threshold":
        mode_threshold(rows, args.target, args.nearby)
    elif args.mode == "acceptance-rate":
        mode_acceptance_rate(rows)


if __name__ == "__main__":
    main()