Spaces:
Sleeping
Sleeping
| """Implementation for the `chat_eval cli tune` command.""" | |
| from __future__ import annotations | |
| import json | |
| from datetime import datetime, timedelta, timezone | |
| def build_audit_entry( | |
| *, | |
| step_id: int, | |
| ts: str, | |
| outcome: str, | |
| identity: dict, | |
| window_start_ts: str, | |
| window_end_ts: str, | |
| cleanup_dropped_doc_ids: list[str], | |
| **extra, | |
| ) -> dict: | |
| return { | |
| "step_id": step_id, | |
| "ts": ts, | |
| "outcome": outcome, | |
| **identity, | |
| "window_start_ts": window_start_ts, | |
| "window_end_ts": window_end_ts, | |
| "cleanup_dropped_doc_ids": cleanup_dropped_doc_ids, | |
| **extra, | |
| } | |
| def build_identity_fields(turns_path_or_list) -> dict: | |
| from scripts.lib.chat_eval import judge as _judge | |
| from scripts.lib.chat_eval.rubric import RUBRIC_VERSION | |
| from scripts.lib.chat_eval.turn_log import iter_turns_compat | |
| fields = { | |
| "rubric_version": int(RUBRIC_VERSION), | |
| "judge_model": _judge._JUDGE_MODEL, | |
| "chat_model": None, | |
| "embedding_model": None, | |
| "catalog_artifact_version_id": None, | |
| } | |
| latest_judged_ts = "" | |
| for t in iter_turns_compat(turns_path_or_list): | |
| if not t.get("judge"): | |
| continue | |
| ts = t.get("ts", "") | |
| if ts > latest_judged_ts: | |
| latest_judged_ts = ts | |
| fields["chat_model"] = t.get("chat_model") | |
| fields["embedding_model"] = t.get("embedding_model") | |
| fields["catalog_artifact_version_id"] = t.get("catalog_artifact_version_id") | |
| return fields | |
| def recent_util_history(turns_path_or_list, doc_id: str, *, end_ts: str, | |
| days: int = 14) -> list[dict]: | |
| from scripts.lib.chat_eval.turn_log import iter_turns_compat | |
| end_dt = datetime.strptime(end_ts, "%Y-%m-%dT%H:%M:%SZ").replace( | |
| tzinfo=timezone.utc | |
| ) | |
| out: list[dict] = [] | |
| for offset in range(days - 1, -1, -1): | |
| bin_start = end_dt - timedelta(days=offset + 1) | |
| bin_end = end_dt - timedelta(days=offset) | |
| bin_start_iso = bin_start.strftime("%Y-%m-%dT%H:%M:%SZ") | |
| bin_end_iso = bin_end.strftime("%Y-%m-%dT%H:%M:%SZ") | |
| retrieved_count = 0 | |
| used_count = 0 | |
| for turn in iter_turns_compat(turns_path_or_list, since=bin_start_iso): | |
| if turn.get("ts", "") >= bin_end_iso or not turn.get("judge"): | |
| continue | |
| doc_ids_retrieved = {r["doc_id"] for r in (turn.get("retrieved") or [])} | |
| if doc_id in doc_ids_retrieved: | |
| retrieved_count += 1 | |
| used = (turn.get("judge") or {}).get("docs_actually_used") or [] | |
| if any(u.get("doc_id") == doc_id for u in used): | |
| used_count += 1 | |
| if retrieved_count == 0: | |
| continue | |
| out.append({ | |
| "bin_start_date": bin_start.date().isoformat(), | |
| "retrieved": retrieved_count, | |
| "used": used_count, | |
| "util_pct": round(100.0 * used_count / retrieved_count, 1), | |
| }) | |
| return out | |
| def _append_jsonl(path, entry: dict) -> None: | |
| with open(path, "a", encoding="utf-8") as f: | |
| f.write(json.dumps(entry) + "\n") | |
| def _load_previous_accepted(log_path) -> dict | None: | |
| previous = None | |
| if not log_path.is_file(): | |
| return previous | |
| for line in log_path.read_text("utf-8").splitlines(): | |
| if not line.strip(): | |
| continue | |
| try: | |
| entry = json.loads(line) | |
| except json.JSONDecodeError: | |
| continue | |
| if entry.get("outcome") == "accepted": | |
| previous = entry | |
| return previous | |
| def _cleanup_stale_penalties(*, state: dict, turns_path, now_ts: str, dry_run: bool) -> list[str]: | |
| cleanup_dropped: list[str] = [] | |
| if not state.get("penalties"): | |
| return cleanup_dropped | |
| catalog_path = turns_path.parent / "catalog.jsonl" | |
| if not catalog_path.is_file(): | |
| return cleanup_dropped | |
| known_ids: set[str] = set() | |
| for line in catalog_path.read_text(encoding="utf-8").splitlines(): | |
| if not line.strip(): | |
| continue | |
| try: | |
| rec = json.loads(line) | |
| except json.JSONDecodeError: | |
| continue | |
| doc_id = rec.get("doc_id") | |
| if doc_id: | |
| known_ids.add(doc_id) | |
| for doc_id in list(state["penalties"].keys()): | |
| if doc_id not in known_ids: | |
| state["penalties"].pop(doc_id) | |
| cleanup_dropped.append(doc_id) | |
| if cleanup_dropped and not dry_run: | |
| state["updated_at"] = now_ts | |
| return cleanup_dropped | |
| def _write_skip( | |
| *, | |
| args, | |
| state: dict, | |
| now_ts: str, | |
| outcome: str, | |
| identity: dict, | |
| window_start_ts: str, | |
| window_end_ts: str, | |
| cleanup_dropped: list[str], | |
| **extra, | |
| ) -> None: | |
| entry = build_audit_entry( | |
| step_id=state["step_count"] + 1, | |
| ts=now_ts, | |
| outcome=outcome, | |
| identity=identity, | |
| window_start_ts=window_start_ts, | |
| window_end_ts=window_end_ts, | |
| cleanup_dropped_doc_ids=cleanup_dropped, | |
| **extra, | |
| ) | |
| if not args.dry_run: | |
| _append_jsonl(args.tuning_log_path, entry) | |
| def run(args) -> int: | |
| from scripts.lib.chat_eval import analytics, tuning, turn_log | |
| state_path = args.tuning_state_path | |
| log_path = args.tuning_log_path | |
| turns_path = args.log_path | |
| all_turns = [ | |
| t for t in turn_log.load_all(turns_path) | |
| if turn_log.is_student_mode_turn(t) | |
| ] | |
| state = ( | |
| json.loads(state_path.read_text("utf-8")) | |
| if state_path.is_file() | |
| else tuning.empty_state() | |
| ) | |
| if args.unfreeze: | |
| if state["frozen"]: | |
| state["frozen"] = False | |
| state["frozen_reason"] = None | |
| state["updated_at"] = tuning.now_utc_iso() | |
| state_path.write_text(json.dumps(state), encoding="utf-8") | |
| print(f"Unfroze tuning state at {state_path}.") | |
| else: | |
| print("State already not frozen; no-op.") | |
| return 0 | |
| now_ts = tuning.now_utc_iso() | |
| now_dt = datetime.strptime(now_ts, "%Y-%m-%dT%H:%M:%SZ").replace(tzinfo=timezone.utc) | |
| window_end_ts = now_ts | |
| window_start_ts = (now_dt - timedelta(days=7)).strftime("%Y-%m-%dT%H:%M:%SZ") | |
| cleanup_dropped = _cleanup_stale_penalties( | |
| state=state, turns_path=turns_path, now_ts=now_ts, dry_run=args.dry_run | |
| ) | |
| if cleanup_dropped and not args.dry_run: | |
| state_path.write_text(json.dumps(state), encoding="utf-8") | |
| identity = build_identity_fields(all_turns) | |
| if state["frozen"]: | |
| _write_skip( | |
| args=args, | |
| state=state, | |
| now_ts=now_ts, | |
| outcome="skipped_frozen", | |
| identity=identity, | |
| window_start_ts=window_start_ts, | |
| window_end_ts=window_end_ts, | |
| cleanup_dropped=cleanup_dropped, | |
| frozen_reason=state.get("frozen_reason"), | |
| ) | |
| print(f"Skipped: state is frozen ({state.get('frozen_reason')!r}).") | |
| return 0 | |
| prev_accepted = _load_previous_accepted(log_path) | |
| before_scores = None | |
| after_scores = None | |
| before_window_start_ts = None | |
| before_window_end_ts = None | |
| regression_check = None | |
| if prev_accepted: | |
| before_window_start_ts = prev_accepted["window_start_ts"] | |
| before_window_end_ts = prev_accepted["window_end_ts"] | |
| before_stats = analytics.score_summary_with_stats( | |
| all_turns, since=before_window_start_ts, until=before_window_end_ts | |
| ) | |
| after_stats = analytics.score_summary_with_stats( | |
| all_turns, since=prev_accepted["ts"], until=window_end_ts | |
| ) | |
| before_scores = {k: v["mean"] for k, v in before_stats.items()} | |
| after_scores = {k: v["mean"] for k, v in after_stats.items()} | |
| after_count = max((v["n"] for v in after_stats.values()), default=0) | |
| regression_check = tuning.check_regression( | |
| before_scores, | |
| after_scores, | |
| judged_turns=after_count, | |
| before_stats=before_stats, | |
| after_stats=after_stats, | |
| ) | |
| if regression_check.get("skipped"): | |
| _write_skip( | |
| args=args, | |
| state=state, | |
| now_ts=now_ts, | |
| outcome="skipped_insufficient_turns", | |
| identity=identity, | |
| window_start_ts=window_start_ts, | |
| window_end_ts=window_end_ts, | |
| cleanup_dropped=cleanup_dropped, | |
| judged_turns_in_window=after_count, | |
| previous_step_id=prev_accepted["step_id"], | |
| ) | |
| print( | |
| f"Skipped: only {after_count} judged turns since step " | |
| f"{prev_accepted['step_id']}; need >=" | |
| f"{tuning.MIN_TURNS_FOR_REGRESSION_CHECK}." | |
| ) | |
| return 0 | |
| if regression_check["any_dimension_regressed"]: | |
| reason = ( | |
| f"{regression_check['max_drop_dimension']} regressed " | |
| f"delta={regression_check['max_drop_value']} " | |
| f"after step {prev_accepted['step_id']}" | |
| ) | |
| new_state = tuning.rollback_step(state, prev_accepted, reason=reason) | |
| rb_step_action = dict(prev_accepted.get("step_action") or {}) | |
| if rb_step_action.get("doc_id"): | |
| rb_step_action["recent_util_history"] = recent_util_history( | |
| all_turns, rb_step_action["doc_id"], end_ts=window_end_ts | |
| ) | |
| rb_entry = build_audit_entry( | |
| step_id=new_state["step_count"] + 1, | |
| ts=now_ts, | |
| outcome="rolled_back", | |
| identity=identity, | |
| window_start_ts=window_start_ts, | |
| window_end_ts=window_end_ts, | |
| cleanup_dropped_doc_ids=cleanup_dropped, | |
| before_window_start_ts=before_window_start_ts, | |
| before_window_end_ts=before_window_end_ts, | |
| rolled_back_step_id=prev_accepted["step_id"], | |
| regression_check=regression_check, | |
| scores_before=before_scores, | |
| scores_after=after_scores, | |
| step_action=rb_step_action, | |
| ) | |
| if not args.dry_run: | |
| state_path.write_text(json.dumps(new_state), encoding="utf-8") | |
| _append_jsonl(log_path, rb_entry) | |
| print(f"Regression detected: {reason}. Rolled back + frozen.") | |
| return 0 | |
| util = analytics.doc_utilization(all_turns) | |
| step_action = tuning.propose_step(util, current_state=state) | |
| if step_action is None: | |
| _write_skip( | |
| args=args, | |
| state=state, | |
| now_ts=now_ts, | |
| outcome="skipped_no_eligible_doc", | |
| identity=identity, | |
| window_start_ts=window_start_ts, | |
| window_end_ts=window_end_ts, | |
| cleanup_dropped=cleanup_dropped, | |
| judged_turns_seen=sum(1 for t in all_turns if t.get("judge")), | |
| ) | |
| print("No eligible doc for a step; skipped.") | |
| return 0 | |
| new_state = tuning.apply_step(state, step_action) | |
| step_action["recent_util_history"] = recent_util_history( | |
| all_turns, step_action["doc_id"], end_ts=window_end_ts | |
| ) | |
| if prev_accepted is None: | |
| cold_after_stats = analytics.score_summary_with_stats( | |
| all_turns, since=window_start_ts, until=window_end_ts | |
| ) | |
| cold_after_scores = {k: v["mean"] for k, v in cold_after_stats.items()} | |
| after_scores = cold_after_scores | |
| regression_check = {"skipped": True, "reason": "cold_start"} | |
| entry = build_audit_entry( | |
| step_id=new_state["step_count"], | |
| ts=now_ts, | |
| outcome="accepted", | |
| identity=identity, | |
| window_start_ts=window_start_ts, | |
| window_end_ts=window_end_ts, | |
| cleanup_dropped_doc_ids=cleanup_dropped, | |
| before_window_start_ts=before_window_start_ts, | |
| before_window_end_ts=before_window_end_ts, | |
| scores_before=before_scores, | |
| scores_after=after_scores, | |
| regression_check=regression_check, | |
| step_action=step_action, | |
| ) | |
| if args.dry_run: | |
| print(json.dumps({"would_apply": step_action, "audit_entry": entry}, indent=2)) | |
| return 0 | |
| state_path.write_text(json.dumps(new_state), encoding="utf-8") | |
| _append_jsonl(log_path, entry) | |
| print( | |
| f"Applied step #{new_state['step_count']} on doc " | |
| f"{step_action['doc_id']}: {step_action['penalty_before']} -> " | |
| f"{step_action['penalty_after']}" | |
| ) | |
| return 0 | |