Update app.py

app.py

@@ -7,11 +7,14 @@ import tensorflow as tf
 from tensorflow.keras.preprocessing.text import Tokenizer
 from tensorflow.keras.preprocessing.sequence import pad_sequences
 import time
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
+model_checkpoint = 'cointegrated/rubert-tiny-toxicity'
+toxicity_tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
+toxicity_model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint)
 clf = load('my_model_filename.pkl')
 vectorizer = load('tfidf_vectorizer.pkl')
 scaler = load('scaler.joblib')
-
 tukinazor = load('tokenizer.pkl')
 rnn_model = load_model('path_to_my_model.h5')
 bert_model = BertForSequenceClassification.from_pretrained('my_bert_model')
@@ -19,6 +22,26 @@ tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased')
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 bert_model = bert_model.to(device)
 
+labels = ["не токсичный", "оскорбляющий", "непристойный", "угрожающий", "опасный"]
+def text2toxicity(text, aggregate=True):
+    """ Calculate toxicity of a text (if aggregate=True) or a vector of toxicity aspects (if aggregate=False)"""
+    with torch.no_grad():
+        inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True).to(model.device)
+        proba = torch.sigmoid(model(**inputs).logits).cpu().numpy()
+    
+    if isinstance(text, str):
+        proba = proba[0]
+    
+    if aggregate:
+        return 1 - proba.T[0] * (1 - proba.T[-1])
+    else:
+        # Добавленный блок кода
+        result = {}
+        for label, prob in zip(labels, proba):
+            result[label] = prob
+        return result
+
+        
 def predict_text(text):
     sequences = tukinazor.texts_to_sequences([text])
     padded_sequences = tf.keras.preprocessing.sequence.pad_sequences(sequences, maxlen=200, padding='post', truncating='post')
@@ -27,8 +50,19 @@ def predict_text(text):
     return predicted_class
 
 
-# Запуск приложения
 def main():
+    page_selection = st.sidebar.selectbox("Выберите страницу:", ["Классификация отзывов", "Анализ токсичности"])
+    
+    if page_selection == "Классификация отзывов":
+        page_reviews_classification()
+    elif page_selection == "Анализ токсичности":
+        page_toxicity_analysis()
+
+if __name__ == "__main__":
+    main()
+
+
+def page_reviews_classification():
     st.title("Модель классификации отзывов")
 
     # Ввод текста
@@ -77,5 +111,14 @@ def main():
             st.write(f"Прогнозируемый класс (BERT): {predictions.item() + 1}")
             st.write(f"Время вычисления: {elapsed_time:.2f} сек.")
 
-if __name__ == "__main__":
-    main()
+def page_toxicity_analysis():
+    # Код для анализа токсичности текста с использованием модели cointegrated/rubert-tiny-toxicity
+    user_input_toxicity = st.text_area("Введите текст для оценки токсичности:")
+
+    if st.button("Оценить токсичность"):
+        start_time = time.time()
+        probs = text2toxicity(user_input_toxicity, aggregate=False)
+        elapsed_time = time.time() - start_time
+        
+        for label, prob in probs.items():
+            st.write(f"Вероятность {label}: {prob:.4f}")