File size: 776 Bytes
1a5f890
bf79d0d
9a33247
 
 
bf79d0d
2268b75
7cf8aab
 
2268b75
 
6767b70
 
7186326
2268b75
9a33247
bf79d0d
9a33247
 
 
 
 
bf79d0d
19c4297
 
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
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')

# e.g. "logits":"tensor([[ 4.8982, -5.1952]], grad_fn=<AddmmBackward0>)"
result = model(batch)

# get probabilities
prediction = nn.functional.softmax(result.logits, dim=-1)

print(prediction)

#fig = px.bar(result, x="", y="", orientation='h')
#fig.show()