import streamlit as st
import torch
from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast
from utils import get_text, get_label

count_labels = 8

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


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

abstract = st.text_area("Введите аннотацию к статье, abstract статьи")

tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-cased')
model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-cased", num_labels=count_labels)
model.load_state_dict(torch.load('weight_model'))

text = get_text(title, abstract)
if text:
  raw_predictions = get_labels(text, model, tokenizer)
  st.markdown(f"{raw_predictions}")
else:
  st.markdown("Ваш запрос пуст. Введите хотя бы название")