"""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())