File size: 1,681 Bytes
4e70841 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 | """Ensemble de probas: usa os 5 folds e devolve a media.
E a melhor pratica cientifica - replica o numero 0.7920 macro-F1 reportado.
Em GPU T4: ~250ms por par. Em CPU: ~30s por par.
"""
import json, torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
from huggingface_hub import snapshot_download
REPO = "histlearn/community-notes-reranker-ptbr"
path = snapshot_download(REPO, allow_patterns=["manifesto.json", "adapter_fold_*/*"])
m = json.load(open(f"{path}/manifesto.json"))
tok = AutoTokenizer.from_pretrained(m["base_model"], padding_side="left")
base = AutoModelForCausalLM.from_pretrained(
m["base_model"], torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
if torch.cuda.is_available(): base.cuda()
def make_text(tw, nt):
return (m["prompt_prefixo"] + "<Instruct>: " + m["instrucao"] +
"\n<Query>: " + tw + "\n<Document>: " + nt + m["prompt_sufixo"])
def score_ensemble(tweet, nota):
probs = []
for k in range(1, 6):
model = PeftModel.from_pretrained(base, f"{path}/adapter_fold_{k}")
model.eval()
enc = tok(make_text(tweet, nota), return_tensors="pt",
truncation=True, max_length=m["max_length"]).to(model.device)
with torch.no_grad():
logits = model(**enc).logits[:, -1, :]
probs.append(float(torch.sigmoid(
logits[:, m["id_yes"]] - logits[:, m["id_no"]]).item()))
model.unload() # libera memoria do adapter
return sum(probs) / 5
print(score_ensemble("Bolsonaro disse que a Terra e plana",
"Bolsonaro nunca afirmou isso; checagem em https://exemplo.org"))
|