scripts/servir_veredicto_multi.py

#!/usr/bin/env python3
"""servir_veredicto_multi.py — API MULTI-ESPECIALISTA do Veredicto ("Frankenstein").

Roteamento por tarefa sobre UM backbone clef-flash (bf16, congelado):

  categoria ← cabeça do B2 (híbrido α=0,5)     — melhor cat da régua: 0,7012
  decisão   ← cabeça do E1 (híbrido α=margem)  — melhor dec: 0,6892
  urgência  ← adapter LoRA v4 + cabeça release — melhor urg: 0,7450

B2 e E1 são frozen_base (só a cabeça muda); o único LoRA é o v4. Logo não é
multi-LoRA de verdade: é 1 backbone + 3 cabeças + 1 adapter, 3 forwards/caso.

Régua esperada (teste de mesa 22/out, preds compostas nos 753 casos seed 42):
cat 0,7012 / dec 0,6892 / urg 0,7450 / geral 0,7118 — empate estatístico com
o juiz de fronteira (0,7127). Este servidor deve reproduzir esses números.

Mesmo contrato da produção (servir_ensemble.py): GET /saude,
POST /decidir {"estado", "limiar"} → {"respostas", "confianca_roteamento", ...}.

Uso (container treino, GPU 0, porta 8200):
  python3 scripts/servir_veredicto_multi.py \
    --checkpoint-cat models/veredicto_b2_step-2208 \
    --checkpoint-dec models/veredicto_e1_step-2576 \
    --base-path /workspace/.cache/huggingface/hub/models--Cloudflare--clef-flash/snapshots/8b2e5fd17c09fd49bc2880b805d5323ec7ff4ff3 \
    --adapter models/veredicto-lora-v4 \
    --porta 8200 --device cuda
"""
import argparse
import json
import sys
import threading
import time
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from chamados_para_decisor import CATEGORIAS, DECISOES, URGENCIA  # noqa: E402

OPCOES = {"categoria": list(CATEGORIAS), "decisao": list(DECISOES),
          "urgencia": list(URGENCIA)}
INSTRUCOES = {
    "categoria": "Qual a categoria principal do chamado?",
    "decisao": "Qual decisão tomar para este caso?",
    "urgencia": "Qual a urgência? (0=baixa, 3=crítica)",
}

LOCK = threading.Lock()  # troca de cabeça + toggle de adapter não são thread-safe
MODEL = {}


def _linha(tarefa: str, estado: str) -> dict:
    ops = OPCOES[tarefa]
    return {"estado": estado, "tarefa": tarefa,
            "pergunta": {"opcoes": ops, "instrucao": INSTRUCOES[tarefa]},
            "rotulo": 0}  # rótulo dummy (índice): exigido pelo contrato, ignorado


def _rotulos_codificados(qenc, tokenizer, canonicas: list[str]) -> list[str]:
    """Posição codificada → texto da opção, NA ORDEM dos option_ids.

    A ordem canônica de negócio (DECISOES/CATEGORIAS) NÃO coincide com a
    ordem codificada (alfabética) — devolver list(canonicas) troca os
    rótulos (bug pego na paridade 22/out: probs certas, nomes trocados;
    urgência escapava porque lá as duas ordens coincidem). Fallback:
    ordem alfabética, a convenção da régua.
    """
    try:
        decod = [tokenizer.decode([oid]).strip() for oid in qenc.option_ids]
        if sorted(decod) == sorted(o.strip() for o in canonicas):
            return decod
    except Exception:
        pass
    return sorted(o.strip() for o in canonicas)


def _forward(head: str, usar_adapter: bool, batch):
    """Um forward com a cabeça e o estado de adapter pedidos. Chamar sob LOCK.

    O retorno é CLONADO tensor a tensor: o clef reutiliza o buffer de saída
    entre forwards, então sem clone os logits do 1º forward seriam
    sobrescritos pelo 3º (bug pego na paridade 22/out: urgência 25/25 —
    era sempre a saída do ÚLTIMO forward — e cat/dec erradas).
    """
    import torch
    m = MODEL
    model = m["model"]
    model.head = m["heads"][head]
    lm = model.language_model
    if not usar_adapter:
        with lm.disable_adapter(), torch.inference_mode():
            return [t.detach().clone() for t in model(batch)[0]]
    with torch.inference_mode():
        return [t.detach().clone() for t in model(batch)[0]]


def decidir(estado: str, limiar: float) -> dict:
    import torch
    from veredicto_contrato import construir_caso, alinhar_perguntas
    m = MODEL
    clef, processor = m["clef"], m["processor"]
    tarefas = ["categoria", "decisao", "urgencia"]
    caso = construir_caso([_linha(t, estado) for t in tarefas],
                          contexto=m["meta"]["contexto"])
    enc = clef.encode_record(processor.tokenizer, caso["record"],
                             processor=processor, **caso["contexto"])
    batch = clef.collate_records([enc], processor.tokenizer.pad_token_id,
                                 torch.device(m["device"]))
    aligned = alinhar_perguntas(caso, enc.questions)

    t0 = time.time()
    with LOCK:
        logits_cat = _forward("b2", False, batch)        # cat ← cabeça B2
        logits_dec = _forward("e1", False, batch)        # dec ← cabeça E1
        logits_urg = _forward("release", True, batch)    # urg ← v4 + release
    latencia_ms = round((time.time() - t0) * 1000)

    fonte = {"categoria": logits_cat, "decisao": logits_dec, "urgencia": logits_urg}
    respostas, confs = {}, []
    for k, (qenc, _, _) in enumerate(aligned):
        tarefa = qenc.question_id
        lg = fonte[tarefa][k]
        probs = lg.float().softmax(-1).cpu().tolist()
        rotulos = _rotulos_codificados(qenc, processor.tokenizer,
                                       caso["alvos"][tarefa]["opcoes_originais"])
        i = max(range(len(probs)), key=probs.__getitem__)
        respostas[tarefa] = {
            "resposta": rotulos[i],
            "confianca": round(probs[i], 4),
            "distribuicao": {r: round(p, 4) for r, p in zip(rotulos, probs)},
        }
        confs.append(probs[i])
    conf = min(confs)
    return {"respostas": respostas,
            "confianca_roteamento": round(conf, 4),
            "limiar": limiar,
            "destino": "automatico" if conf >= limiar else "system2",
            "latencia_ms": latencia_ms,
            "shadow": {"modo": m["modo"], "especialistas": m["especialistas"]}}


class Handler(BaseHTTPRequestHandler):
    def _json(self, obj, status=200):
        corpo = json.dumps(obj, ensure_ascii=False).encode()
        self.send_response(status)
        self.send_header("Content-Type", "application/json; charset=utf-8")
        self.send_header("Content-Length", str(len(corpo)))
        self.end_headers()
        self.wfile.write(corpo)

    def log_message(self, *a):
        pass

    def do_GET(self):
        if self.path == "/saude":
            self._json({"status": "ok", "shadow": True,
                        "modo": MODEL["modo"],
                        "especialistas": MODEL["especialistas"],
                        "calibrado": False})
            return
        self._json({"erro": "use GET /saude ou POST /decidir"}, 404)

    def do_POST(self):
        if self.path != "/decidir":
            self._json({"erro": "rota desconhecida"}, 404)
            return
        try:
            corpo = json.loads(self.rfile.read(int(self.headers["Content-Length"])))
            estado = str(corpo["estado"]).strip()
            if not estado:
                raise ValueError("estado vazio")
            limiar = float(corpo.get("limiar", 0.8))
            if not 0.0 <= limiar <= 1.0:
                raise ValueError("limiar fora de [0,1]")
        except (KeyError, ValueError, TypeError) as e:
            self._json({"erro": f"payload inválido: {e}",
                        "formato": {"estado": "texto", "limiar": 0.8}}, 400)
            return
        try:
            self._json(decidir(estado, limiar))
        except Exception as e:  # erro vira 500 isolado, nunca derruba o servidor
            self._json({"erro": f"falha de inferência no multi: {e}"}, 500)


def _carregar_head(clef, ckpt_dir: str, device, dtype):
    """Instancia uma JointSchemaHead de um checkpoint frozen_base e carrega os pesos."""
    import torch
    d = Path(ckpt_dir)
    cfg = json.loads((d / "joint_head_config.json").read_text(encoding="utf-8"))
    head = clef.JointSchemaHead(**cfg)
    head.load_state_dict(
        torch.load(d / "joint_head.pt", map_location=device, weights_only=True),
        strict=True)
    return head.to(device=device, dtype=dtype).eval()


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--checkpoint-cat", required=True,
                    help="checkpoint frozen_base cuja cabeça responde categoria (B2)")
    ap.add_argument("--checkpoint-dec", required=True,
                    help="checkpoint frozen_base cuja cabeça responde decisão (E1)")
    ap.add_argument("--base-path", required=True)
    ap.add_argument("--adapter", required=True,
                    help="adapter LoRA de urgência (v4) — usado com a cabeça de release")
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--porta", type=int, default=8200)
    args = ap.parse_args()

    import torch
    from safetensors.torch import load_file
    from veredicto_contrato import carregar_checkpoint, carregar_codigo_clef

    # O checkpoint do B2 fornece backbone + processor + meta + cabeça de categoria.
    model, processor, meta = carregar_checkpoint(
        args.checkpoint_cat, args.device, base_path=args.base_path)
    clef = carregar_codigo_clef(args.base_path)
    model.eval()
    dtype = next(model.head.parameters()).dtype

    heads = {"b2": model.head,
             "e1": _carregar_head(clef, args.checkpoint_dec, args.device, dtype)}

    # Adapter de urgência: backbone NÃO é PeftModel (frozen_base) — wrap limpo.
    from peft import PeftModel
    model.language_model = PeftModel.from_pretrained(
        model.language_model, args.adapter, is_trainable=False,
        local_files_only=True)
    base = Path(args.base_path)
    cfg = json.loads((base / "joint_head_config.json").read_text(encoding="utf-8"))
    head_release = clef.JointSchemaHead(**cfg)
    head_release.load_state_dict(load_file(str(base / "joint_head.safetensors")),
                                 strict=True)
    heads["release"] = head_release.to(device=args.device, dtype=dtype).eval()
    model.eval()

    MODEL.update({
        "model": model, "processor": processor, "clef": clef, "meta": meta,
        "heads": heads, "device": args.device, "modo": "multi_especialista",
        "especialistas": {
            "categoria": args.checkpoint_cat.rstrip("/").split("/")[-1],
            "decisao": args.checkpoint_dec.rstrip("/").split("/")[-1],
            "urgencia": args.adapter.rstrip("/").split("/")[-1] + " + head_release",
        },
    })
    print(f"✓ Veredicto MULTI na porta {args.porta} — "
          f"cat←{MODEL['especialistas']['categoria']} "
          f"dec←{MODEL['especialistas']['decisao']} "
          f"urg←{MODEL['especialistas']['urgencia']}", flush=True)
    ThreadingHTTPServer(("0.0.0.0", args.porta), Handler).serve_forever()


if __name__ == "__main__":
    main()