Update Models/bert_strim.py

Models/bert_strim.py

@@ -30,6 +30,7 @@ def predict_sentiment(text):
     with torch.no_grad():
         output = model(input_ids, attention_mask)
         prediction = torch.round(output).cpu().numpy()[0][0]
+        return prediction
     if prediction == 1:
         return "Позитивный отзыв 😀"
     else:
@@ -39,12 +40,13 @@ def bert_model_page():
     st.title("Классификация отзывов")
     user_input = st.text_area("Введите отзыв:")
     if st.button("Классифицировать"):
-        if user_input:
-            start_time = time.time()
-            prediction = predict_sentiment(user_input)
-            end_time = time.time()
-            execution_time = end_time - start_time
-            st.write(prediction)
-            st.write(f'Время предсказания: {execution_time:.4f} секунд')
+        start_time = time.time()
+        prediction = predict_sentiment(user_input)
+        end_time = time.time()
+        execution_time = end_time - start_time
+        if prediction > 0.5:
+            st.write("Отзыв положительный 🌟")
         else:
-            st.write("Пожалуйста, введите отзыв для классификации.")
+            st.write("Отзыв отрицательный 😞")
+        st.write(f'Время предсказания: {execution_time:.4f} секунд')
+