home_work / app.py
arkmartov's picture
Update app.py
3de3a79
raw
history blame
1.94 kB
import streamlit as st
import torch
@st.cache
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')
model.load_state_dict(bn_state_dict)
return model, tokenizer
def Predict(model, tokenizer, text):
res = tokenizer(s, padding=True, truncation=True, return_tensors="pt", max_length=512)
#var.to("cuda:0")
res = model(**res)
logits = res.logits.softmax(dim=1)
logits = logits.numpy()[0]#logits.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)
sum, idx = 0, 0
while sum < 0.95:
st.markdown(f"{idx + 1}. ", dictionary[z[idx][1]])
sum += z[idx][0]
idx += 1
def filter(title, abstract):
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)
Print(logits, dictionary)