294 lines
10 KiB
Python
Executable file
294 lines
10 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
"""Build a candidate matrix from retrieval optimization JSONL artifacts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import math
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
|
|
@dataclass
|
|
class Row:
|
|
candidate: str
|
|
source_file: str
|
|
line_number: int
|
|
success: bool
|
|
objective: float
|
|
hybrid_precision: float
|
|
hybrid_mrr: float
|
|
hybrid_ndcg: float
|
|
hybrid_recall: float
|
|
hybrid_map: float
|
|
hybrid_eval_ms: float
|
|
keyword_eval_ms: float
|
|
hybrid_keyword_mrr_delta: float
|
|
hybrid_duplicate_rate_at_k: float
|
|
objective_penalty_hybrid_regression: float
|
|
objective_penalty_duplicate_rate: float
|
|
trace_coverage: float
|
|
graph_rerank_apply_rate: float
|
|
semantic_rescue_rate_mean: float
|
|
vector_only_below_threshold_mean: float
|
|
query_without_relevant_hit_rate: float
|
|
trace_top_n: float
|
|
trace_component_top_n: float
|
|
tuning_state: str
|
|
tuning_reason: str
|
|
env_overrides: dict[str, Any]
|
|
|
|
|
|
def _safe_float(value: Any, default: float = 0.0) -> float:
|
|
try:
|
|
parsed = float(value)
|
|
if math.isnan(parsed) or math.isinf(parsed):
|
|
return default
|
|
return parsed
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
def _is_success(row: dict[str, Any]) -> bool:
|
|
if isinstance(row.get("success"), bool):
|
|
return bool(row.get("success"))
|
|
status = str(row.get("status", "")).strip().lower()
|
|
return status == "ok"
|
|
|
|
|
|
def _objective(row: dict[str, Any]) -> float:
|
|
if "objective" in row:
|
|
return _safe_float(row.get("objective"), 0.0)
|
|
if "objective_score" in row:
|
|
return _safe_float(row.get("objective_score"), 0.0)
|
|
return 0.0
|
|
|
|
|
|
def _nested(data: dict[str, Any], path: list[str], default: Any = None) -> Any:
|
|
cur: Any = data
|
|
for key in path:
|
|
if not isinstance(cur, dict):
|
|
return default
|
|
if key not in cur:
|
|
return default
|
|
cur = cur[key]
|
|
return cur
|
|
|
|
|
|
def _load_rows(inputs: list[Path]) -> list[Row]:
|
|
out: list[Row] = []
|
|
for path in inputs:
|
|
if not path.exists():
|
|
raise FileNotFoundError(f"Input JSONL not found: {path}")
|
|
|
|
with path.open("r", encoding="utf-8") as handle:
|
|
for line_number, line in enumerate(handle, start=1):
|
|
line = line.strip()
|
|
if not line:
|
|
continue
|
|
data = json.loads(line)
|
|
if not isinstance(data, dict):
|
|
continue
|
|
|
|
candidate = str(data.get("candidate", "")).strip()
|
|
if not candidate:
|
|
continue
|
|
|
|
hybrid = data.get("hybrid", {})
|
|
hybrid_debug = data.get("hybrid_debug_summary", {})
|
|
objective_penalties = data.get("objective_penalties", {})
|
|
out.append(
|
|
Row(
|
|
candidate=candidate,
|
|
source_file=str(path),
|
|
line_number=line_number,
|
|
success=_is_success(data),
|
|
objective=_objective(data),
|
|
hybrid_precision=_safe_float(hybrid.get("precision_at_k"), 0.0),
|
|
hybrid_mrr=_safe_float(hybrid.get("mrr"), 0.0),
|
|
hybrid_ndcg=_safe_float(hybrid.get("ndcg_at_k"), 0.0),
|
|
hybrid_recall=_safe_float(hybrid.get("recall_at_k"), 0.0),
|
|
hybrid_map=_safe_float(hybrid.get("map"), 0.0),
|
|
hybrid_duplicate_rate_at_k=_safe_float(
|
|
hybrid.get("duplicate_rate_at_k"), 0.0
|
|
),
|
|
hybrid_eval_ms=_safe_float(data.get("hybrid_eval_ms"), 0.0),
|
|
keyword_eval_ms=_safe_float(data.get("keyword_eval_ms"), 0.0),
|
|
hybrid_keyword_mrr_delta=_safe_float(
|
|
data.get("hybrid_keyword_mrr_delta"), 0.0
|
|
),
|
|
objective_penalty_hybrid_regression=_safe_float(
|
|
objective_penalties.get("hybrid_regression"), 0.0
|
|
),
|
|
objective_penalty_duplicate_rate=_safe_float(
|
|
objective_penalties.get("duplicate_rate"), 0.0
|
|
),
|
|
trace_coverage=_safe_float(
|
|
hybrid_debug.get("trace_coverage"), 0.0
|
|
),
|
|
graph_rerank_apply_rate=_safe_float(
|
|
hybrid_debug.get("graph_rerank_apply_rate"), 0.0
|
|
),
|
|
semantic_rescue_rate_mean=_safe_float(
|
|
_nested(
|
|
hybrid_debug, ["semantic_rescue_rate", "mean"], 0.0
|
|
),
|
|
0.0,
|
|
),
|
|
vector_only_below_threshold_mean=_safe_float(
|
|
_nested(
|
|
hybrid_debug,
|
|
["vector_only_below_threshold", "mean"],
|
|
0.0,
|
|
),
|
|
0.0,
|
|
),
|
|
query_without_relevant_hit_rate=_safe_float(
|
|
hybrid_debug.get("query_without_relevant_hit_rate"), 0.0
|
|
),
|
|
trace_top_n=_safe_float(data.get("trace_top_n"), 0.0),
|
|
trace_component_top_n=_safe_float(
|
|
data.get("trace_component_top_n"), 0.0
|
|
),
|
|
tuning_state=str(data.get("tuning_state", "")),
|
|
tuning_reason=str(data.get("tuning_reason", "")),
|
|
env_overrides=data.get("env_overrides", {}),
|
|
)
|
|
)
|
|
return out
|
|
|
|
|
|
def _latest_success(rows: list[Row]) -> dict[str, Row]:
|
|
latest: dict[str, Row] = {}
|
|
for row in rows:
|
|
if row.success:
|
|
latest[row.candidate] = row
|
|
return latest
|
|
|
|
|
|
def _matrix_rows(latest_success: dict[str, Row]) -> list[dict[str, Any]]:
|
|
matrix: list[dict[str, Any]] = []
|
|
for candidate, row in latest_success.items():
|
|
matrix.append(
|
|
{
|
|
"candidate": candidate,
|
|
"objective": row.objective,
|
|
"precision_at_k": row.hybrid_precision,
|
|
"mrr": row.hybrid_mrr,
|
|
"ndcg_at_k": row.hybrid_ndcg,
|
|
"recall_at_k": row.hybrid_recall,
|
|
"map": row.hybrid_map,
|
|
"hybrid_eval_ms": row.hybrid_eval_ms,
|
|
"keyword_eval_ms": row.keyword_eval_ms,
|
|
"hybrid_keyword_mrr_delta": row.hybrid_keyword_mrr_delta,
|
|
"duplicate_rate_at_k": row.hybrid_duplicate_rate_at_k,
|
|
"objective_penalty_hybrid_regression": row.objective_penalty_hybrid_regression,
|
|
"objective_penalty_duplicate_rate": row.objective_penalty_duplicate_rate,
|
|
"trace_coverage": row.trace_coverage,
|
|
"graph_rerank_apply_rate": row.graph_rerank_apply_rate,
|
|
"semantic_rescue_rate_mean": row.semantic_rescue_rate_mean,
|
|
"vector_only_below_threshold_mean": row.vector_only_below_threshold_mean,
|
|
"query_without_relevant_hit_rate": row.query_without_relevant_hit_rate,
|
|
"trace_top_n": row.trace_top_n,
|
|
"trace_component_top_n": row.trace_component_top_n,
|
|
"source_file": row.source_file,
|
|
"line_number": row.line_number,
|
|
"tuning_state": row.tuning_state,
|
|
"tuning_reason": row.tuning_reason,
|
|
}
|
|
)
|
|
|
|
matrix.sort(
|
|
key=lambda r: (r["precision_at_k"], r["ndcg_at_k"], r["mrr"], r["objective"]),
|
|
reverse=True,
|
|
)
|
|
return matrix
|
|
|
|
|
|
def _write_csv(path: Path, matrix: list[dict[str, Any]]) -> None:
|
|
fieldnames = [
|
|
"candidate",
|
|
"objective",
|
|
"precision_at_k",
|
|
"mrr",
|
|
"ndcg_at_k",
|
|
"recall_at_k",
|
|
"map",
|
|
"hybrid_eval_ms",
|
|
"keyword_eval_ms",
|
|
"hybrid_keyword_mrr_delta",
|
|
"duplicate_rate_at_k",
|
|
"objective_penalty_hybrid_regression",
|
|
"objective_penalty_duplicate_rate",
|
|
"trace_coverage",
|
|
"graph_rerank_apply_rate",
|
|
"semantic_rescue_rate_mean",
|
|
"vector_only_below_threshold_mean",
|
|
"query_without_relevant_hit_rate",
|
|
"trace_top_n",
|
|
"trace_component_top_n",
|
|
"source_file",
|
|
"line_number",
|
|
"tuning_state",
|
|
"tuning_reason",
|
|
]
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", encoding="utf-8", newline="") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=fieldnames)
|
|
writer.writeheader()
|
|
writer.writerows(matrix)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument(
|
|
"--input", action="append", required=True, help="Input JSONL file"
|
|
)
|
|
parser.add_argument("--json-out", required=True, help="Output matrix JSON path")
|
|
parser.add_argument("--csv-out", required=True, help="Output matrix CSV path")
|
|
args = parser.parse_args()
|
|
|
|
input_paths = [Path(v) for v in args.input]
|
|
rows = _load_rows(input_paths)
|
|
latest = _latest_success(rows)
|
|
matrix = _matrix_rows(latest)
|
|
|
|
perfect = [r for r in matrix if r["precision_at_k"] >= 1.0]
|
|
best = matrix[0] if matrix else None
|
|
|
|
output = {
|
|
"schema_version": "retrieval_opt_matrix_v1",
|
|
"inputs": [str(p) for p in input_paths],
|
|
"record_count": len(rows),
|
|
"successful_latest_candidates": len(matrix),
|
|
"perfect_precision_candidates": perfect,
|
|
"best_precision_candidate": best,
|
|
"matrix": matrix,
|
|
}
|
|
|
|
json_path = Path(args.json_out)
|
|
json_path.parent.mkdir(parents=True, exist_ok=True)
|
|
json_path.write_text(json.dumps(output, indent=2, sort_keys=True), encoding="utf-8")
|
|
|
|
_write_csv(Path(args.csv_out), matrix)
|
|
|
|
if perfect:
|
|
names = ",".join(entry["candidate"] for entry in perfect)
|
|
print(f"perfect_precision_candidates={names}")
|
|
elif best is not None:
|
|
print(
|
|
"best_precision_candidate="
|
|
f"{best['candidate']} precision_at_k={best['precision_at_k']:.6f} "
|
|
f"location={best['source_file']}:{best['line_number']}"
|
|
)
|
|
else:
|
|
print("best_precision_candidate=none")
|
|
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|