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

app.py

@@ -22,6 +22,8 @@ import json
 import gensim
 import torch.nn.functional as F
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
+from transformers import AutoModelForSequenceClassification
+
 
 
 st.title('10-я неделя DS. Классификация отзывов, определение токсичности и генерация текста')
@@ -230,43 +232,57 @@ if page == "Определение токсичности":
         return model
 
     # Загрузка обученной модели
-    clf = load_model('toxic/logistic_regression_model_toxic.pkl')  # Укажите путь к файлу модели
+    clf_c = load_model('toxic/logistic_regression_model_toxic.pkl')  # Укажите путь к файлу модели
 
     # Загрузка токенизатора и модели BERT
-    tokenizer = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny-toxicity")
-    model = AutoModel.from_pretrained("cointegrated/rubert-tiny-toxicity")
+    tokenizer_c = AutoTokenizer.from_pretrained("cointegrated/rubert-tiny-toxicity")
+    model_c = AutoModel.from_pretrained("cointegrated/rubert-tiny-toxicity")
 
     # Функция для предсказания токсичности сообщения
     def predict_toxicity(text):
-        encoded = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        encoded = tokenizer_c(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
         with torch.no_grad():
-            outputs = model(**encoded)
+            outputs = model_c(**encoded)
         features = outputs.last_hidden_state[:, 0, :].numpy()
-        prediction = clf.predict(features)
+        prediction = clf_c.predict_proba(features)
         return prediction[0]
 
+    model_checkpoint = 'cointegrated/rubert-tiny-toxicity'
+    tokenizer_b = AutoTokenizer.from_pretrained(model_checkpoint)
+    model_b = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
+
+    def text2toxicity(text):
+        with torch.no_grad():
+            inputs = tokenizer_b(text, return_tensors='pt', truncation=True, padding=True)
+            proba = torch.sigmoid(model_b(**inputs).logits).cpu().numpy()
+            return proba[0][1]
+
     # Создание интерфейса Streamlit
     st.title("Оценка токсичности сообщения")
 
     # Текстовое поле для ввода сообщения
     user_input = st.text_area("Введите сообщение для оценки")
 
-    if st.button("Оценить"):
+    if st.button("Оценить токсичность сообщения кастомизированной моделью"):
         if user_input:
             # Оценка токсичности сообщения
-            prediction = predict_toxicity(user_input)
-            if prediction > 0.5:
-                st.write("Сообщение токсично")
-                st.write(prediction)
-            else:
-                st.write("Сообщение не токсично")
-                st.write(prediction)
+            prediction = predict_toxicity(user_input)[1]
+            st.write(f'Вероятность токсичности согласно кастомизированной модели: {prediction:.4f}')
         else:
             st.write("Пожалуйста, введите сообщение")
+    
+    if st.button('Определить токсичность базовой моделью'):
+        if user_input:
+            # Определение токсичности сообщения
+            proba_toxicity = text2toxicity(user_input)
+            st.write(f'Вероятность токсичности rubert-tiny-toxicity.pretrained: {proba_toxicity:.4f}')
+        else:
+            st.write('Пожалуйста, введите сообщение')
 
     
     
 if page == "Генерация текста":
+
     # Путь к вашим весам модели
     model_weights_path = 'gpt-2/model.pt'
 
@@ -288,7 +304,7 @@ if page == "Генерация текста":
         st.title("Генератор плохих отзывов больниц от ruGPT3")
 
         # Ввод текста от пользователя
-        user_prompt = st.text_area("Введите текст-промпт:", "Я была в этой клинике..")
+        user_prompt = st.text_area("Введите текст-промпт:", "Я была в этой клинике")
 
         # Виджеты для динамической регуляции параметров
         max_length = st.slider("Выберите max_length:", 10, 300, 100)
@@ -312,4 +328,8 @@ if page == "Генерация текста":
                 ).cpu().numpy()
                 generated_text = tokenizer.decode(out[0], skip_special_tokens=True)
                 st.subheader("Сгенерированный текст:")
-                st.write(generated_text)
+                st.write(generated_text)
+
+    if __name__ == "__main__":
+        main()
+