File size: 2,895 Bytes
b27d10d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
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)}")
|