Update app.py

app.py

@@ -18,7 +18,6 @@ def main():
     import eli5
     import shap
     from custom_models import HF_DistilBertBasedModelAppDocs, HF_BertBasedModelAppDocs
-    import docx
 
     # Initialize Spacy
     nlp = spacy.load("en_core_web_sm")
@@ -93,24 +92,8 @@ def main():
         else:
             model = HF_DistilBertBasedModelAppDocs.from_pretrained("ferdmartin/HF_DistilBertBasedModelAppDocs").to(device)
         return model
-
-    def app_model(option, model, text):
-        with st.spinner("Wait for the magic 🪄🔮"):
-            # Use model
-            if option in ("Naive Bayes", "Logistic Regression"):
-                prediction, predict_proba = nb_lr(model, text)
-                st.session_state["sklearn"] = True
-            else:
-                prediction, predict_proba = torch_pred(tokenizer, model, text)
-                st.session_state["torch"] = True
     
-            # Store the result in session state
-            st.session_state["color_pred"] = "blue" if prediction == 0 else "red"
-            prediction = pred_str(prediction)
-            st.session_state["prediction"] = prediction
-            st.session_state["predict_proba"] = predict_proba
-            st.session_state["text"] = text
-            
+
     # Streamlit app:
 
     models_available = {"Logistic Regression":"models/baseline_model_lr2.joblib", 
@@ -121,7 +104,7 @@ def main():
 
     st.set_page_config(page_title="AI/Human GradAppDocs", page_icon="🤖", layout="wide")
     st.title("Academic Application Document Classifier")
-    st.header("Is it human-made 📝 or Generated with AI 🦾 ?  ")
+    st.header("Is it human-made 📝 or Generated with AI 🤖 ?  ")
     
     # Check the model to use
     def restore_prediction_state():
@@ -148,18 +131,32 @@ def main():
             </style>
             """
     st.markdown(hide_st_style, unsafe_allow_html=True)
-    col1, col2 = st.columns(2)
+    
     # Use model
-    with col1:
-        if st.button("Let's check this text!"):
-            if text.strip() == "":
-                st.error("Please enter some text")
-            else:
-                app_model(option, model, text)
+    if st.button("Let's check this text!"):
+        if text.strip() == "":
+            st.error("Please enter some text")
+        else:
+            with st.spinner("Wait for the magic 🪄🔮"):
+                # Use model
+                if option in ("Naive Bayes", "Logistic Regression"):
+                    prediction, predict_proba = nb_lr(model, text)
+                    st.session_state["sklearn"] = True
+                else:
+                    prediction, predict_proba = torch_pred(tokenizer, model, text)
+                    st.session_state["torch"] = True
+
+            # Store the result in session state
+            st.session_state["color_pred"] = "blue" if prediction == 0 else "red"
+            prediction = pred_str(prediction)
+            st.session_state["prediction"] = prediction
+            st.session_state["predict_proba"] = predict_proba
+            st.session_state["text"] = text
+            
             # Print result
-            #st.markdown(f"I think this text is: **:{st.session_state['color_pred']}[{st.session_state['prediction']}]** (Prediction probability: {st.session_state['predict_proba'] * 100}%)")
-   
-    if "prediction" in st.session_state:
+            st.markdown(f"I think this text is: **:{st.session_state['color_pred']}[{st.session_state['prediction']}]** (Prediction probability: {st.session_state['predict_proba'] * 100}%)")
+
+    elif "prediction" in st.session_state:
         # Display the stored result if available        
         st.markdown(f"I think this text is: **:{st.session_state['color_pred']}[{st.session_state['prediction']}]** (Prediction probability: {st.session_state['predict_proba'] * 100}%)")