исправлена первая модель

app.py

@@ -152,21 +152,30 @@ if page == "Классификация отзывов":
             vector = [vocab_to_int.get(word, unknown_token_id) for word in words]  # здесь unknown_token_id - это ID для "неизвестных" слов
             return np.array(vector, dtype=np.int64)  # Убедитесь, что тип данных int64
 
-    def classify_review_lstm(review):
-        # Векторизация отзыва
-        review_vector = text_to_vector(review)
-        # Преобразование в тензор PyTorch и добавление размерности пакета (batch)
-        review_tensor = torch.tensor(review_vector).unsqueeze(0)
-        
-        # Предсказание
+
+    def classify_review_lstm(review: str, SEQ_LEN: int, model: nn.Module, threshold: float = 0.5):
+        """Predict sentiment class for a review
+
+        Args:
+            review (str): Review text
+            SEQ_LEN (int): sequence length
+            model (nn.Module): trained model
+            threshold (float): threshold for class prediction
+
+        Returns:
+            str: Predicted sentiment ('positive' or 'negative')
+        """
+        inp = text_to_vector(review)
+        inp_tensor = torch.tensor(inp, dtype=torch.int64)
         start_time = time.time()
-        with torch.no_grad():
-            prediction, _ = model_lstm(review_tensor)
+        with torch.inference_mode():
+            pred, _ = model(inp_tensor.long().unsqueeze(0))
         end_time = time.time()
-
-        # Время предсказания
         prediction_time = end_time - start_time
-        return prediction, prediction_time
+        # Convert prediction to sentiment label
+        sentiment = 'positive' if pred.sigmoid().item() > threshold else 'negative'
+
+        return sentiment, prediction_time
     
     tokenizer_rubert = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny2")
     model_rubert = AutoModel.from_pretrained("cointegrated/rubert-tiny2")
@@ -199,9 +208,9 @@ if page == "Классификация отзывов":
             prediction_tf, pred_time_tf = classify_review_tf(user_review)
             st.write(f'Предсказанный класс TF-IDF: {prediction_tf}')
             st.write(f'Время предсказания TF-IDF: {pred_time_tf:.4f} секунд')
-            prediction_lstm, pred_time_lstm = classify_review_lstm(user_review)
-            st.write(f'Предсказанный класс LSTM: {prediction_tf}')
-            st.write(f'Время предсказания LSTM: {pred_time_tf:.4f} секунд')
+            prediction_lstm, pred_time_lstm = classify_review_lstm(user_review, SEQ_LEN=SEQ_LEN, model=model_lstm)
+            st.write(f'Предсказанный класс LSTM: {prediction_lstm}')
+            st.write(f'Время предсказания LSTM: {pred_time_lstm:.4f} секунд')
             prediction_rubert, pred_time_rubert = make_prediction(user_review)
             prediction_ru = 'negative' if prediction_rubert == 0 else 'positive'
             st.write(f'Предсказанный класс RuBERT: {prediction_ru}')