scripts/avaliar_veredicto.py

#!/usr/bin/env python3
"""Avalia checkpoint COMPLETO do Veredicto; T=1, sem calibrar no ouro.

python scripts/avaliar_veredicto.py --checkpoint RUN/checkpoint --base-path SNAPSHOT
Predicoes/logits nao incluem textos de chamados. O split oficial e mantido.
"""
import argparse
from collections import Counter, defaultdict
import hashlib
import json
import math
from pathlib import Path
import random
import sys

ROOT = Path(__file__).resolve().parents[1]


def chave_caso(estado):
    return hashlib.md5(estado.encode("utf-8")).hexdigest()


def split_por_caso(exemplos, seed=42):
    from veredicto_contrato import identificar_tarefa
    estratos = defaultdict(list)
    for row in exemplos:
        if identificar_tarefa(row) == "categoria":
            estratos[row["rotulo"]].append(chave_caso(row["estado"]))
    todos = {chave_caso(row["estado"]) for row in exemplos}
    com_cat = set().union(*map(set, estratos.values())) if estratos else set()
    estratos["sem_categoria"] = sorted(todos - com_cat)
    rng, calib, teste = random.Random(seed), set(), set()
    for _, keys in sorted(estratos.items(), key=lambda kv: str(kv[0])):
        keys = sorted(set(keys))
        rng.shuffle(keys)
        meio = len(keys) // 2
        calib.update(keys[:meio]); teste.update(keys[meio:])
    return calib, teste


def _validar(rows):
    vistos = set()
    if not rows:
        raise ValueError("Nenhuma predicao para avaliar")
    for row in rows:
        key = (row["caso"], row["q"])
        if key in vistos:
            raise ValueError("Predicao caso/pergunta duplicada")
        vistos.add(key)
        probs, ops = row["probs"], row["option_ids"]
        if len(ops) != len(probs) or not probs or len(set(ops)) != len(ops):
            raise ValueError("Opcoes/probabilidades invalidas")
        if any(not math.isfinite(x) or x < 0 for x in probs) or abs(sum(probs)-1) > 1e-5:
            raise ValueError("Probabilidades invalidas")
        if any(type(row[k]) is not int or not 0 <= row[k] < len(ops) for k in ("rotulo", "pred")):
            raise ValueError("Indice de predicao/rotulo invalido")
        if row["pred"] != max(range(len(probs)), key=probs.__getitem__):
            raise ValueError("Predicao diverge do argmax")


def metricas(rows):
    _validar(rows)
    tarefas, subsets = defaultdict(list), defaultdict(list)
    for row in rows:
        tarefas[row["q"]].append(row)
        for key in (f"{row.get('subset', 'nao_informado')}|{row['q']}",
                    f"{row.get('subset', 'nao_informado')}|todas"):
            subsets[key].append(row)
    def acc(group):
        return sum(r["pred"] == r["rotulo"] for r in group) / len(group)
    por_q = {}
    for qid, group in tarefas.items():
        bins = defaultdict(list)
        for r in group:
            confidence = max(r["probs"])
            bins[min(14, max(0, math.ceil(confidence * 15) - 1))].append(r)
        ece = sum(len(g) / len(group) * abs(sum(max(r["probs"]) for r in g) / len(g) - acc(g))
                  for g in bins.values())
        nll = -sum(math.log(max(r["probs"][r["rotulo"]], 1e-12)) for r in group) / len(group)
        brier = sum(sum((p - float(i == r["rotulo"]))**2 for i, p in enumerate(r["probs"]))
                    for r in group) / len(group)
        por_q[qid] = {"n": len(group), "acc": acc(group), "ece": ece, "nll": nll, "brier": brier}
    result = {"teste_n": len(rows), "casos": len({r["caso"] for r in rows}),
              "acc_geral": acc(rows), "por_pergunta": por_q,
              "acc_por_subset": {k: {"n": len(v), "acc": acc(v)} for k,v in subsets.items()},
              "temperatura": 1.0, "calibracao_ajustada": False}
    classes, conf = defaultdict(lambda: [0, 0, 0]), Counter()
    for r in tarefas.get("categoria", []):
        real, pred = r["option_ids"][r["rotulo"]], r["option_ids"][r["pred"]]
        if real == pred:
            classes[real][0] += 1
        else:
            classes[real][2] += 1; classes[pred][1] += 1
            conf[(real, pred)] += 1
    if classes:
        f1 = {k: 2*v[0]/max(1, 2*v[0]+v[1]+v[2]) for k,v in classes.items()}
        result["categoria_macro_f1"] = sum(f1.values())/len(f1)
        result["categoria_f1_por_classe"] = f1
    result["confusoes"] = [{"real": a, "predita": b, "n": n} for (a,b),n in conf.most_common()]
    return result


def bootstrap_pareado(a, b, repeticoes=10000, seed=42):
    _validar(a); _validar(b)
    if repeticoes < 1:
        raise ValueError("repeticoes deve ser positiva")
    pa, pb = {(r["caso"],r["q"]): r for r in a}, {(r["caso"],r["q"]): r for r in b}
    if set(pa) != set(pb):
        raise ValueError("Pares caso/pergunta incompativeis")
    grupos = defaultdict(list)
    for key, ra in pa.items():
        rb = pb[key]
        if ra["option_ids"][ra["rotulo"]] != rb["option_ids"][rb["rotulo"]]:
            raise ValueError("Rotulos de referencia divergentes")
        grupos[key[0]].append(int(ra["pred"] == ra["rotulo"]) - int(rb["pred"] == rb["rotulo"]))
    keys, rng, boot = sorted(grupos), random.Random(seed), []
    for _ in range(repeticoes):
        sample = [grupos[rng.choice(keys)] for _ in keys]
        boot.append(sum(map(sum,sample))/sum(map(len,sample)))
    boot.sort()
    return {"delta": sum(map(sum,grupos.values()))/sum(map(len,grupos.values())),
            "ic95": [boot[int(.025*(repeticoes-1))],boot[int(.975*(repeticoes-1))]],
            "casos": len(keys), "repeticoes": repeticoes, "seed": seed}


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--checkpoint", required=True)
    ap.add_argument("--base-path", required=True)
    ap.add_argument("--gold", default=str(ROOT / "data/processed/ouro_gold.jsonl"))
    ap.add_argument("--out", required=True)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--seed", type=int, default=42)
    ap.add_argument("--comparar-preds")
    args = ap.parse_args()
    import torch
    from veredicto_contrato import construir_caso, alinhar_perguntas, carregar_checkpoint, carregar_codigo_clef
    out = Path(args.out); out.mkdir(parents=True, exist_ok=False)
    examples = [json.loads(l) for l in Path(args.gold).open()]
    _, test = split_por_caso(examples, args.seed)
    groups = defaultdict(list)
    for idx, row in enumerate(examples):
        if chave_caso(row["estado"]) in test:
            groups[chave_caso(row["estado"])].append((idx,row))
    model, processor, meta = carregar_checkpoint(args.checkpoint,args.device,base_path=args.base_path)
    # Loader verified this exact public source and revision before importing it.
    clef = carregar_codigo_clef(args.base_path)
    model.eval()
    predictions = []
    with torch.inference_mode():
        for case_key, group in groups.items():
            case = construir_caso([row for _,row in group], contexto=meta["contexto"])
            enc = clef.encode_record(processor.tokenizer,case["record"],processor=processor,**case["contexto"])
            batch = clef.collate_records([enc],processor.tokenizer.pad_token_id,torch.device(args.device))
            logits = model(batch)[0]
            aligned = alinhar_perguntas(case,enc.questions)
            if len(logits) != len(aligned):
                raise ValueError("Logits/perguntas divergentes")
            by_q = {}
            from veredicto_contrato import identificar_tarefa
            for idx,row in group:
                by_q[identificar_tarefa(row)] = (idx,row)
            for (qenc,target,_),lg in zip(aligned,logits):
                if len(qenc.option_ids) != lg.numel():
                    raise ValueError("Logits/opcoes divergentes")
                idx,row = by_q[qenc.question_id]
                probs = lg.float().softmax(-1).cpu().tolist()
                predictions.append({"idx":idx,"caso":case_key,"q":qenc.question_id,
                    "subset":row.get("subset","nao_informado"),"option_ids":list(qenc.option_ids),
                    "rotulo":target,"pred":max(range(len(probs)),key=probs.__getitem__),
                    "probs":probs,"logits":lg.float().cpu().tolist()})
    result = metricas(predictions)
    result.update({"gold_sha256":hashlib.sha256(Path(args.gold).read_bytes()).hexdigest(),
                   "checkpoint":str(Path(args.checkpoint).resolve()),"manifesto_checkpoint":meta,"seed":args.seed})
    if args.comparar_preds:
        other = [json.loads(l) for l in Path(args.comparar_preds).open()]
        result["comparacao_pareada"] = bootstrap_pareado(predictions,other,seed=args.seed)
    (out/"preds.jsonl").write_text("".join(json.dumps(r,ensure_ascii=False)+"\n" for r in predictions))
    (out/"relatorio.json").write_text(json.dumps(result,ensure_ascii=False,indent=2)+"\n")
    print(json.dumps({"acc_geral":result["acc_geral"],"por_pergunta":result["por_pergunta"]},ensure_ascii=False))


if __name__ == "__main__":
    main()