import streamlit as st
from joblib import load
from transformers import BertTokenizer, BertForSequenceClassification
import torch
from tensorflow.keras.models import load_model
import tensorflow as tf
from tensorflow.keras.preprocessing.text import Tokenizer
from tensorflow.keras.preprocessing.sequence import pad_sequences
import time

clf = load('my_model_filename.pkl')
vectorizer = load('tfidf_vectorizer.pkl')
scaler = load('scaler.joblib')

tukinazor = load('tokenizer.pkl')
rnn_model = load_model('path_to_my_model.h5')
bert_model = BertForSequenceClassification.from_pretrained('my_bert_model')
tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
bert_model = bert_model.to(device)

def predict_text(text):
    sequences = tukinazor.texts_to_sequences([text])
    padded_sequences = tf.keras.preprocessing.sequence.pad_sequences(sequences, maxlen=200, padding='post', truncating='post')
    predictions = rnn_model.predict(padded_sequences)
    predicted_class = tf.argmax(predictions, axis=-1).numpy()[0]
    return predicted_class


# Запуск приложения
def main():
    st.title("Модель классификации отзывов")

    # Ввод текста
    user_input = st.text_area("Введите текст отзыва:")

    if st.button("Классифицировать"):
        start_time = time.time()
        user_input_vec = vectorizer.transform([user_input])
        sentence_vector_scaled = scaler.transform(user_input_vec)
        prediction = clf.predict(
            sentence_vector_scaled)
        elapsed_time = time.time() - start_time
        st.write(f"Прогнозируемый класс: {prediction[0]}")
        st.write(f"Время вычисления: {elapsed_time:.2f} сек.")

    user_input_rnn = st.text_area("Введите текст отзыва для Keras RNN модели:")

    if st.button("Классифицировать с RNN"):
        start_time = time.time()
        prediction_rnn = predict_text(user_input_rnn)
        elapsed_time = time.time() - start_time
        st.write(f"Прогнозируемый класс с RNN: {prediction_rnn}")
        st.write(f"Время вычисления: {elapsed_time:.2f} сек.")

    user_input_bert = st.text_area("Введите текст отзыва для BERT:")

    if st.button("Классифицировать (BERT)"):
        start_time = time.time()
        encoding = tokenizer.encode_plus(
            user_input_bert,
            add_special_tokens=True,
            max_length=200,
            return_token_type_ids=False,
            padding='max_length',
            truncation=True,
            return_attention_mask=True,
            return_tensors='pt'
        )
        input_ids = encoding['input_ids'].to(device)
        attention_mask = encoding['attention_mask'].to(device)

        with torch.no_grad():
            outputs = bert_model(input_ids=input_ids, attention_mask=attention_mask)
            predictions = torch.argmax(outputs.logits, dim=1)
            elapsed_time = time.time() - start_time
            st.write(f"Прогнозируемый класс (BERT): {predictions.item() + 1}")
            st.write(f"Время вычисления: {elapsed_time:.2f} сек.")

if __name__ == "__main__":
    main()