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]
toxicProb = prediction.data[0][1]

neutralProb = torch.round(neutralProb, decimals=4)
toxicProb = torch.round(toxicProb, decimals=4)

# default text input ought to return:
# Neutral: 0.0052
# Toxic: 0.9948

st.write("Classification Probabilities")
st.write(f"{neutralProb:.4}  -  NEUTRAL")
st.write(f"{toxicProb:.4}  -  TOXIC")