Spaces:
Runtime error
Runtime error
| 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 | |
| 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)}") | |