import streamlit as st
from transformers import pipeline
import spacy
from spacy import displacy
import plotly.express as px
import numpy as np
st.set_page_config(page_title="Text Classification")
st.title("Text Classification'")
st.write("_This web application is intended for educational use, please do not upload any sensitive information._")
st.write("Placing a piece of text into one or more categories.")

@st.cache(allow_output_mutation=True, show_spinner=False)
def Loading_Classifier():
    class1 = pipeline("zero-shot-classification",framework="pt")
    return class1

def plot_result(top_topics, scores):
    top_topics = np.array(top_topics)
    scores = np.array(scores)
    scores *= 100
    fig = px.bar(x=scores, y=top_topics, orientation='h',
                 labels={'x': 'Probability', 'y': 'Category'},
                 text=scores,
                 range_x=(0,115),
                 title='Top Predictions',
                 color=np.linspace(0,1,len(scores)),
                 color_continuous_scale="Bluered")
    fig.update(layout_coloraxis_showscale=False)
    fig.update_traces(texttemplate='%{text:0.1f}%', textposition='outside')
    st.plotly_chart(fig)

with st.spinner(text="Please wait for the models to load. This could take up to 60 seconds."):
    class1 = Loading_Classifier()

cat1 = st.text_input('Enter each possible category name (separated by a comma). Maximum 5 categories.')
text = st.text_area('Enter Text Below:', height=200)
submit = st.button('Generate')
if submit:
    st.subheader("Classification Results:")
    labels1 = cat1.strip().split(',')
    result = class1(text, candidate_labels=labels1)
    cat1name = result['labels'][0]
    cat1prob = result['scores'][0]
    st.write('Category: {} | Probability: {:.1f}%'.format(cat1name,(cat1prob*100)))
    plot_result(result['labels'][::-1][-10:], result['scores'][::-1][-10:])