Update app.py

app.py

@@ -23,40 +23,7 @@ def load_model():
     model = fasttext.load_model('./model_fasttext.bin')
     return model
 
-model = load_model()
-
-bin_str = get_image_base64("./билли.png")
-page_bg_img = '''
-<style>
-.stApp{
-background-image: linear-gradient(rgba(255, 255, 255, 0.7), 
-                        rgba(255, 255, 255, 0.7)),
-                        url("data:image/png;base64,%s");
-background-size: cover;
-background-position: center;
-background-repeat: no-repeat;
-background-attachment: fixed;
-}
-</style>
-''' % bin_str
-st.markdown(page_bg_img, unsafe_allow_html=True)
-
-logo_base64 = get_image_base64("./top_logo1.png")
-
-# Используем HTML для вставки логотипа в заголовок
-st.markdown(
-    f"""
-    <div style="display: flex; justify-content: center;">
-        <img src="data:image/png;base64,{logo_base64}" width="400">
-    </div>
-    """,
-    unsafe_allow_html=True
-)
-
-# Описание
-st.write("""<p style='text-align: center; font-size: 24px;'>Это приложение сделает твою речь менее токсичной. 
-И даже не придётся платить 300 bucks.</p>""", unsafe_allow_html=True)
-
+@st.cache_resource
 class ModelWrapper(object):
     MODELS_DIR: str = "./new_models/"
     MODEL_NAME: str = "model"
@@ -90,19 +57,40 @@ class ModelWrapper(object):
         predicted_id = torch.argmax(predicted, dim=1).numpy()[0]
         return self.id2label[predicted_id]
 
+model_fasttext = load_model()
+model_BERT=ModelWrapper()
 
-def highlight_obscene_words(text):
-    label,_=model.predict(text.lower())
-    if label[0]=='__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
-        )
+bin_str = get_image_base64("./билли.png")
+page_bg_img = '''
+<style>
+.stApp{
+background-image: linear-gradient(rgba(255, 255, 255, 0.7), 
+                        rgba(255, 255, 255, 0.7)),
+                        url("data:image/png;base64,%s");
+background-size: cover;
+background-position: center;
+background-repeat: no-repeat;
+background-attachment: fixed;
+}
+</style>
+''' % bin_str
+st.markdown(page_bg_img, unsafe_allow_html=True)
+
+logo_base64 = get_image_base64("./top_logo1.png")
+
+# Используем HTML для вставки логотипа в заголовок
+st.markdown(
+    f"""
+    <div style="display: flex; justify-content: center;">
+        <img src="data:image/png;base64,{logo_base64}" width="400">
+    </div>
+    """,
+    unsafe_allow_html=True
+)
+
+# Описание
+st.write("""<p style='text-align: center; font-size: 24px;'>Это приложение сделает твою речь менее токсичной. 
+И даже не придётся платить 300 bucks.</p>""", unsafe_allow_html=True)
 
 # Боковая панель
 with st.sidebar:
@@ -119,6 +107,43 @@ st.write("""<p style='text-align: center;'>Введите текст ниже,
 
 user_input = st.text_area('',height=200)
 
+model_type = st.sidebar.selectbox(
+    "Выберите модель",
+    ("fasttext", "ru-BERT", "FRIDA")
+)
+
+def highlight_obscene_words(text, model_type):
+    if model_type=="fasttext":
+        label,_=model.predict(text.lower())
+        if label[0]=='__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
+            )
+    elif model_type=="ru-BERT":
+        label=model_BERT(text.lower())
+        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
+            )
+    else:
+        #ЗАГЛУШКА
+        st.markdown(
+                "<span style='background:#47916B;'>{}|приемлемо</span>".format(text), 
+                unsafe_allow_html=True
+            )
+
 if st.button("Проверить текст"):
     if user_input.strip():
         st.subheader("Результат:")