import streamlit as st import torch from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification from safetensors.torch import load_file as safe_load target_to_ind = {'cs': 0, 'econ': 1, 'eess': 2, 'math': 3, 'phys': 4, 'q-bio': 5, 'q-fin': 6, 'stat': 7} target_to_label = {'cs': 'Computer Science', 'econ': 'Economics', 'eess': 'Electrical Engineering and Systems Science', 'math': 'Mathematics', 'phys': 'Physics', 'q-bio': 'Quantitative Biology', 'q-fin': 'Quantitative Finance', 'stat': 'Statistics'} ind_to_target = {ind: target for target, ind in target_to_ind.items()} ![image/gif](https://cdn-uploads.huggingface.co/production/uploads/67f0e021e17e5360e7c18a52/oNcjUz7MbWW0Bymah31y7.gif) st.title('papers_classifier - your best friend in 🤓') st.text("You need to give me paper's title and (if you want to) it's abstract. Also you need to choose classification mode - there are 2 of them:\ best prediction and top 95% which means that you'll see as many classes as model needs to show to be confident with probability at least 0.95 that the correct one is among them. \ After that you need to press the Get prediction button and I'll tell you to which fields of study this paper is related. \ According to arXiv there are 8 different fields - Computer Science, Economics, Electrical Engineering and Systems Science, Mathematics, Physics, Quantitative Biology, \ Quantitative Finance and Statistics.\ ") @st.cache_resource def load_model_and_tokenizer(): model_name = 'distilbert/distilbert-base-cased' tokenizer = AutoTokenizer.from_pretrained(model_name) model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=len(target_to_ind)) state_dict = safe_load("model.safetensors") model.load_state_dict(state_dict) return model, tokenizer model, tokenizer = load_model_and_tokenizer() def get_predict(title: str, abstract: str) -> (str, float, dict): text = [title + tokenizer.sep_token + abstract[:128]] tokens_info = tokenizer( text, padding=True, truncation=True, return_tensors="pt", ) with torch.no_grad(): out = model(**tokens_info) probs = torch.nn.functional.softmax(out.logits, dim=-1).tolist()[0] return list(sorted([(p, ind_to_target[i]) for i, p in enumerate(probs)]))[::-1] title = st.text_area("Title ", "", height=100) abstract = st.text_area("Abstract ", "", height=150) mode = st.radio("Mode: ", ("Best prediction", "Top 95%")) if st.button("Get prediction", key="manual"): if len(title) == 0: st.error("Please, provide paper's title") else: with st.spinner("Be patient, I'm doing my best"): predict = get_predict(title, abstract) tags = [] threshold = 0 if mode == "Best prediction" else 0.95 sum_p = 0 for p, tag in predict: sum_p += p tags.append(target_to_label[tag]) if sum_p >= threshold: break tags = '\n'.join(tags) st.success(tags)