Predictions for sequenceclassification task

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()

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