Spaces:
Runtime error
Runtime error
import streamlit as st | |
import plotly.express as px | |
import torch | |
from torch import nn | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
deftxt = "I hate you cancerous insects so much" | |
txt = st.text_area('Text to analyze', deftxt) | |
# load tokenizer and model weights | |
tokenizer = AutoTokenizer.from_pretrained("s-nlp/roberta_toxicity_classifier") | |
model = AutoModelForSequenceClassification.from_pretrained("s-nlp/roberta_toxicity_classifier") | |
batch = tokenizer.encode(txt, return_tensors='pt') | |
# run model e.g. "logits": tensor([[ 4.8982, -5.1952]], grad_fn=<AddmmBackward0>) | |
result = model(batch) | |
# get probabilities e.g. tensor([[9.9996e-01, 4.2627e-05]], grad_fn=<SoftmaxBackward0>) | |
# first indice is neutral, second is toxic | |
prediction = nn.functional.softmax(result.logits, dim=-1) | |
neutralProb = prediction.data[0][0].double() | |
toxicProb = prediction.data[0][1].double() | |
neutralProb | |
toxicProb | |
print("Classification Probabilities") | |
print(f"Neutral: {neutralProb}") | |
print(f"Toxic: {toxicProb}") | |