home_work / app.py
arkmartov's picture
Update app.py
bb61482
import streamlit as st
import torch
import numpy as np
@st.cache(allow_output_mutation=True)
def Model():
from transformers import DebertaTokenizer, DebertaForSequenceClassification
tokenizer = DebertaTokenizer.from_pretrained("microsoft/deberta-base")
model = DebertaForSequenceClassification.from_pretrained("microsoft/deberta-base", num_labels=8)
bn_state_dict = torch.load('model_weights.pt', map_location=torch.device('cpu'))
model.load_state_dict(bn_state_dict)
return model, tokenizer
def Predict(model, tokenizer, text):
res = tokenizer(text, padding=True, truncation=True, return_tensors="pt", max_length=512)
res = model(**res)
logits = res.logits.softmax(dim=1)
logits = logits.detach().numpy()[0]#.cpu().detach().numpy()[0]
return logits
def Print(logits, dictionary):
z = zip(logits, np.arange(0, 8))
z = sorted(z, key=lambda x: x[0], reverse=True)
summ, idx = 0, 0
while summ < 0.95:
string = dictionary[z[idx][1]]
st.markdown(f"{idx + 1}. {string}")
summ += z[idx][0]
idx += 1
def filter(title, abstract):
if len(title) < 10 or (len(abstract) > 0 and len(abstract) < 20):
st.markdown("Хммм... Вы точно не ошиблись? Слишком маленькое название или описание.")
return False
return True
st.title('Классификация статьи по названию и описанию')
# ^-- можно показывать пользователю текст, картинки, ограниченное подмножество html - всё как в jupyter
title = st.text_area("Введите название статьи:")
abstract = st.text_area("Введите описание статьи:")
# ^-- показать текстовое поле. В поле text лежит строка, которая находится там в данный момент
text = title + '. ' + abstract
dictionary = ['computer science', 'economics', 'Electrical Engineering and Systems Science',
'math', 'physics', 'quantitative biology', 'quantitative finance',
'statistics']
if filter(title, abstract):
model, tokenizer = Model()
logits = Predict(model, tokenizer, text)
st.header("Топ 95%:")
Print(logits, dictionary)