raulmodena/leire-corpus

1"""Item 6 da Fase 0: harness de avaliacao validado contra numero publico.2 3Protocolo do lm-evaluation-harness: para cada (contexto, continuacao) soma o4log-prob dos tokens da continuacao; acc usa a soma crua, acc_norm divide pelo5numero de caracteres da continuacao.6 7Lei 1: o harness so vale depois de reproduzir um numero publico conhecido.8"""9import json, math, re, sys, time10from pathlib import Path11 12import numpy as np13import torch14import torch.nn.functional as F15 16RAIZ = Path('/kaggle/working')17SAIDA = RAIZ / 'relatorios' / 'harness.json'18SAIDA.parent.mkdir(parents=True, exist_ok=True)19 20 21class ModeloHF:22    def __init__(self, nome, device='cuda', max_len=2048):23        from transformers import AutoModelForCausalLM, AutoTokenizer24        self.nome = nome25        self.tok = AutoTokenizer.from_pretrained(nome)26        self.model = AutoModelForCausalLM.from_pretrained(nome, torch_dtype=torch.float32).to(device)27        self.model.eval()28        self.device = device29        self.max_len = max_len30 31    def encode(self, texto):32        return self.tok.encode(texto, add_special_tokens=False)33 34    @torch.no_grad()35    def logits(self, lote):36        return self.model(lote).logits37 38 39@torch.no_grad()40def loglikelihoods(modelo, pedidos, batch_size=16):41    codificados = []42    for i, (ctx, cont) in enumerate(pedidos):43        a = modelo.encode(ctx)44        b = modelo.encode(ctx + cont)[len(a):]45        if not b:46            b = modelo.encode(cont)47        seq = (a + b)[-(modelo.max_len + 1):]48        codificados.append((i, seq, min(len(b), len(seq) - 1)))49    codificados.sort(key=lambda x: -len(x[1]))50    saida = [None] * len(pedidos)51    for ini in range(0, len(codificados), batch_size):52        lote = codificados[ini:ini + batch_size]53        maior = max(len(s) for _, s, _ in lote)54        entrada = torch.zeros(len(lote), maior, dtype=torch.long)55        for b, (_, seq, _) in enumerate(lote):56            entrada[b, maior - len(seq):] = torch.tensor(seq, dtype=torch.long)57        entrada = entrada.to(modelo.device)58        lp = F.log_softmax(modelo.logits(entrada).float(), dim=-1)59        for b, (idx, seq, n_cont) in enumerate(lote):60            fim = maior - 161            comeco = fim - n_cont62            alvo = entrada[b, comeco + 1:fim + 1]63            l = lp[b, comeco:fim]64            saida[idx] = float(l.gather(-1, alvo.unsqueeze(-1)).sum().item())65        if (ini // batch_size) % 40 == 0:66            print(f'  {ini + len(lote)}/{len(codificados)}', flush=True)67    return saida68 69 70def avaliar(modelo, docs, batch_size=16):71    pedidos = []72    for ctx, conts, gold in docs:73        for c in conts:74            pedidos.append((ctx, c))75    res = loglikelihoods(modelo, pedidos, batch_size)76    acc, accn, k = [], [], 077    for ctx, conts, gold in docs:78        n = len(conts)79        lls = np.array(res[k:k + n], dtype=np.float64)80        comp = np.array([float(len(c)) for c in conts])81        acc.append(float(int(lls.argmax()) == gold))82        accn.append(float(int((lls / comp).argmax()) == gold))83        k += n84    a, an = np.asarray(acc), np.asarray(accn)85    n = len(a)86    return {'n': n, 'acc': float(a.mean()),87            'acc_stderr': float(a.std(ddof=1) / math.sqrt(n)),88            'acc_norm': float(an.mean()),89            'acc_norm_stderr': float(an.std(ddof=1) / math.sqrt(n))}90 91 92def limpar_hellaswag(t):93    t = t.strip().replace(' [title]', '. ')94    t = re.sub(r'\[.*?\]', '', t)95    return t.replace('  ', ' ')96 97 98def hellaswag(limite=None):99    from datasets import load_dataset100    ds = load_dataset('Rowan/hellaswag', split='validation')101    docs = []102    for row in ds:103        ctx = row['ctx_a'] + ' ' + row['ctx_b'].capitalize()104        q = limpar_hellaswag(row['activity_label'] + ': ' + ctx)105        finais = [' ' + limpar_hellaswag(e) for e in row['endings']]106        docs.append((q, finais, int(row['label'])))107        if limite and len(docs) >= limite:108            break109    return docs110 111 112def arc(subset='ARC-Easy', limite=None):113    from datasets import load_dataset114    ds = load_dataset('allenai/ai2_arc', subset, split='test')115    docs = []116    for row in ds:117        rot, txt = row['choices']['label'], row['choices']['text']118        if row['answerKey'] not in rot:119            continue120        docs.append(('Question: ' + row['question'] + chr(10) + 'Answer:',121                     [' ' + t for t in txt], rot.index(row['answerKey'])))122        if limite and len(docs) >= limite:123            break124    return docs125 126 127# Numeros publicos do card oficial do SmolLM2-135M (HuggingFaceTB)128REF = {'hellaswag': 42.1, 'arc_media': 43.9}129 130EST = {'referencia': REF, 'modelo': 'HuggingFaceTB/SmolLM2-135M'}131t0 = time.time()132m = ModeloHF('HuggingFaceTB/SmolLM2-135M')133print('modelo carregado', round(time.time() - t0), 's', flush=True)134 135print('hellaswag...', flush=True)136EST['hellaswag'] = avaliar(m, hellaswag(), batch_size=32)137EST['hellaswag']['publicado'] = REF['hellaswag']138EST['hellaswag']['erro_relativo_pct'] = round(139    abs(EST['hellaswag']['acc_norm'] * 100 - REF['hellaswag']) / REF['hellaswag'] * 100, 3)140print('HELLASWAG acc_norm', round(EST['hellaswag']['acc_norm'] * 100, 2),141      'publicado', REF['hellaswag'], 'erro%', EST['hellaswag']['erro_relativo_pct'], flush=True)142SAIDA.write_text(json.dumps(EST, indent=1), encoding='utf-8')143 144for sub, chave in [('ARC-Easy', 'arc_easy'), ('ARC-Challenge', 'arc_challenge')]:145    EST[chave] = avaliar(m, arc(sub), batch_size=32)146    print(chave, 'acc_norm', round(EST[chave]['acc_norm'] * 100, 2), flush=True)147    SAIDA.write_text(json.dumps(EST, indent=1), encoding='utf-8')148 149media_arc = (EST['arc_easy']['acc_norm'] + EST['arc_challenge']['acc_norm']) / 2 * 100150EST['arc_media'] = {'acc_norm': media_arc, 'publicado': REF['arc_media'],151                    'erro_relativo_pct': round(abs(media_arc - REF['arc_media']) / REF['arc_media'] * 100, 3)}152print('ARC media', round(media_arc, 2), 'publicado', REF['arc_media'],153      'erro%', EST['arc_media']['erro_relativo_pct'], flush=True)154 155EST['validado_dentro_de_1pct'] = bool(156    EST['hellaswag']['erro_relativo_pct'] <= 1.0 or EST['arc_media']['erro_relativo_pct'] <= 1.0)157EST['fim'] = True158SAIDA.write_text(json.dumps(EST, indent=1), encoding='utf-8')159print('HARNESS', 'VALIDADO' if EST['validado_dentro_de_1pct'] else 'FORA DE 1%', flush=True)160