xlm-roberta-large-english-cap-v3

Este fragmento imprime las tres etiquetas más probables y sus correspondientes puntuaciones de softmax:

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForSequenceClassification.from_pretrained("poltextlab/xlm-roberta-large-english-cap-v3")
tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-large")

sentence = "This is an example."

inputs = tokenizer(sentence,
return_tensors="pt",
max_length=512,
padding="do_not_pad",
truncation=True
)

logits = model(**inputs).logits

probs = torch.softmax(logits, dim=1).tolist()[0]
probs = {model.config.id2label[index]: round(probability, 2) for index, probability in enumerate(probs)}
top3_probs = dict(sorted(probs.items(), key=lambda item: item[1], reverse=True)[:3])

print(top3_probs)