spb/llmindex Public

1#!/usr/bin/env python32# llmindex.io — fit CLI: response matrices JSON → fitted 2PL parameters JSON3# Author:  Simon-Pierre Boucher4# Contact: contact@spboucher.ai5# License: Proprietary — © Simon-Pierre Boucher, all rights reserved6#7# Pure compute: reads the matrix file written by the worker (no DB access),8# fits a 2PL per domain, writes thetas/SEs and item parameters back to JSON.910from __future__ import annotations1112import argparse13import json14import sys1516import numpy as np1718from llmindex_psycho.bt import fit_bradley_terry19from llmindex_psycho.irt import fit_2pl202122def main() -> int:23    parser = argparse.ArgumentParser(description="llmindex 2PL fit")24    parser.add_argument("--input", required=True)25    parser.add_argument("--output", required=True)26    args = parser.parse_args()2728    with open(args.input) as f:29        payload = json.load(f)3031    hyper = payload.get("hyperparams", {})32    max_iter = int(hyper.get("maxIterations", 500))33    tol = float(hyper.get("tolerance", 1e-6))3435    out: dict = {"domains": {}}36    for domain, data in payload.get("domains", {}).items():37        models = data["models"]38        items = data["items"]39        matrix = np.array(40            [[np.nan if v is None else float(v) for v in row] for row in data["responses"]],41            dtype=float,42        )43        result = fit_2pl(matrix, max_iterations=max_iter, tolerance=tol)44        out["domains"][domain] = {45            "models": [46                {"slug": slug, "theta": float(t), "se": float(se)}47                for slug, t, se in zip(models, result.theta, result.theta_se)48            ],49            "items": [50                {"id": item_id, "a": float(a), "b": float(b)}51                for item_id, a, b in zip(items, result.a, result.b)52            ],53            "diagnostics": {54                "iterations": result.iterations,55                "converged": result.converged,56                "final_loglik": result.final_loglik,57                **result.diagnostics,58            },59        }60        print(61            f"[fit] {domain}: {len(models)} models × {len(items)} items — "62            f"{'converged' if result.converged else 'max iterations'} @ {result.iterations}",63            file=sys.stderr,64        )6566    # Judged duel domains: Bradley-Terry over wins matrices; log-strengths are67    # standardized to a theta-like scale so downstream scoring is uniform.68    out["duel_domains"] = {}69    for domain, data in payload.get("duel_domains", {}).items():70        models = data["models"]71        wins = np.array(data["wins"], dtype=float)72        result = fit_bradley_terry(wins)73        sd = result.log_strength.std() or 1.074        theta = result.log_strength / sd75        theta_se = result.log_strength_se / sd76        out["duel_domains"][domain] = {77            "models": [78                {"slug": slug, "theta": float(t), "se": float(min(se, 3.0))}79                for slug, t, se in zip(models, theta, theta_se)80            ],81            "diagnostics": {82                "iterations": result.iterations,83                "converged": result.converged,84                "comparisons": float(wins.sum()),85            },86        }87        print(f"[fit] duel {domain}: {len(models)} models, {wins.sum():.0f} comparisons", file=sys.stderr)8889    with open(args.output, "w") as f:90        json.dump(out, f, indent=2)91    return 0929394if __name__ == "__main__":95    raise SystemExit(main())96