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%})")