audit/audit_datasets.py

"""Dry-run audit des 27 datasets staging.yaml.

Stream max 200 samples ou 60s par dataset. Logger :
- load_dataset success/fail
- yielded count
- cast_errors
- formats matched (1-6 de to_chatml_messages)
- lang distribution
- drops par raison
- kept count + %

Rapport JSON : audit_report.json
"""
from __future__ import annotations
import json
import os
import signal
import sys
import time
import traceback
from collections import defaultdict, Counter
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

import yaml
from datasets import load_dataset

# Import les helpers de prepare_sft.py SANS exec main
from prepare_sft import (
    to_chatml_messages,
    render_chatml,
    get_sample_lang,
    is_truncated,
    get_assistant_text,
    has_nonempty_assistant,
    passes_filter,
    approx_token_count,
    StagingConfig,
    DatasetSpec,
    load_config,
    get_langid_model,
)

MAX_SAMPLES_PER_DS = 200
MAX_TIME_PER_DS = 60.0  # seconds


def detect_format(sample: dict) -> str:
    """Renvoie quel Format de to_chatml_messages matchera ce sample."""
    if "messages" in sample and isinstance(sample["messages"], list):
        return "F1_messages"
    if "conversations" in sample and isinstance(sample["conversations"], list):
        return "F2_conversations"
    if "instruction" in sample:
        return "F3_instruction"
    if "prompt" in sample and ("response" in sample or "completion" in sample):
        return "F4_prompt_response"
    if "problem" in sample and ("solution" in sample or "answer" in sample):
        return "F4b_problem_solution"
    if ("query" in sample or "question" in sample) and "answer" in sample:
        return "F4c_query_answer"
    if "query" in sample and "response" in sample:
        return "F4d_query_response"
    if "chosen" in sample and "rejected" in sample:
        return "F5_dpo"
    if "text" in sample:
        return "F6_text"
    return "UNKNOWN"


def audit_dataset(spec: DatasetSpec, cfg: StagingConfig, langid_model) -> dict:
    out: dict = {
        "id": spec.id,
        "split": spec.split,
        "group": spec.group,
        "load_ok": False,
        "load_err": None,
        "yielded": 0,
        "cast_errors": 0,
        "formats": Counter(),
        "first_keys": [],
        "normalize_fail": 0,
        "empty_assistant": 0,
        "lang_excluded": 0,
        "truncated": 0,
        "too_short": 0,
        "too_long": 0,
        "kept": 0,
        "langs_kept": Counter(),
        "elapsed_s": 0.0,
    }
    t0 = time.time()
    try:
        ds = load_dataset(spec.id, split=spec.split, streaming=True, trust_remote_code=False)
        out["load_ok"] = True
    except Exception as e:
        out["load_err"] = f"{type(e).__name__}: {str(e)[:200]}"
        out["elapsed_s"] = round(time.time() - t0, 1)
        return out

    it = iter(ds)
    while True:
        if out["yielded"] >= MAX_SAMPLES_PER_DS:
            break
        if (time.time() - t0) >= MAX_TIME_PER_DS:
            out["timeout"] = True
            break
        try:
            sample = next(it)
        except StopIteration:
            break
        except Exception as e:
            out["cast_errors"] += 1
            if out["cast_errors"] == 1:
                out["first_cast_err"] = f"{type(e).__name__}: {str(e)[:200]}"
            if out["cast_errors"] > 50:
                out["aborted_cast"] = True
                break
            continue

        out["yielded"] += 1
        if not out["first_keys"]:
            out["first_keys"] = sorted(list(sample.keys()))[:15]

        fmt = detect_format(sample)
        out["formats"][fmt] += 1

        # Try full normalize
        msgs = to_chatml_messages(sample, spec.id, cfg.system_default)
        if msgs is None:
            out["normalize_fail"] += 1
            continue
        if cfg.drop_empty_assistant and not has_nonempty_assistant(msgs):
            out["empty_assistant"] += 1
            continue
        text = render_chatml(msgs)
        lang = get_sample_lang(sample, msgs, text, langid_model)

        assistant_text = get_assistant_text(msgs)
        if is_truncated(assistant_text):
            out["truncated"] += 1
            continue

        ok, reason = passes_filter(text, lang, cfg)
        if not ok:
            if "too_short" in reason:
                out["too_short"] += 1
            elif "too_long" in reason:
                out["too_long"] += 1
            elif "lang_excluded" in reason:
                out["lang_excluded"] += 1
            continue

        out["kept"] += 1
        out["langs_kept"][lang] += 1

    out["elapsed_s"] = round(time.time() - t0, 1)
    out["formats"] = dict(out["formats"])
    out["langs_kept"] = dict(out["langs_kept"])
    out["kept_pct"] = round(100.0 * out["kept"] / max(1, out["yielded"]), 1)
    # Verdict
    if not out["load_ok"]:
        out["verdict"] = "DEAD_load_fail"
    elif out["yielded"] < 50:
        out["verdict"] = "DEAD_no_yield"
    elif out["kept_pct"] < 10:
        out["verdict"] = "DEAD_no_kept"
    elif out["yielded"] < 100:
        out["verdict"] = "WEAK"
    else:
        out["verdict"] = "OK"
    return out


def main() -> int:
    cfg = load_config("staging.yaml")
    cfg.cache.mkdir(parents=True, exist_ok=True)
    langid_model = get_langid_model(cfg.cache)
    print(f"langid model loaded: {langid_model is not None}")

    results = []
    for i, spec in enumerate(cfg.datasets, 1):
        print(f"\n[{i}/{len(cfg.datasets)}] {spec.id} (group={spec.group})")
        try:
            r = audit_dataset(spec, cfg, langid_model)
        except Exception as e:
            r = {"id": spec.id, "group": spec.group, "audit_err": f"{type(e).__name__}: {e}"}
        results.append(r)
        # Print summary
        v = r.get("verdict", "ERROR")
        print(f"  -> {v} | yielded={r.get('yielded',0)} kept={r.get('kept',0)} ({r.get('kept_pct',0)}%) cast_err={r.get('cast_errors',0)}")
        if r.get("load_err"):
            print(f"     LOAD_ERR: {r['load_err']}")

    out_path = Path("audit_report.json")
    out_path.write_text(json.dumps(results, indent=2, default=str), encoding="utf-8")
    print(f"\n[REPORT] {out_path}")

    # Summary by group
    print("\n=== SUMMARY BY GROUP ===")
    by_group: dict = defaultdict(list)
    for r in results:
        by_group[r.get("group", "?")].append(r)
    for g, items in by_group.items():
        ok = sum(1 for r in items if r.get("verdict") == "OK")
        weak = sum(1 for r in items if r.get("verdict") == "WEAK")
        dead = sum(1 for r in items if r.get("verdict", "").startswith("DEAD"))
        total = len(items)
        print(f"  {g:20s} OK={ok}/{total} WEAK={weak} DEAD={dead}")
        for r in items:
            v = r.get("verdict", "?")
            print(f"     [{v:14s}] {r['id']:60s} kept={r.get('kept_pct',0):.1f}%")

    return 0


if __name__ == "__main__":
    sys.exit(main())