import streamlit as st
import torch
from torch import nn
from transformers import pipeline
from metric import top_95

st.markdown('### Привет! Я нейросеть, которая распознает тематику научной статьи по ее названию и описанию.')
st.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)


labels = {
    'LABEL_0' : 'Computer Science',
    'LABEL_1' : 'Economics',
    'LABEL_2' : 'Electrical Engineering and Systems Science',
    'LABEL_3' : 'Mathematics',
    'LABEL_4' : 'Physics',
    'LABEL_5' : 'Quantitative Biology',
    'LABEL_6' : 'Quantitative Finance',
    'LABEL_7' : 'Statistics'
}

@st.cache(allow_output_mutation=True)
def get_model():
    params = torch.load('params_finetune.pth', map_location=torch.device('cpu'))

    classifier = pipeline('sentiment-analysis', model="distilbert-base-cased")
    classifier.model.classifier = nn.Linear(in_features=768, out_features=8, bias=True)
    classifier.model.config.num_labels = 8
    classifier.model.load_state_dict(params)
    return classifier

classifier = get_model()

title = st.text_area("Название и описание статьи. Их можно вводить в произвольном порядке.")

if title:
    predict = classifier(title, return_all_scores=True)[0]
    pred = list(map(lambda x: (x['label'], x['score']), predict))
    result = []

    for el in pred:
        result.append((labels[el[0]], el[1]))
    result = top_95(result)
    st.markdown('Думаю, что ваша статья подходит под следующие топики: ' + ', '.join(result))