import streamlit as st

st.markdown("### Article Classifier")
st.markdown("<img width=200px src='https://media.istockphoto.com/photos/funny-cat-is-studying-chemistry-picture-id526831620'>", unsafe_allow_html=True)
st.markdown("This is a tool for classifying article category by it's title and summary. \n Follow the instructions below")
from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
import torch
import numpy as np
tokenizer = AutoTokenizer.from_pretrained("Wi/arxiv-topics-distilbert-base-cased")
model = AutoModelForSequenceClassification.from_pretrained("Wi/arxiv-topics-distilbert-base-cased")

title = st.text_area("Put the title here")
abstract = st.text_area("Put the abstract here")
if st.button('Press when ready'):
    text = 'Title:' + title + '\n' + 'Abstract:' + abstract
    
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits

    probs = logits.softmax(dim=-1).detach().cpu().flatten().numpy().tolist()
    order = np.argsort(probs)[::-1]
    i = 0
    sum = 0
    predicted_class_id = []
    while sum < 0.95:
        predicted_class_id.append(order[i])
        sum += probs[order[i]]
        i+=1
    

    for id in predicted_class_id:
        st.markdown(model.config.id2label[id])