ki33elev/Transformer_arxiv_classification

1import streamlit as st2import numpy as np3import pandas as pd4import torch5import transformers6import tokenizers7 8@st.cache(suppress_st_warning=True, hash_funcs={tokenizers.Tokenizer: lambda _: None})9def load_model():10    from transformers import AutoTokenizer, AutoModelForSequenceClassification11    model_name = 'distilbert-base-cased'12    tokenizer = AutoTokenizer.from_pretrained(model_name)13    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=8)14    model.load_state_dict(torch.load('model_weights2.pt', map_location=torch.device('cpu')))15    model.eval()16    return tokenizer, model17    18@st.cache(suppress_st_warning=True, hash_funcs={tokenizers.Tokenizer: lambda _: None})  19def predict(title, summary, tokenizer, model):20    text = title + "\n" + summary21    tokens = tokenizer.encode(text)22    with torch.no_grad():23        logits = model(torch.as_tensor([tokens]))[0]24        probs = torch.softmax(logits[-1, :], dim=-1).data.cpu().numpy()25    26    classes = np.flip(np.argsort(probs))27    sum_probs = 028    ind = 029    prediction = []30    prediction_probs = []31    while sum_probs < 0.95:32        prediction.append(label_to_theme[classes[ind]])33        prediction_probs.append(str("{:.2f}".format(100 * probs[classes[ind]])) + "%")34        sum_probs += probs[classes[ind]]35        ind += 136    37    return prediction, prediction_probs38   39@st.cache(suppress_st_warning=True) 40def get_results(prediction, prediction_probs):41    frame =  pd.DataFrame({'Category': prediction, 'Confidence': prediction_probs})42    frame.index = np.arange(1, len(frame) + 1)43    return frame44    45label_to_theme = {0: 'Computer science', 1: 'Economics', 2: 'Electrical Engineering and Systems Science', 3: 'Math',46                  4: 'Quantitative biology', 5: 'Quantitative Finance', 6: 'Statistics', 7: 'Physics'}47 48st.title("Arxiv articles classification")49st.markdown("<h1 style='text-align: center;'><img width=300px src='https://media.wired.com/photos/592700e3cfe0d93c474320f1/191:100/w_1200,h_630,c_limit/faces-icon.jpg'>", unsafe_allow_html=True)50st.markdown("This is an interface that can determine the article's category based on its title and summary. Though it can work with title only, it is recommended that you provide summary if possible - this will result in a better prediction quality.")51 52tokenizer, model = load_model()53 54title = st.text_area(label='Title', height=100)55summary = st.text_area(label='Summary (optional)', height=250)56button = st.button('Run')57 58if button:59    prediction, prediction_probs = predict(title, summary, tokenizer, model)60    ans = get_results(prediction, prediction_probs)61    if len(title + "\n" + summary) < 20:62        st.error("Your input is too short. It is probably not a real article, please try again.")63    else:64        st.subheader('Results:')65        st.write(ans)