import streamlit as st
import torch
import re
import json
import time
from nltk.corpus import stopwords
from Models.model_file import data_preprocessing, preprocess_single_string, LSTMBahdanauAttention
from nltk.corpus import stopwords
stop_words = set(stopwords.words('russian'))

with open('Weights/vocab_to_int.json', 'r') as file:
    vocab_to_int = json.load(file)

SEQ_LEN = 96
model_bah = LSTMBahdanauAttention()
model_bah.load_state_dict(torch.load('Weights/final_model_bah.pth'))
model_bah.eval()

def analyze_sentiment(text):
    preprocessed_text = data_preprocessing(text)
    sample = preprocess_single_string(preprocessed_text, SEQ_LEN, vocab_to_int)
    with torch.no_grad():
        probability = model_bah(sample.unsqueeze(0))[0].sigmoid().item()
        return probability

def lstm_model_page():
    st.title("Классификация отзывов")
    user_input = st.text_area("Введите ваш отзыв:")
    if st.button("Классифицировать"):
        start_time = time.time()
        probability = analyze_sentiment(user_input)
        end_time = time.time()
        execution_time = end_time - start_time
        if probability > 0.5:
            st.write("Отзыв положительный 🌟")
        else:
            st.write("Отзыв отрицательный 😞")
        st.write(f'Время предсказания: {execution_time:.4f} секунд')