code/self_healing_finetuning/data_pipeline/paired_eval_diagnostics.py

from __future__ import annotations

import hashlib
from collections import Counter
from pathlib import Path
from typing import Any

from data_pipeline.learning_pdf_to_jsonl import write_jsonl
from data_pipeline.pairwise_preference_memory import (
    DIAGNOSTIC_BUCKETS,
    candidate_critical_failed,
    candidate_lost,
    clean_model_response,
    corrected_chosen_answer,
    defects_for_prediction,
    failure_bucket_for_prediction,
    judge_rationale_for_prediction,
    read_json,
    read_jsonl,
    repair_acceptance_checks,
    repair_admitted_to_training,
    score_number,
    utc_now,
    winner_for_prediction,
)
from n21.config import write_json
from n21.settings import SHFT_WORKSPACE_ROOT


DIAGNOSTIC_SCHEMA_VERSION = "shft_paired_eval_diagnostics_v1"


def sha256_file(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def diagnostic_row(*, prediction: dict[str, Any], run_id: str, asset_class: str, role: str, record_index: int) -> dict[str, Any]:
    defects = defects_for_prediction(prediction)
    prompt = str(prediction.get("prompt") or "").strip()
    candidate = clean_model_response(str(prediction.get("candidate_answer") or prediction.get("candidate_response") or ""))
    baseline = clean_model_response(str(prediction.get("baseline_answer") or prediction.get("baseline_response") or ""))
    repair_answer_raw, repair_source = corrected_chosen_answer(
        prediction=prediction, asset_class=asset_class, role=role, defects=defects
    )
    repair_answer = clean_model_response(repair_answer_raw)
    acceptance = repair_acceptance_checks(prompt=prompt, chosen=repair_answer, rejected=candidate, task=prediction.get("task"))
    is_loss = candidate_lost(prediction)
    is_critical = candidate_critical_failed(prediction)
    return {
        "schema_version": DIAGNOSTIC_SCHEMA_VERSION,
        "source_run_id": run_id,
        "record_index": record_index,
        "prompt_id": prediction.get("id"),
        "task": prediction.get("task"),
        "prompt": prompt,
        "baseline_answer": baseline,
        "candidate_answer": candidate,
        "winner": winner_for_prediction(prediction),
        "failure_bucket": failure_bucket_for_prediction(prediction, defects),
        "defect_types": defects,
        "critical_failure": is_critical,
        "pairwise_loss": is_loss,
        "baseline_score": score_number(prediction.get("baseline_score")),
        "candidate_score": score_number(prediction.get("candidate_score")),
        "delta": prediction.get("delta"),
        "judge_rationale": judge_rationale_for_prediction(prediction, defects),
        "repair_target": {
            "answer": repair_answer,
            "source": repair_source,
            "acceptance_checks": acceptance,
            "admitted_to_training": bool((is_loss or is_critical) and repair_admitted_to_training(acceptance)),
        },
        "created_at": utc_now(),
    }


def validate_diagnostic_row(row: dict[str, Any]) -> list[str]:
    errors: list[str] = []
    for field in ("prompt", "baseline_answer", "candidate_answer", "winner", "failure_bucket"):
        if not isinstance(row.get(field), str) or not row[field].strip():
            errors.append(f"{field} must be non-empty text")
    if row.get("failure_bucket") not in DIAGNOSTIC_BUCKETS:
        errors.append(f"failure_bucket must be one of {', '.join(DIAGNOSTIC_BUCKETS)}")
    if not isinstance(row.get("judge_rationale"), list) or not row["judge_rationale"]:
        errors.append("judge_rationale must be a non-empty list")
    repair = row.get("repair_target")
    if not isinstance(repair, dict):
        errors.append("repair_target must be an object")
    else:
        answer = repair.get("answer")
        checks = repair.get("acceptance_checks")
        if not isinstance(answer, str) or not answer.strip():
            errors.append("repair_target.answer must be non-empty")
        if not isinstance(checks, dict) or not checks:
            errors.append("repair_target.acceptance_checks must be non-empty")
        if "<think" in str(answer).lower() or "</think" in str(answer).lower():
            errors.append("repair_target.answer must not contain think tags")
    return errors


def write_markdown_summary(path: Path, result: dict[str, Any]) -> None:
    summary = result["summary"]
    lines = [
        "# SHFT Paired Eval Diagnostics",
        "",
        f"- Run: `{result['run_id']}`",
        f"- Asset/role: `{result['asset_class']}/{result['role']}`",
        f"- Predictions: `{summary['prediction_count']}`",
        f"- Losses: `{summary['pairwise_loss_count']}`",
        f"- Critical failures: `{summary['critical_failure_count']}`",
        f"- Accepted repair targets: `{summary['accepted_repair_target_count']}`",
        f"- Diagnostics JSONL: `{result['diagnostics_jsonl_path']}`",
        f"- Repair targets JSONL: `{result['repair_targets_jsonl_path']}`",
        "",
        "## Failure Buckets",
        "",
    ]
    for bucket, count in sorted(summary["failure_bucket_counts"].items(), key=lambda item: (-item[1], item[0])):
        lines.append(f"- `{bucket}`: `{count}`")
    lines.append("")
    path.write_text("\n".join(lines), encoding="utf-8")


def build_paired_eval_diagnostics(
    *,
    run_id: str,
    asset_class: str,
    role: str,
    predictions_path: Path | None = None,
    output_dir: Path | None = None,
    max_records: int | None = None,
) -> dict[str, Any]:
    run_path = SHFT_WORKSPACE_ROOT / "runs" / run_id
    predictions = predictions_path or run_path / "eval" / "paired_predictions.jsonl"
    if not predictions.exists():
        raise FileNotFoundError(f"paired predictions not found: {predictions}")
    report_path = predictions.parent / "paired_eval_report.json"
    report = read_json(report_path) if report_path.exists() else {}
    out_dir = output_dir or run_path / "diagnostics"
    diagnostics_path = out_dir / "paired_eval_diagnostics.jsonl"
    repair_targets_path = out_dir / "repair_targets.jsonl"
    manifest_path = out_dir / "paired_eval_diagnostics_manifest.json"
    markdown_path = out_dir / "paired_eval_diagnostics_summary.md"

    rows: list[dict[str, Any]] = []
    repair_targets: list[dict[str, Any]] = []
    schema_errors: list[str] = []
    bucket_counts: Counter[str] = Counter()
    loss_bucket_counts: Counter[str] = Counter()
    defect_counts: Counter[str] = Counter()
    source_loss_ids: set[str] = set()
    diagnosed_loss_ids: set[str] = set()

    for idx, prediction in enumerate(read_jsonl(predictions), start=1):
        if max_records is not None and len(rows) >= max_records:
            break
        if candidate_lost(prediction):
            source_loss_ids.add(str(prediction.get("id") or idx))
        row = diagnostic_row(prediction=prediction, run_id=run_id, asset_class=asset_class, role=role, record_index=idx)
        errors = validate_diagnostic_row(row)
        if errors:
            schema_errors.extend(f"{row.get('prompt_id')}: {error}" for error in errors)
            continue
        rows.append(row)
        bucket_counts[row["failure_bucket"]] += 1
        if row["pairwise_loss"]:
            loss_bucket_counts[row["failure_bucket"]] += 1
            diagnosed_loss_ids.add(str(row.get("prompt_id") or idx))
        defect_counts.update(row["defect_types"])
        if row["repair_target"]["admitted_to_training"]:
            repair_targets.append(
                {
                    "schema_version": "shft_repair_target_v1",
                    "source_run_id": run_id,
                    "source_prediction_id": row["prompt_id"],
                    "failure_bucket": row["failure_bucket"],
                    "defect_types": row["defect_types"],
                    "source_task": row["task"],
                    "source_delta": row["delta"],
                    "prompt": row["prompt"],
                    "repair_answer": row["repair_target"]["answer"],
                    "acceptance_checks": row["repair_target"]["acceptance_checks"],
                    "created_at": row["created_at"],
                }
            )

    out_dir.mkdir(parents=True, exist_ok=True)
    write_jsonl(diagnostics_path, rows)
    write_jsonl(repair_targets_path, repair_targets)
    summary = {
        "prediction_count": len(rows),
        "source_pairwise_loss_count": len(source_loss_ids),
        "pairwise_loss_count": sum(1 for row in rows if row["pairwise_loss"]),
        "diagnosed_pairwise_loss_count": len(diagnosed_loss_ids),
        "latest_loss_coverage_ratio": round(len(diagnosed_loss_ids) / max(1, len(source_loss_ids)), 6),
        "latest_loss_coverage_ok": len(source_loss_ids) == len(diagnosed_loss_ids),
        "critical_failure_count": sum(1 for row in rows if row["critical_failure"]),
        "accepted_repair_target_count": len(repair_targets),
        "failure_bucket_counts": {bucket: int(bucket_counts.get(bucket, 0)) for bucket in DIAGNOSTIC_BUCKETS},
        "loss_failure_bucket_counts": {bucket: int(loss_bucket_counts.get(bucket, 0)) for bucket in DIAGNOSTIC_BUCKETS},
        "defect_counts": dict(defect_counts),
        "schema_error_count": len(schema_errors),
    }
    result = {
        "ok": bool(rows) and not schema_errors and summary["latest_loss_coverage_ok"],
        "schema_version": DIAGNOSTIC_SCHEMA_VERSION,
        "run_id": run_id,
        "asset_class": asset_class,
        "role": role,
        "predictions_path": str(predictions),
        "paired_eval_report_path": str(report_path) if report_path.exists() else None,
        "paired_eval_improvement": report.get("improvement", {}),
        "diagnostics_jsonl_path": str(diagnostics_path),
        "diagnostics_jsonl_sha256": sha256_file(diagnostics_path),
        "repair_targets_jsonl_path": str(repair_targets_path),
        "repair_targets_jsonl_sha256": sha256_file(repair_targets_path),
        "manifest_path": str(manifest_path),
        "markdown_path": str(markdown_path),
        "summary": summary,
        "schema_errors": schema_errors,
        "created_at": utc_now(),
    }
    write_json(manifest_path, result)
    write_markdown_summary(markdown_path, result)
    return result