Update app.py

app.py

@@ -19,10 +19,9 @@ def get_image_base64(path):
 
 @st.cache_resource  # Кэширование модели для ускорения работы
 def load_model():
-    model = fasttext.load_model('./model_fasttext.bin')
-    return model
+    return ModelWrapper()
 
-@st.cache_resource
+"""@st.cache_resource
 class ModelWrapper(object):
     MODELS_DIR: str = "./new_models/"
     MODEL_NAME: str = "model"
@@ -54,10 +53,9 @@ class ModelWrapper(object):
         )  # output is logits for huggingfcae transformers
         predicted = torch.nn.functional.softmax(outputs.logits, dim=1)
         predicted_id = torch.argmax(predicted, dim=1).numpy()[0]
-        return self.id2label[predicted_id]
+        return self.id2label[predicted_id]"""
 
-model_fasttext = load_model()
-model_BERT=ModelWrapper()
+model_wrapper= load_model()
 
 bin_str = get_image_base64("./билли.png")
 page_bg_img = '''
@@ -112,41 +110,22 @@ model_type = st.radio(
 )
 
 def highlight_obscene_words(text, model_type):
-    if model_type=="fasttext":
-        label,_=model_fasttext.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
-            )
+    label=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:#47916B;'>{}|приемлемо</span>".format(text), 
-                unsafe_allow_html=True
-            )
-
+            "<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 = re.split(r'[.\n!?]+', user_input)
         result = [part for part in result if part.strip() != ""]
         if result!=[]:
             for text in result: