omega-squad-v2

System Omega — SQuAD v2 (ModernBERT-base, 150M, dual-head)

Modelo de QA extrativo com duas cabeças: decisão (a pergunta é respondível no contexto?) e extração de span (onde está a resposta). Teste 50/50 de answerability no SQuAD v2: acurácia de decisão de 87.71%, à frente do Laya (0.766) e do Jev 1.13 (0.727) — com menos da metade dos parâmetros do Laya (421M) e ~4x menos latência.

Resultados (teste 50/50 próprio, 7199 amostras)

Formato de entrada (importante!)

O modelo espera: INSTRUCTION: <pergunta> | CONTEXT: <contexto>

Arquitetura

ModernBERT-base + cabeça de decisão ([mean-pool; CLS] -> MLP) + cabeça de span (Linear(hidden, 2) sobre cada token). Decodificação por argmax na matriz de pares (start, end) válidos; (0,0) = sem resposta. Pesos em bf16.

Uso

modeling_omega.py (neste repositório) contém a classe, o decode_spans e o helper load_omega. Baixe-o para o mesmo diretório e então:

import torch
from modeling_omega import load_omega, decode_spans
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("Kodjaoglanian/omega-squad-v2")
model = load_omega("Kodjaoglanian/omega-squad-v2").eval().to("cuda")

contexto = "The Eiffel Tower is a wrought-iron lattice tower in Paris, France."
pergunta = "Where is the Eiffel Tower located?"
texto = "INSTRUCTION: " + pergunta + " | CONTEXT: " + contexto

enc = tok(texto, return_tensors="pt", truncation=True, max_length=512)
ids, mask = enc["input_ids"].to("cuda"), enc["attention_mask"].to("cuda")
with torch.amp.autocast("cuda", dtype=torch.bfloat16):
    cat, s_log, e_log = model(ids, mask)
s, e = decode_spans(s_log, e_log, mask)

if int(cat.argmax()) == 1 and int(s[0]) > 0:
    print(tok.decode(ids[0][int(s[0]):int(e[0]) + 1], skip_special_tokens=True))
else:
    print("(sem resposta)")

Treino

• SQuAD v2 oficial, 48,000 amostras balanceadas 50/50 (40,800 treino / 7,200 teste)
• 4 épocas, batch 64, lr 5e-05, warmup 5% + decaimento linear
• Loss = CE(decisão, label smoothing 0.05) + (CE(start) + CE(end)) / 2
• Length bucketing + trim dinâmico, bf16 autocast, seed 42

Nota de benchmark

O teste 50/50 de answerability é o protocolo interno do System Omega. A comparação com Laya/Jev usa os números públicos deles (typed-decisions); latência, tamanho e classe de arquitetura (encoder, 1 forward pass) são comparáveis entre os três.