File size: 3,916 Bytes
81f30c1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """Fases 0 (tarea 2) y 4 de docs/PLAN_MEJORA_RAG.md — la batería CLOSE sobre el goal raíz de cada
problema, 0 llamadas LLM: misma función (`close_goal`) y misma selección de premisas
(`select_premises`) que usa el grafo. Reanudable; una fila por (problema, variante).
Variantes:
battery_plain sin premisas (lo que corren `no_rag` y `baseline_auto`)
battery_rag premisas por goal: prior + índice denso (fusión RRF, canal C encendido) + hints + `grind +suggestions`
battery_rag_prior igual, sólo con el canal A (la configuración por defecto, `close.dense_channel: false`)
Una prueba cuenta si, además de cerrar el goal en el REPL, el archivo final pasa `LeanChecker.check`
(validador + kernel + auditoría de axiomas).
Uso: RAGSOLVER_MATHLIB_TABLE=mathlib_full python scripts/probe_close.py [--out data/eval/results/probe_close/battery.jsonl]
"""
from __future__ import annotations
import argparse
import json
import time
from pathlib import Path
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--problems", default="data/eval/problems/eval_v1.jsonl")
ap.add_argument("--out", default="data/eval/results/probe_close/battery.jsonl")
ap.add_argument("--variants", default="battery_plain,battery_rag,battery_rag_prior")
args = ap.parse_args()
from ragsolver.agents.agents import load_agents_config
from ragsolver.lean.verify import LeanChecker
from ragsolver.orchestrator.graph import close_goal, select_premises
from ragsolver.rag import premises as P
from ragsolver.rag.retrieve import Retriever
cfg = load_agents_config()
ccfg, lean = cfg["close"], cfg["lean"]
problems = [json.loads(l) for l in open(args.problems, encoding="utf-8") if l.strip()]
out = Path(args.out)
out.parent.mkdir(parents=True, exist_ok=True)
done = {(r["problem_id"], r["variant"]) for r in map(json.loads, out.open())} if out.exists() else set()
variants = args.variants.split(",")
retriever = Retriever() if "battery_rag" in variants else None
if retriever is not None:
retriever.trace_dir = str(out.parent / "traces")
checker = LeanChecker(imports=lean["imports"], timeout_s=600)
with out.open("a", encoding="utf-8") as f:
for p in problems:
todo = [v for v in variants if (p["problem_id"], v) not in done]
if not todo:
continue
goals, errors = checker.goals(p["formal_statement"], "sorry", p.get("header", ""), lean["max_heartbeats"])
for v in todo:
t0 = time.time()
block, log, names = None, [], []
if len(goals) == 1 and not errors:
g = goals[0]
if v == "battery_plain":
block, log = close_goal(checker, g, ccfg["tactics"], False, [], [], ccfg["timeout_s"])
else:
names = P.prior_channel(g["goal"], k=ccfg["k_premises"]) if v == "battery_rag_prior" else select_premises(retriever, g["goal"], ccfg["k_premises"], ctx={"problem_id": p["problem_id"], "agent": "probe"}, dense_channel=True)[0]
block, log = close_goal(checker, g, ccfg["tactics"], True, names, P.nlinarith_hints(g["goal"]), ccfg["timeout_s"])
ok = bool(block) and checker.check(p["formal_statement"], block, p.get("header", ""), lean["max_heartbeats"]).ok
f.write(json.dumps({"problem_id": p["problem_id"], "variant": v, "ok": ok, "block": block, "premises": names, "tried": log, "s": round(time.time() - t0, 1)}, ensure_ascii=False) + "\n")
f.flush()
print(p["problem_id"], v, ok, flush=True)
checker.close()
rows = [json.loads(l) for l in out.open()]
print(json.dumps({v: sum(r["ok"] for r in rows if r["variant"] == v) for v in variants}), flush=True)
if __name__ == "__main__":
main()
|