Streamlit adjustment

pages/review_predictor.py

@@ -7,37 +7,58 @@ import torch
 import torch.nn as nn
 import transformers
 
-from model.funcs import (create_model_and_tokenizer, load_model,
-                         predict_sentiment)
+from model.funcs import (create_model_and_tokenizer, execution_time,
+                         load_model, predict_sentiment)
 from model.model import LSTMConcatAttentionEmbed
 from preprocessing.preprocessing import data_preprocessing
 from preprocessing.rnn_preprocessing import preprocess_single_string
 
-# Load preprocessing steps
-with open("vectorizer.pkl", "rb") as f:
-    logreg_vectorizer = pickle.load(f)
 
-# Load trained model
-with open("logreg_model.pkl", "rb") as f:
-    logreg_predictor = pickle.load(f)
+@st.cache_resource
+def load_logreg():
+    with open("vectorizer.pkl", "rb") as f:
+        logreg_vectorizer = pickle.load(f)
 
-model_concat_embed = LSTMConcatAttentionEmbed()
-model_concat_embed.load_state_dict(torch.load("model/model_weights.pt"))
+    with open("logreg_model.pkl", "rb") as f:
+        logreg_predictor = pickle.load(f)
+    return logreg_vectorizer, logreg_predictor
 
-with open("model/vocab.json", "r") as f:
-    vocab_to_int = json.load(f)
 
-with open("model/int_vocab.json", "r") as f:
-    int_to_vocab = json.load(f)
+logreg_vectorizer, logreg_predictor = load_logreg()
 
-model_class = transformers.AutoModel
-tokenizer_class = transformers.AutoTokenizer
-pretrained_weights = "cointegrated/rubert-tiny2"
-weights_path = "model/best_bert_weights.pth"
-model = load_model(model_class, pretrained_weights, weights_path)
-tokenizer = tokenizer_class.from_pretrained(pretrained_weights)
 
+@st.cache_resource
+def load_lstm():
+    with open("model/vocab.json", "r") as f:
+        vocab_to_int = json.load(f)
 
+    with open("model/int_vocab.json", "r") as f:
+        int_to_vocab = json.load(f)
+    model_concat_embed = LSTMConcatAttentionEmbed()
+    model_concat_embed.load_state_dict(torch.load("model/model_weights.pt"))
+
+    return vocab_to_int, int_to_vocab, model_concat_embed
+
+
+vocab_to_int, int_to_vocab, model_concat_embed = load_lstm()
+
+
+@st.cache_resource
+def load_bert():
+    model_class = transformers.AutoModel
+    tokenizer_class = transformers.AutoTokenizer
+    pretrained_weights = "cointegrated/rubert-tiny2"
+    weights_path = "model/best_bert_weights.pth"
+    model = load_model(model_class, pretrained_weights, weights_path)
+    tokenizer = tokenizer_class.from_pretrained(pretrained_weights)
+
+    return model, tokenizer
+
+
+model, tokenizer = load_bert()
+
+
+@execution_time
 def plot_and_predict(review: str, SEQ_LEN: int, model: nn.Module):
     inp = preprocess_single_string(review, SEQ_LEN, vocab_to_int)
     model.eval()
@@ -52,12 +73,12 @@ def preprocess_text_logreg(text):
     clean_text = data_preprocessing(
         text
     )  # Assuming data_preprocessing is your preprocessing function
-    print("Clean text ", clean_text)
     vectorized_text = logreg_vectorizer.transform([" ".join(clean_text)])
     return vectorized_text
 
 
 # Define function for making predictions
+@execution_time
 def predict_sentiment_logreg(text):
     # Preprocess input text
     processed_text = preprocess_text_logreg(text)
@@ -68,7 +89,7 @@ def predict_sentiment_logreg(text):
 
 metrics = {
     "Models": ["Logistic Regression", "LSTM + attention", "ruBERTtiny2"],
-    "f1-macro score": [0.94376, 1, 0.94070],
+    "f1-macro score": [0.94376, 0.93317, 0.94070],
 }
 
 
@@ -94,7 +115,6 @@ if st.button("Predict"):
         )
     elif model_type == "BERT":
         prediction = predict_sentiment(text_input, model, tokenizer, "cpu")
-        st.write(prediction)
 
     if prediction == 1:
         st.write("prediction")

pages/text_generator.py

@@ -2,6 +2,8 @@ import streamlit as st
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
+from model.funcs import execution_time
+
 
 @st.cache_data
 def load_model():
@@ -13,11 +15,10 @@ def load_model():
 
 
 tokenizer, model = load_model()
-promt = st.text_input("Ask a question")
-generate = st.button("Generate")
-if generate:
-    if not promt:
-        st.write("42")
+
+
+@execution_time
+def generate_text(promt):
     promt = tokenizer.encode(promt, return_tensors="pt")
     model.eval()
     with torch.no_grad():
@@ -27,6 +28,16 @@ if generate:
             num_beams=2,
             temperature=1.5,
             top_p=0.9,
+            max_length=150,
         )
     out = list(map(tokenizer.decode, out))[0]
-    st.write(out)
+    return out
+
+
+promt = st.text_input("Ask a question")
+generate = st.button("Generate")
+if generate:
+    if not promt:
+        st.write("42")
+    else:
+        st.write(generate_text(promt))