scripts/veredicto_contrato.py

"""Semantic supervision, shared context policy and complete local Clef checkpoints.

Importing this module requires only the standard library. PyTorch, PEFT and the
pinned public Clef implementation are loaded only when saving/loading weights.
No function downloads a model. State text is never truncated by character count;
pass ``**caso['contexto']`` to the pinned ``encode_record`` in every consumer.
"""
from __future__ import annotations

from collections.abc import Mapping
import copy
import hashlib
import importlib.util
import json
import math
from pathlib import Path
import shutil
import sys
import tempfile
from typing import Any

BASE_REVISION = '8b2e5fd17c09fd49bc2880b805d5323ec7ff4ff3'
CLEF_SOURCE_SHA256 = '0e304cf7c6500e8bb59bef7e2afd2c6373f82596dfb3b57d1aa93c175e2dc3a3'
CHECKPOINT_VERSION = 1
CATEGORIAS = ('reembolso', 'troca_devolucao', 'fraude_suspeita', 'prazo_envio',
              'pagamento', 'frete', 'produto_defeito', 'chargeback',
              'duvida_comercial', 'outro', 'incidente', 'whatsapp_meta',
              'bug_produto', 'integracao', 'conta_config', 'plano_creditos', 'duvida_produto')
DECISOES = ('aprovar', 'rejeitar', 'escalar', 'executar_acao', 'pedir_mais_info')
URGENCIA = ('0 - baixa', '1 - moderada', '2 - alta', '3 - crítica')
SCHEMA = {'categoria': list(CATEGORIAS), 'decisao': list(DECISOES), 'urgencia': list(URGENCIA)}
CONTEXTO_PADRAO = {'max_length': 16384, 'max_state_tokens': None}


def _opcoes_validas(opcoes: Any, nome: str) -> list[str]:
    if not isinstance(opcoes, (list, tuple)) or not opcoes:
        raise ValueError(f'{nome}: lista de opções vazia ou inválida')
    if any(not isinstance(op, str) or not op for op in opcoes):
        raise ValueError(f'{nome}: opções devem ser strings não vazias')
    if len(set(opcoes)) != len(opcoes):
        raise ValueError(f'{nome}: opções duplicadas')
    return list(opcoes)


def alinhar_alvo(rotulo: int, probs: list[float] | None,
                 opcoes_originais: list[str], option_ids: list[str]
                 ) -> tuple[int, list[float] | None]:
    """Map hard/soft targets by option identity, validating before permutation.

    Soft weights need not initially sum to one. Normalize once after checking
    finite, nonnegative values and positive finite total mass.
    """
    originais = _opcoes_validas(opcoes_originais, 'opcoes_originais')
    codificadas = _opcoes_validas(option_ids, 'option_ids')
    if set(originais) != set(codificadas):
        raise ValueError('conjuntos de opções divergentes')
    if type(rotulo) is not int or not 0 <= rotulo < len(originais):
        raise ValueError('rótulo fora do intervalo ou não inteiro')
    indice = codificadas.index(originais[rotulo])
    if probs is None:
        return indice, None
    if not isinstance(probs, (list, tuple)) or len(probs) != len(originais):
        raise ValueError('distribuição com tamanho incorreto')
    if any(isinstance(p, bool) or not isinstance(p, (int, float)) or
           not math.isfinite(p) or p < 0 for p in probs):
        raise ValueError('distribuição exige pesos finitos e não negativos')
    try:
        massa = math.fsum(probs)
    except OverflowError as exc:
        raise ValueError('massa da distribuição não finita') from exc
    if not math.isfinite(massa) or massa <= 0:
        raise ValueError('massa da distribuição deve ser positiva e finita')
    por_opcao = dict(zip(originais, probs))  # lengths already checked above
    return indice, [por_opcao[op] / massa for op in codificadas]


def politica_contexto(contexto: Mapping[str, Any] | None = None) -> dict[str, Any]:
    """The only context limits used by both training and inference encoders."""
    if contexto is not None and not isinstance(contexto, Mapping):
        raise ValueError('contexto deve ser um mapeamento')
    contexto = dict(contexto or {})
    if set(contexto) - set(CONTEXTO_PADRAO):
        raise ValueError('política de contexto desconhecida; usar limites de tokens')
    resultado = {**CONTEXTO_PADRAO, **contexto}
    max_length, max_state = resultado['max_length'], resultado['max_state_tokens']
    if type(max_length) is not int or max_length <= 0:
        raise ValueError('max_length deve ser inteiro positivo')
    if max_state is not None and (type(max_state) is not int or max_state < 0):
        raise ValueError('max_state_tokens deve ser None ou inteiro não negativo')
    return resultado


def identificar_tarefa(linha: Mapping[str, Any]) -> str:
    """Resolve only complete canonical option sets; legacy conversion is separate."""
    try:
        opcoes = _opcoes_validas(linha['pergunta']['opcoes'], 'opcoes')
    except (KeyError, TypeError) as exc:
        raise ValueError('pergunta sem opções') from exc
    candidatas = [qid for qid, canonicas in SCHEMA.items() if set(opcoes) == set(canonicas)]
    if len(candidatas) != 1:
        raise ValueError('esquema de opções desconhecido ou legado; converter explicitamente')
    tarefa = candidatas[0]
    declarada = linha.get('tarefa', linha.get('question_id'))
    if declarada is not None and declarada != tarefa:
        raise ValueError('tarefa declarada diverge das opções')
    return tarefa


def construir_caso(linhas: list[dict[str, Any]], *,
                   contexto: Mapping[str, Any] | None = None) -> dict[str, Any]:
    """Build one joint record with stable question IDs and unencoded targets.

    Missing tasks remain absent. ``alvos[qid]`` contains ``rotulo``,
    ``probs_alvo`` and ``opcoes_originais``; consumers align via question_id.
    Only group rows from a single case before invoking this function.
    """
    if not linhas:
        raise ValueError('caso sem perguntas')
    policy = politica_contexto(contexto)
    try:
        estado = linhas[0]['estado']
    except (KeyError, TypeError) as exc:
        raise ValueError('caso sem estado') from exc
    perguntas, alvos = {}, {}
    for linha in linhas:
        if linha.get('estado') != estado:
            raise ValueError('caso contém estados divergentes')
        qid = identificar_tarefa(linha)
        if qid in perguntas:
            raise ValueError('caso contém tarefa duplicada')
        ops = list(linha['pergunta']['opcoes'])
        rotulo, probs = linha.get('rotulo'), linha.get('probs_alvo')
        # Validation is shared with the encoded target path; keep original order.
        alinhar_alvo(rotulo, probs, ops, ops)
        perguntas[qid] = {'type': 'choice',
                          'instructions': linha['pergunta'].get('instrucao') or qid,
                          'criteria': {op: op for op in ops}}
        alvos[qid] = {'rotulo': rotulo, 'probs_alvo': copy.deepcopy(probs),
                      'opcoes_originais': ops}
    identidade = json.dumps(estado, ensure_ascii=False, sort_keys=True, separators=(',', ':'))
    record = {'id': hashlib.sha256(identidade.encode('utf-8')).hexdigest(), 'state': estado,
              'questions': {qid: perguntas[qid] for qid in SCHEMA if qid in perguntas}}
    return {'record': record, 'alvos': {qid: alvos[qid] for qid in SCHEMA if qid in alvos},
            'contexto': policy}


def alinhar_perguntas(caso: Mapping[str, Any], perguntas_codificadas: Any
                     ) -> list[tuple[Any, int, list[float] | None]]:
    """Return targets in encoded-question order without positional assumptions."""
    perguntas = list(perguntas_codificadas)
    ids = [q.question_id for q in perguntas]
    alvos = caso['alvos']
    if len(ids) != len(set(ids)) or set(ids) != set(alvos):
        raise ValueError('perguntas codificadas ausentes, extras ou duplicadas')
    resultado = []
    for pergunta in perguntas:
        alvo = alvos[pergunta.question_id]
        indice, probs = alinhar_alvo(alvo['rotulo'], alvo['probs_alvo'],
                                    alvo['opcoes_originais'], pergunta.option_ids)
        resultado.append((pergunta, indice, probs))
    return resultado


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open('rb') as arquivo:
        for bloco in iter(lambda: arquivo.read(1024 * 1024), b''):
            digest.update(bloco)
    return digest.hexdigest()


def _json_write(path: Path, value: Any) -> None:
    path.write_text(json.dumps(value, ensure_ascii=False, sort_keys=True,
                               indent=2, allow_nan=False) + '\n', encoding='utf-8')


def _validar_manifesto(manifesto: Mapping[str, Any]) -> dict[str, Any]:
    if not isinstance(manifesto, Mapping):
        raise ValueError('manifesto deve ser mapeamento')
    meta = copy.deepcopy(dict(manifesto))
    if meta.get('base_revision') != BASE_REVISION:
        raise ValueError('revisão da base deve ser a revisão fixa auditada')
    if not isinstance(meta.get('base_path'), str) or not meta['base_path']:
        raise ValueError('manifesto exige base_path local explícito')
    if not isinstance(meta.get('contexto'), Mapping):
        raise ValueError('manifesto exige política de contexto explícita')
    meta['contexto'] = politica_contexto(meta['contexto'])
    if 'contrato_schema' in meta and meta['contrato_schema'] != SCHEMA:
        raise ValueError('schema do checkpoint diverge do contrato fixo')
    if 'clef_source_sha256' in meta and meta['clef_source_sha256'] != CLEF_SOURCE_SHA256:
        raise ValueError('fonte Clef diverge da revisão auditada')
    meta['contrato_schema'] = copy.deepcopy(SCHEMA)
    meta['clef_source_sha256'] = CLEF_SOURCE_SHA256
    return meta


def _estado_cpu(module: Any) -> dict[str, Any]:
    return {name: value.detach().cpu().contiguous().clone()
            for name, value in module.state_dict().items()}


def salvar_checkpoint(model: Any, processor: Any, destino: str | Path,
                      manifesto: Mapping[str, Any], estado_treino: Any = None) -> Path:
    """Atomically save every trained component; never replace an existing run.

    Head weights are mandatory even when frozen. PEFT adapters are saved with
    PEFT, an unfrozen non-PEFT backbone gets a full state_dict, and a completely
    frozen backbone refers to the fixed local release. Optimizer/RNG states may
    be supplied in estado_treino and are saved separately for explicit resume.
    """
    import torch

    meta = _validar_manifesto(manifesto)
    destino = Path(destino)
    if destino.exists():
        raise FileExistsError(f'checkpoint já existe: {destino}')
    base_path = Path(meta['base_path']).expanduser()
    if not base_path.is_dir():
        raise ValueError('base_path precisa apontar para um diretório local existente')
    if not hasattr(model, 'head') or not hasattr(model, 'language_model'):
        raise ValueError('modelo exige head e language_model')
    config = meta.get('head_config')
    if config is None:
        config_path = base_path / 'joint_head_config.json'
        if not config_path.is_file():
            raise ValueError('configuração da head ausente na base e no manifesto')
        config = json.loads(config_path.read_text(encoding='utf-8'))
    if not isinstance(config, dict):
        raise ValueError('configuração da head deve ser objeto JSON')
    meta['head_config'] = config
    lm = model.language_model
    if getattr(lm, 'peft_config', None):
        mode = 'peft_adapter'
        if len(lm.peft_config) != 1 or set(lm.peft_config) != {'default'}:
            raise ValueError('checkpoint suporta exatamente um adapter PEFT chamado default')
        # Standard PEFT persists LoRA and modules_to_save, not arbitrary unfrozen
        # base parameters. Refuse a state that would silently lose those weights.
        trainable = [name for name, p in lm.named_parameters() if p.requires_grad]
        all_biases_saved = getattr(lm.peft_config['default'], 'bias', 'none') == 'all'
        if any('lora_' not in name and 'modules_to_save.' not in name
               and not (all_biases_saved and name.endswith('.bias')) for name in trainable):
            raise ValueError('PEFT contém parâmetros base treináveis que não serão persistidos')
    elif any(p.requires_grad for p in lm.parameters()):
        mode = 'state_dict'
    else:
        mode = 'frozen_base'
    dtype = str(next(model.head.parameters()).dtype).removeprefix('torch.')
    meta['dtype'] = dtype
    meta['checkpoint'] = {
        'version': CHECKPOINT_VERSION, 'language_model_mode': mode,
        'head_weights': 'joint_head.pt', 'head_config': 'joint_head_config.json',
        'head_trainable': [name for name, p in model.head.named_parameters() if p.requires_grad],
        'lm_trainable': [name for name, p in lm.named_parameters() if p.requires_grad],
        'processor_path': 'processor', 'training_state': estado_treino is not None,
    }
    destino.parent.mkdir(parents=True, exist_ok=True)
    staging = Path(tempfile.mkdtemp(prefix=f'.{destino.name}.tmp-', dir=destino.parent))
    try:
        torch.save(_estado_cpu(model.head), staging / 'joint_head.pt')
        _json_write(staging / 'joint_head_config.json', config)
        processor_path = staging / 'processor'
        processor_path.mkdir()
        processor.save_pretrained(processor_path)
        tokenizer = getattr(processor, 'tokenizer', None)
        if tokenizer is not None:
            tokenizer.save_pretrained(processor_path)
        if mode == 'peft_adapter':
            lm.save_pretrained(staging / 'adapter', safe_serialization=True)
            if not (staging / 'adapter' / 'adapter_config.json').is_file():
                raise ValueError('PEFT não produziu configuração do adapter')
        elif mode == 'state_dict':
            torch.save(_estado_cpu(lm), staging / 'language_model.pt')
        if estado_treino is not None:
            torch.save(estado_treino, staging / 'estado_treino.pt')
        files = [p for p in staging.rglob('*') if p.is_file()]
        if not any(processor_path in p.parents for p in files):
            raise ValueError('processor não foi persistido')
        meta['checkpoint']['files_sha256'] = {
            str(p.relative_to(staging)): _sha256(p) for p in files}
        _json_write(staging / 'manifesto.json', meta)
        # Renaming a complete staging dir keeps interrupted saves invisible.
        if destino.exists():
            raise FileExistsError(f'checkpoint já existe: {destino}')
        staging.rename(destino)
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise
    return destino


def carregar_codigo_clef(base_path: str | Path) -> Any:
    """Return the hash-verified pinned Clef module, without loading model weights.

    Train and evaluation consumers use this same module for encode/collate and
    model construction. Cache entries are keyed by absolute local release path;
    the source hash is checked again even for a cached module.
    """
    base_path = Path(base_path).expanduser().resolve()
    fonte = base_path / 'joint_schema_model.py'
    if not fonte.is_file() or _sha256(fonte) != CLEF_SOURCE_SHA256:
        raise ValueError('código Clef local ausente ou SHA-256 diferente da revisão fixa')
    nome = '_veredicto_clef_' + hashlib.sha256(str(base_path).encode()).hexdigest()[:16]
    if nome in sys.modules:
        return sys.modules[nome]
    spec = importlib.util.spec_from_file_location(nome, fonte)
    if spec is None or spec.loader is None:
        raise ValueError('não foi possível carregar o código Clef local')
    modulo = importlib.util.module_from_spec(spec)
    sys.modules[nome] = modulo
    try:
        spec.loader.exec_module(modulo)
    except BaseException:
        sys.modules.pop(nome, None)
        raise
    return modulo


def _carregar_release(base_path: Path, device: Any, dtype: Any) -> tuple[Any, Any]:
    """Load an existing local release with the audited code and no hub access."""
    modulo = carregar_codigo_clef(base_path)
    return modulo.load_release_model(base_path, device=device, dtype=dtype, local_files_only=True)


def carregar_checkpoint(destino: str | Path, device: Any, *,
                        base_path: str | Path | None = None) -> tuple[Any, Any, dict[str, Any]]:
    """Restore a complete checkpoint strictly. ``base_path`` is a local-only override.

    The returned model is in eval mode. Original requires_grad flags are restored
    for an explicit training resume; estado_treino.pt can be read with torch.load
    using weights_only=True. This function never silently substitutes base head.
    """
    import torch

    destino = Path(destino)
    try:
        raw = json.loads((destino / 'manifesto.json').read_text(encoding='utf-8'))
        meta = _validar_manifesto(raw)
    except (OSError, json.JSONDecodeError) as exc:
        raise ValueError('checkpoint sem manifesto JSON legível') from exc
    ckpt = meta.get('checkpoint', {})
    if ckpt.get('version') != CHECKPOINT_VERSION:
        raise ValueError('versão do checkpoint ausente ou não suportada')
    if raw.get('contrato_schema') != SCHEMA or raw.get('clef_source_sha256') != CLEF_SOURCE_SHA256:
        raise ValueError('checkpoint não declara schema e fonte Clef fixos')
    mode = ckpt.get('language_model_mode')
    if mode not in {'peft_adapter', 'state_dict', 'frozen_base'}:
        raise ValueError('modo do backbone desconhecido')
    files = ckpt.get('files_sha256')
    if not isinstance(files, dict) or not {'joint_head.pt', 'joint_head_config.json'} <= set(files):
        raise ValueError('checkpoint incompleto: inventário da head ausente')
    if not any(name.startswith('processor/') for name in files):
        raise ValueError('checkpoint incompleto: processor ausente')
    if mode == 'state_dict' and 'language_model.pt' not in files:
        raise ValueError('checkpoint incompleto: backbone treinado ausente')
    if mode == 'peft_adapter' and ('adapter/adapter_config.json' not in files or
            not any(name.startswith('adapter/') and name.endswith(('.safetensors', '.bin')) for name in files)):
        raise ValueError('checkpoint incompleto: adapter ausente')
    if ckpt.get('training_state') and 'estado_treino.pt' not in files:
        raise ValueError('checkpoint incompleto: estado de treino ausente')
    for name, expected in files.items():
        path = destino / name
        if Path(name).is_absolute() or '..' in Path(name).parts or not path.is_file():
            raise ValueError('arquivo ausente ou caminho inválido no checkpoint')
        if _sha256(path) != expected:
            raise ValueError(f'checksum divergente no checkpoint: {name}')
    if json.loads((destino / 'joint_head_config.json').read_text(encoding='utf-8')) != meta.get('head_config'):
        raise ValueError('configuração da head diverge do manifesto')
    local_base = Path(base_path if base_path is not None else meta['base_path']).expanduser()
    if not local_base.is_dir():
        raise ValueError('base local indisponível; fornecer base_path explícito')
    dtype = getattr(torch, meta.get('dtype', ''), None)
    if not isinstance(dtype, torch.dtype) or not dtype.is_floating_point:
        raise ValueError('dtype de checkpoint ausente ou inválido')
    model, processor_base = _carregar_release(local_base, device, dtype)
    base_config_path = local_base / 'joint_head_config.json'
    if base_config_path.is_file():
        base_head_config = json.loads(base_config_path.read_text(encoding='utf-8'))
        if base_head_config != meta['head_config']:
            # A candidate may explicitly change head architecture; instantiate
            # its saved configuration instead of silently using the base head.
            clef = carregar_codigo_clef(local_base)
            model.head = clef.JointSchemaHead(**meta['head_config']).to(device=device, dtype=dtype)
    if mode == 'peft_adapter':
        from peft import PeftModel
        model.language_model = PeftModel.from_pretrained(
            model.language_model, destino / 'adapter', is_trainable=True,
            local_files_only=True)
    elif mode == 'state_dict':
        model.language_model.load_state_dict(
            torch.load(destino / 'language_model.pt', map_location=device, weights_only=True), strict=True)
    model.head.load_state_dict(
        torch.load(destino / 'joint_head.pt', map_location=device, weights_only=True), strict=True)
    model.head.to(device=device, dtype=dtype)
    for module, field in [(model.head, 'head_trainable'), (model.language_model, 'lm_trainable')]:
        expected = ckpt.get(field)
        parameters = dict(module.named_parameters())
        if not isinstance(expected, list) or not set(expected) <= set(parameters):
            raise ValueError('inventário de parâmetros treináveis incompatível')
        for name, parameter in parameters.items():
            parameter.requires_grad_(name in expected)
    processor = type(processor_base).from_pretrained(destino / 'processor', local_files_only=True)
    return model.eval(), processor, meta