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

# Load vocabulary mapping
with open('vocab_to_int.json', 'r') as file:
    vocab_to_int = json.load(file)

# Load the pre-trained model
SEQ_LEN = 96
model_bah = LSTMBahdanauAttention()
# Set the new vocabulary size in the model
model_bah.load_state_dict(torch.load('final_model_bah.pth'))
model_bah.eval()

# Function to analyze sentiment
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

# Streamlit UI
def lstm_model_page():
    st.title("Классификация отзывов")
    user_input = st.text_area("Введите ваш отзыв:")
    if st.button("Классифицировать"):
        probability = analyze_sentiment(user_input)
        if probability > 0.5:
            st.write("Отзыв положительный 🌟")
        else:
            st.write("Отзыв отрицательный 😞")