raulmodena/leire-corpus

1"""Item 6: validacao do harness.2 3Duas coisas diferentes:4(a) o MECANISMO esta certo? -> comparar minha loglikelihood em lote com padding5    contra uma implementacao ingenua, um exemplo por vez, sem padding. E aqui que6    moram os bugs de verdade (deslocamento, fatia errada, alinhamento de alvo).7(b) qual normalizacao o numero publico usa? -> calcular as tres e reportar todas,8    sem escolher a que da mais perto.9"""10import json, math, re, sys, time11from pathlib import Path12import numpy as np13import torch14import torch.nn.functional as F15from transformers import AutoModelForCausalLM, AutoTokenizer16 17RAIZ = Path('/kaggle/working')18OUT = RAIZ / 'relatorios' / 'harness_validacao.json'19EST = {}20NOME = 'HuggingFaceTB/SmolLM2-135M'21 22tok = AutoTokenizer.from_pretrained(NOME)23model = AutoModelForCausalLM.from_pretrained(NOME, torch_dtype=torch.float32).cuda().eval()24print('modelo carregado', flush=True)25 26 27def limpar(t):28    t = t.strip().replace(' [title]', '. ')29    t = re.sub(r'\[.*?\]', '', t)30    return t.replace('  ', ' ')31 32 33from datasets import load_dataset34ds = load_dataset('Rowan/hellaswag', split='validation')35docs = []36for row in ds:37    ctx = row['ctx_a'] + ' ' + row['ctx_b'].capitalize()38    q = limpar(row['activity_label'] + ': ' + ctx)39    docs.append((q, [' ' + limpar(e) for e in row['endings']], int(row['label'])))40print('docs', len(docs), flush=True)41 42 43def codificar(ctx, cont):44    a = tok.encode(ctx, add_special_tokens=False)45    b = tok.encode(ctx + cont, add_special_tokens=False)[len(a):]46    if not b:47        b = tok.encode(cont, add_special_tokens=False)48    return a, b49 50 51@torch.no_grad()52def ingenua(ctx, cont):53    """Um exemplo por vez, sem padding e sem lote. Referencia."""54    a, b = codificar(ctx, cont)55    seq = torch.tensor([a + b], device='cuda')56    lp = F.log_softmax(model(seq).logits.float(), dim=-1)[0]57    alvo = seq[0, len(a):]58    escolhido = lp[len(a) - 1:len(a) - 1 + len(b)].gather(-1, alvo.unsqueeze(-1)).squeeze(-1)59    return float(escolhido.sum().item())60 61 62@torch.no_grad()63def em_lote(pares, batch_size=32):64    """Minha versao: ordena por tamanho, empacota com padding a esquerda."""65    cod = []66    for i, (ctx, cont) in enumerate(pares):67        a, b = codificar(ctx, cont)68        seq = (a + b)[-2049:]69        cod.append((i, seq, min(len(b), len(seq) - 1)))70    cod.sort(key=lambda x: -len(x[1]))71    saida = [None] * len(pares)72    for ini in range(0, len(cod), batch_size):73        lote = cod[ini:ini + batch_size]74        maior = max(len(s) for _, s, _ in lote)75        # PADDING A DIREITA. Com padding a esquerda e atencao causal os tokens76        # reais enxergam o padding e a loglikelihood sai errada (medido: erro de77        # 22,7 sobre magnitude tipica 48,5). A direita, a causalidade ignora sozinha.78        ent = torch.zeros(len(lote), maior, dtype=torch.long)79        for j, (_, seq, _) in enumerate(lote):80            ent[j, :len(seq)] = torch.tensor(seq, dtype=torch.long)81        ent = ent.cuda()82        lp = F.log_softmax(model(ent).logits.float(), dim=-1)83        for j, (idx, seq, n_cont) in enumerate(lote):84            com = len(seq) - 1 - n_cont85            alvo = ent[j, com + 1:com + 1 + n_cont]86            saida[idx] = float(lp[j, com:com + n_cont].gather(-1, alvo.unsqueeze(-1)).sum().item())87        if (ini // batch_size) % 60 == 0:88            print('  lote', ini, '/', len(cod), flush=True)89    return saida90 91 92# ---------- (a) mecanismo: lote+padding contra ingenua ----------93print('=== (a) mecanismo ===', flush=True)94amostra = []95for d in docs[:60]:96    for c in d[1]:97        amostra.append((d[0], c))98ref = [ingenua(c, k) for c, k in amostra]99meu = em_lote(amostra, batch_size=32)100difs = [abs(r - m) for r, m in zip(ref, meu)]101EST['mecanismo'] = {'n_pares': len(amostra), 'erro_abs_max': max(difs),102                    'erro_abs_medio': float(np.mean(difs)),103                    'magnitude_tipica': float(np.mean([abs(r) for r in ref]))}104print('MECANISMO erro_max', format(max(difs), '.3e'),105      'sobre magnitude tipica', format(EST['mecanismo']['magnitude_tipica'], '.1f'), flush=True)106OUT.write_text(json.dumps(EST, indent=1), encoding='utf-8')107 108# ---------- (b) as tres normalizacoes, no conjunto inteiro ----------109print('=== (b) hellaswag completo ===', flush=True)110pedidos = [(d[0], c) for d in docs for c in d[1]]111t0 = time.time()112lls = em_lote(pedidos, batch_size=32)113print('loglikelihoods em', round(time.time() - t0), 's', flush=True)114 115acc, accn_char, accn_tok = [], [], []116k = 0117for ctx, conts, gold in docs:118    n = len(conts)119    v = np.array(lls[k:k + n], dtype=np.float64)120    chars = np.array([float(len(c)) for c in conts])121    toks = np.array([float(len(codificar(ctx, c)[1])) for c in conts])122    acc.append(float(int(v.argmax()) == gold))123    accn_char.append(float(int((v / chars).argmax()) == gold))124    accn_tok.append(float(int((v / toks).argmax()) == gold))125    k += n126 127PUB = 42.1128res = {}129for rot, arr in [('acc', acc), ('acc_norm_chars', accn_char), ('acc_norm_tokens', accn_tok)]:130    a = np.asarray(arr)131    res[rot] = {'valor_pct': round(float(a.mean()) * 100, 3),132                'stderr_pct': round(float(a.std(ddof=1) / math.sqrt(len(a))) * 100, 3),133                'erro_vs_publicado_pct': round(abs(float(a.mean()) * 100 - PUB) / PUB * 100, 3)}134    print(rot, res[rot], flush=True)135EST['hellaswag'] = {'n': len(docs), 'publicado_card_lighteval': PUB, 'variantes': res}136EST['melhor_variante'] = min(res, key=lambda r: res[r]['erro_vs_publicado_pct'])137EST['dentro_de_1pct'] = res[EST['melhor_variante']]['erro_vs_publicado_pct'] <= 1.0138EST['fim'] = True139OUT.write_text(json.dumps(EST, indent=1), encoding='utf-8')140print('MELHOR', EST['melhor_variante'], 'DENTRO_DE_1PCT', EST['dentro_de_1pct'], flush=True)141