File size: 1,863 Bytes
22c0b81 |
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 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
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%})")
|