Nlp_ / lstm.py
osheina's picture
Update lstm.py
e01d67b verified
raw
history blame contribute delete
No virus
1.34 kB
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("Отзыв отрицательный 😞")