File size: 3,180 Bytes
a31a731 | 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 76 77 78 79 80 81 82 83 84 85 86 87 88 | """Langfuse logging for checkpoint eval batches (Kaggle + local)."""
from __future__ import annotations
import os
from typing import Any
def init_langfuse():
from langfuse import Langfuse, get_client
Langfuse(
public_key=os.environ["LANGFUSE_PUBLIC_KEY"],
secret_key=os.environ["LANGFUSE_SECRET_KEY"],
host=os.getenv("LANGFUSE_HOST")
or os.getenv("LANGFUSE_BASE_URL")
or "https://cloud.langfuse.com",
)
lf = get_client()
if not lf.auth_check():
raise RuntimeError("Langfuse auth_check failed")
return lf
def log_checkpoint_eval(
*,
step: int,
hf_revision: str,
train_model_id: str,
results: list[dict[str, Any]],
wandb_run_id: str | None = None,
trace_name: str = "kaggle-ckpt-eval",
) -> dict[str, Any]:
"""
Write one Langfuse agent trace for a checkpoint eval suite.
Each problem becomes a nested span; aggregate scores on the trace.
MUST be called BEFORE pushing/saving the checkpoint artifact.
"""
lf = init_langfuse()
n = len(results)
mean_pass = sum(float(r.get("pass_frac", 0.0)) for r in results) / n if n else 0.0
mean_reward = sum(float(r.get("reward_total", 0.0)) for r in results) / n if n else 0.0
with lf.start_as_current_observation(
name=trace_name,
as_type="evaluator",
metadata={
"step": step,
"hf_revision": hf_revision,
"train_model_id": train_model_id,
"wandb_run_id": wandb_run_id,
"phase": "ckpt_eval",
"n_problems": n,
},
input={"step": step, "hf_revision": hf_revision, "n_problems": n},
) as root:
for r in results:
with lf.start_as_current_observation(
name=f"task:{r.get('task_id')}",
as_type="span",
input={"prompt": r.get("prompt", "")[:2000]},
metadata={"task_id": r.get("task_id")},
) as span:
span.update(
output={
"completion": (r.get("completion") or "")[:2000],
"pass_frac": r.get("pass_frac"),
"reward_total": r.get("reward_total"),
"ok": r.get("pass_frac", 0) >= 1.0,
}
)
summary = {
"mean_pass_frac": mean_pass,
"mean_reward_total": mean_reward,
"n_problems": n,
"n_perfect": sum(1 for r in results if r.get("pass_frac", 0) >= 1.0),
"step": step,
"hf_revision": hf_revision,
}
root.update(output=summary)
lf.score_current_trace(name="mean_pass_frac", value=mean_pass)
lf.score_current_trace(name="mean_reward_total", value=mean_reward)
lf.score_current_trace(name="n_perfect", value=float(summary["n_perfect"]))
lf.score_current_trace(name="checkpoint_step", value=float(step))
trace_id = lf.get_current_trace_id() or ""
lf.flush()
return {"langfuse_trace_id": trace_id, **summary}
|