|
import streamlit as st
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
|
|
import json
|
|
import numpy as np
|
|
|
|
|
|
|
|
@st.cache_data
|
|
def load_subject_dict():
|
|
with open('subject_dict.json') as json_file:
|
|
subject_dict = json.load(json_file)
|
|
return subject_dict
|
|
|
|
@st.cache_resource
|
|
def load_model_and_tokenizer():
|
|
model_name = 'DeepPavlov/rubert-base-cased'
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
model = torch.load('rubert_26650.pt', weights_only=False)
|
|
return model, tokenizer
|
|
|
|
model, tokenizer = load_model_and_tokenizer()
|
|
subject_dict = load_subject_dict()
|
|
|
|
st.title("Прогнозирование темы по условию математической задачи")
|
|
|
|
query = st.text_input("Текст задачи", value="Все зебры полосатые. Все полосатые животные - веселые. Верно ли, что все зебры веселые?")
|
|
|
|
st.header("Предлагаемые темы:")
|
|
if query:
|
|
tokens_info = tokenizer(
|
|
query,
|
|
padding=True,
|
|
truncation=False,
|
|
return_tensors="pt",
|
|
)
|
|
|
|
model.eval()
|
|
model.cpu()
|
|
with torch.no_grad():
|
|
out = model(**tokens_info)
|
|
probs = torch.nn.functional.softmax(out.logits, dim=-1)
|
|
top_k = 3
|
|
if top_k:
|
|
best_prob, best_idx = torch.topk(probs, top_k, dim=-1)
|
|
best_prob = best_prob[0].numpy()
|
|
best_prob /= best_prob.sum()
|
|
best_idx = best_idx[0].numpy()
|
|
for i in range(top_k):
|
|
st.write(f"**{subject_dict[str(best_idx[i])]}** (с вероятностью {best_prob[i]:2.1%})")
|
|
|
|
|
