import streamlit as st
import torch
import numpy as np

st.markdown("### A dummy site for classifying article topics by title and abstract.")
st.markdown("It can predict the following topics: Computer Science, Economics, Electrical Engineering and Systems Science, Mathematics, Quantitative Biology, Quantitative Finance, Statistics, Physics")


from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForSequenceClassification

@st.cache(suppress_st_warning=True)
def model_tokenizer():
    model_name = 'distilbert-base-cased'
    #tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased", problem_type="multi_label_classification")
    model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-cased", num_labels=8, problem_type="multi_label_classification")
    weights = torch.load('model.pt', map_location=torch.device('cpu'))
    model.load_state_dict(weights)
    return model#, tokenizer

def make_prediction(model, tokenizer, text):
    #print(text)
    tokens = tokenizer.encode(text)
    with torch.no_grad():
        logits = model.cpu()(torch.as_tensor([tokens]))[0]
        #print(logits)
        probs = np.array(torch.softmax(logits[-1, :], dim=-1))
        #print(probs)
    
    sorted_classes, sorted_probs = np.flip(np.argsort(probs)), sorted(probs, reverse=True)
    prediction_classes, prediction_probs = [], []
    probs_sum = 0
    i=0
    res = []
    while probs_sum <= 0.95:
#         print(i)
#         print(sorted_classes)
#         print(sorted_classes[i])
#         print(to_category)
#         print(sorted_classes[i], to_category[sorted_classes[i]])
        prediction_classes.append(to_category[sorted_classes[i]])
        prediction_probs.append(100*sorted_probs[i])
        probs_sum += sorted_probs[i]
        i += 1
    for pr, cl in zip(prediction_probs, prediction_classes):
        print(str("{:.2f}".format(pr) + "%"), cl)
        res.append((str("{:.2f}".format(pr) + "%"), cl))
    return res

model = model_tokenizer()
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased", problem_type="multi_label_classification")

categories_full = ['Computer Science', 'Economics', 'Electrical Engineering and Systems Science', 'Mathematics', 'Quantitative Biology', 'Quantitative Finance', 'Statistics', 'Physics']

to_category = {}

for i in range(len(categories_full)):
    to_category[i] = categories_full[i]



title = st.text_area("Type the title of the article here")
abstract = st.text_area("Type the abstract of the article here")


if st.button('Analyse'):
    if title or abstract:
        text = '[TITLE] ' + title + ' [ABSTRACT] ' + abstract
        res = make_prediction(model, tokenizer, text)
        for cat in res:
            st.markdown(f"{cat[0],  cat[1]}")
    else:
        st.error(f"Write title or abstract")

#st.markdown(f"{make_prediction(model, tokenizer, text)}")