import streamlit as st
import torch
import plotly.express as px
from transformers import AutoTokenizer, AutoModelForSequenceClassification

txt = st.text_area('Text to analyze', '''
    It was the best of times, it was the worst of times, it was
    the age of wisdom, it was the age of foolishness, it was
    the epoch of belief, it was the epoch of incredulity, it
    was the season of Light, it was the season of Darkness, it
    was the spring of hope, it was the winter of despair, (...)
    ''')

# load tokenizer and model weights
tokenizer = AutoTokenizer.from_pretrained("s-nlp/roberta_toxicity_classifier")
model = AutoModelForSequenceClassification.from_pretrained("s-nlp/roberta_toxicity_classifier")

# prepare the input
batch = tokenizer.encode('txt', return_tensors='pt')

# inference
result = model(batch)

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