Spaces:
Sleeping
Sleeping
File size: 924 Bytes
26ff02c 1180a53 26ff02c 1180a53 26ff02c 1180a53 26ff02c 1180a53 |
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 |
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
MODEL_NAME = "cardiffnlp/twitter-roberta-base-sentiment-latest"
LABELS = ["negative", "neutral", "positive"]
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
def test_model_loads():
assert model is not None
assert tokenizer is not None
def test_model_prediction_shape():
text = "I love this product!"
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
assert outputs.logits.shape[-1] == len(LABELS)
def test_sentiment_confidence():
text = "I hate this"
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
probs = torch.nn.functional.softmax(model(**inputs).logits, dim=-1)
assert torch.isclose(probs.sum(), torch.tensor(1.0), atol=1e-3)
|