File size: 1,226 Bytes
6b5a5a6
7c5a283
 
ce814c7
7c5a283
 
6b5a5a6
b53e809
4ccb522
7c5a283
b53e809
 
7c5a283
b53e809
9c634ec
 
 
 
 
 
 
 
 
 
 
 
7c5a283
6b5a5a6
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
import streamlit as st
import torch
from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast
from utils import get_text, get_labels, load_model

count_labels = 8

st.markdown("## Классификатор статей")
st.markdown("Сервис классифицирует статьи по названию и аннотации. Нужно ввести в каждое окошко свою сущность и вам выдадут к какому классу относится статья")


title = st.text_area("Введите название статьи")

abstract = st.text_area("Введите аннотацию к статье, abstract статьи")
if st.button("Классифицируй!"):
  model, tokenizer = load_model()
  #model.load_state_dict(torch.load('weight_model'))
  
  text = get_text(title, abstract)
  if text:
    raw_predictions = get_labels(text, model, tokenizer)
    st.markdown("Список классов к которым может относится данная статья")
    for raw in raw_predictions:
      st.markdown(f"{raw}")
  else:
    st.markdown("Ваш запрос пуст. Введите хотя бы название")