import re
import streamlit as st
import torch
st.title("film_review")

input_text = st.text_area("Enter your text")
from pages.film_review.model.model_lstm import *
from pages.film_review.model.model_logreg import *
from pages.film_review.model.model_bert import *


import time

class Timer:
    def __enter__(self):
        self.start_time = time.time()
        return self

    def __exit__(self, *args):
        self.end_time = time.time()
        self.execution_time = self.end_time - self.start_time


@st.cache_resource
def get_model():
    return  torch.load("pages/film_review/model/model_lstm.pt",map_location=torch.device('cpu'))
model = get_model()
model.eval()
dec = {0:'отрицательный',1:'нейтральный',2:'положительный'}

if input_text:
    
    with Timer() as t:
        with torch.no_grad():
            ans = torch.nn.functional.softmax(model(input_text), dim=1)
            idx = torch.argmax(ans, dim=1).item()
        st.write(f'LSTM - отзыв: {dec[idx]}, уверенность: { round(ans[0][idx].item(),2)}')
    st.write("Время выполнения:", round(t.execution_time*1000, 2), "миллисекунд")
    
    st.write("------------")
    with Timer() as t:
        st.write(f'Logreg - отзыв: {dec[ predict_tfidf(input_text)[0]]}')
    st.write("Время выполнения:", round(t.execution_time*1000, 2), "миллисекунд")

    st.write("------------")
    with Timer() as t:
        st.write(f'Bert - отзыв: {dec[ predict_bert(input_text)]}')
    st.write("Время выполнения:", round(t.execution_time*1000, 2), "миллисекунд")