Vinodel/classify_papers

1import streamlit as st2import torch3from torch import nn4from transformers import pipeline5from metric import top_956 7st.markdown('### Привет! Я нейросеть, которая распознает тематику научной статьи по ее названию и описанию.')8st.markdown('<img width=200px src="https://sci-fi-news.ru/wp-content/uploads/2017/06/Transformers-5-poster-Optimus-Prime-850x545.jpg">', unsafe_allow_html=True)9 10 11labels = {12    'LABEL_0' : 'Computer Science',13    'LABEL_1' : 'Economics',14    'LABEL_2' : 'Electrical Engineering and Systems Science',15    'LABEL_3' : 'Mathematics',16    'LABEL_4' : 'Physics',17    'LABEL_5' : 'Quantitative Biology',18    'LABEL_6' : 'Quantitative Finance',19    'LABEL_7' : 'Statistics'20}21 22@st.cache(allow_output_mutation=True)23def get_model():24    params = torch.load('params_finetune.pth', map_location=torch.device('cpu'))25 26    classifier = pipeline('sentiment-analysis', model="distilbert-base-cased")27    classifier.model.classifier = nn.Linear(in_features=768, out_features=8, bias=True)28    classifier.model.config.num_labels = 829    classifier.model.load_state_dict(params)30    return classifier31 32classifier = get_model()33 34title = st.text_area("Название и описание статьи. Их можно вводить в произвольном порядке.")35 36if title:37    predict = classifier(title, return_all_scores=True)[0]38    pred = list(map(lambda x: (x['label'], x['score']), predict))39    result = []40 41    for el in pred:42        result.append((labels[el[0]], el[1]))43    result = top_95(result)44    st.markdown('Думаю, что ваша статья подходит под следующие топики: ' + ', '.join(result))