import streamlit as st
from transformers import pipeline
import tokenizers

@st.cache(hash_funcs={tokenizers.Tokenizer: lambda _: None, tokenizers.AddedToken: lambda _: None}, allow_output_mutation=True)
def get_pipe():
    return pipeline("sentiment-analysis", "tmp/arxiv")

pipe = get_pipe()

labels = {
    0: "cs.AI (Artificial Intelligence)",
    1: "cs.CE (Computational Engineering)",
    2: "cs.CV (Computer Vision)",
    3: "cs.DS (Data Structures)",
    4: "cs.IT (Information Theory )",
    5: "cs.NE (Neural and Evolutionary)",
    6: "cs.PL (Programming Languages)",
    7: "cs.SY (Systems and Control)",
    8: "math.AC (Commutative Algebra )",
    9: "math.GR (Group Theory)",
    10: "math.ST (Statistics Theory)"
}


st.markdown("# Try this SOTA article text classifier for free!")
st.markdown("<img width=300px src='https://incrussia.ru/wp-content/uploads/2018/12/Snimok-ekrana-2018-12-12-v-16aaa.jpg'>", unsafe_allow_html=True)

st.markdown("Input example: \n\n Title: \n\nAttention Based Relation Network for Facial Action Units Recognition \n\n Summary: \n\nFacial action unit (AU) recognition is essential to facial expression analysis. Since there are highly positive or negative correlations between AUs, some existing AU recognition works have focused on modeling AU relations. However, previous relationship-based approaches typically embed predefined rules into their models and ignore the impact of various AU relations in different crowds. In this paper, we propose a novel Attention Based Relation Network (ABRNet) for AU recognition, which can automatically capture AU relations without unnecessary or even disturbing predefined rules. ABRNet uses several relation learning layers to automatically capture different AU relations. The learned AU relation features are then fed into a self-attention fusion module, which aims to refine individual AU features with attention weights to enhance the feature robustness. Furthermore, we propose an AU relation dropout strategy and AU relation loss (AUR-Loss) to better model AU relations, which can further improve AU recognition. Extensive experiments show that our approach achieves state-of-the-art performance on the DISFA and DISFA+ datasets. ")


title = st.text_area("Enter article title:", height=20)
summary = st.text_area("Enter article summary:", height=50)

if st.button('Run'):
    input_text = title + " Abstract " + summary
    raw_predictions = pipe(input_text, return_all_scores=True)
    sorted_preds = sorted(raw_predictions[0], key=lambda x: x['score'], reverse=True)

    if len(input_text) > 1:
        s = 0
        i = 0
        while s < 0.95:
            label, score = sorted_preds[i]['label'], sorted_preds[i]['score']
            string = f'{labels[label]} {score}'
            st.write(string)
            s += score
            i += 1

else:
    st.write('Press a run button...') #displayed when the button is unclicked