I am working on SequenceClassification task and dont know how to see the predictions
here is my code
from transformers import DistilBertTokenizer, TFDistilBertForMaskedLM
import tensorflow as tf
import torch
loaded_model = DistilBertForSequenceClassification.from_pretrained(“/content/results/distillbert/model1”,return_dict=True)
tokenizer = DistilBertTokenizer.from_pretrained(‘distilbert-base-cased’)
inputs = tokenizer(“Hello, this news is Not good for learning”, return_tensors=“pt”)
labels = torch.tensor([1]).unsqueeze(0) # Batch size 1
outputs = loaded_model(**inputs, labels=labels)
loss = outputs.loss
logits = outputs.logits
Checkout this docs.
First, you won’t need to pass labels
for predictions and for prediction put the code under torch.no_grad()
to avoid calculating gradients.
The returned logits
have shape [batch_size, num_classes]
, you can then apply argmax
to get the index highest scored class.
classes = ["NEGATIVE", "POSITIVE"]
enc = tokenizer("positive text", return_tensors="pt")
with torch.no_grad():
logits = model(**enc, return_dict=True).logits
pred_class_idx = torch.argmax(logits).item()
classes[pred_class_idx]
you can also apply softmax
on logits
to get the probabilities.
and better yet, just use pipeline
as shown in the docs above, which does all of this for you.
1 Like
solved, very useful
thanks