Update app.py

app.py

@@ -72,29 +72,46 @@ user_input = st.text_area('',height=200)
 
 model_type = st.radio(
     "Выберите модель",
-    ("fasttext", "ru-BERT")
+    ("fasttext", "ru-BERT","FRIDA")
 )
 
 def highlight_obscene_words(text, model_type):
-    label, prob=model_wrapper(text.lower(),model_type)
-    if label=='__label__positive':
-        st.markdown(
-            "<span style='background:#47916B;'>{}:приемлемо</span>".format(text), 
-            unsafe_allow_html=True
-        )
-    else:
-        st.markdown(
-            "<span style='background:#ffcccc;'>{}:токсично</span>".format(text), 
-            unsafe_allow_html=True
-        )
+    if model_type=="FRIDA":
+        result=model_wrapper(text)
+        for item in result:
+            if item.label=="non-toxic":
+                st.markdown(
+                    "<span style='background:#47916B;'>{}:приемлемо</span>".format(text), 
+                    unsafe_allow_html=True
+                )
+            else:
+                st.markdown(
+                    "<span style='background:#ffcccc;'>{}:токсично</span>".format(text), 
+                    unsafe_allow_html=True
+                )
+    else:    
+        label, prob=model_wrapper(text.lower(),model_type)
+        if label=='__label__positive':
+            st.markdown(
+                "<span style='background:#47916B;'>{}:приемлемо</span>".format(text), 
+                unsafe_allow_html=True
+            )
+        else:
+            st.markdown(
+                "<span style='background:#ffcccc;'>{}:токсично</span>".format(text), 
+                unsafe_allow_html=True
+            )
     
 if st.button("Проверить текст"):
     if user_input.strip():
         st.subheader("Результат:")
         result = re.split(r'[.\n!?]+', user_input)
         result = [part for part in result if part.strip() != ""]
-        if result!=[]:
-            for text in result:
-                highlight_obscene_words(text,model_type)
+        if model_type=="FRIDA":
+            highlight_obscene_words(text,model_type)
+        else:
+            if result!=[]:
+                for text in result:
+                    highlight_obscene_words(text,model_type)
     else:
         st.warning("Пожалуйста, введите текст для проверки")