File size: 1,359 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 | """Inferencia rapida usando um unico fold (fold 1). Ideal para demos.
Em CPU: ~5-10s por par. Em GPU T4: ~50ms.
"""
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_1/*"])
m = json.load(open(f"{path}/manifesto.json"))
tok = AutoTokenizer.from_pretrained(m["base_model"], padding_side="left")
model = AutoModelForCausalLM.from_pretrained(
m["base_model"], torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32)
model = PeftModel.from_pretrained(model, f"{path}/adapter_fold_1")
if torch.cuda.is_available(): model.cuda()
model.eval()
def score(tweet, nota):
text = (m["prompt_prefixo"] + "<Instruct>: " + m["instrucao"] +
"\n<Query>: " + tweet + "\n<Document>: " + nota + m["prompt_sufixo"])
enc = tok(text, return_tensors="pt", truncation=True, max_length=m["max_length"]).to(model.device)
with torch.no_grad():
logits = model(**enc).logits[:, -1, :]
return float(torch.sigmoid(logits[:, m["id_yes"]] - logits[:, m["id_no"]]).item())
print(score("Bolsonaro disse que a Terra e plana",
"Bolsonaro nunca afirmou isso; checagem em https://exemplo.org"))
|