Update main.py

main.py

@@ -53,7 +53,7 @@ if all([
     st.session_state.gpt_detector_probability = [1, 0]
 
     st.session_state.untrue_detector_result = ''
-    st.session_state.untrue_detector_probability = [1, 0]
+    st.session_state.untrue_detector_probability = 1
     
     st.session_state.bert_result = [{'label': '', 'score': 1}]
 
@@ -99,6 +99,7 @@ if st.session_state.agree:
             user_tfidf_untrue_inf = st.session_state.tfidf_vectorizer_untrue_inf.transform([user_input])
             st.session_state.untrue_detector_result = st.session_state.untrue_detector.predict(user_tfidf_untrue_inf)[0]
             st.session_state.untrue_detector_probability = st.session_state.untrue_detector.predict_proba(user_tfidf_untrue_inf)[0]
+            st.session_state.untrue_detector_probability = max(st.session_state.untrue_detector_probability[0], st.session_state.untrue_detector_probability[1]) 
 
             st.session_state.bert_result = st.session_state.bert(user_input)
     
@@ -113,11 +114,11 @@ if st.session_state.agree:
 
         if st.session_state.untrue_detector_result == 1:
             st.warning(content['untrue_getect_yes'][st.session_state.lang] +
-                       str(round(st.session_state.untrue_detector_probability[1] * 100, 2)) +
+                       str(round(st.session_state.untrue_detector_probability * 100, 2)) +
                        content['untrue_yes_proba'][st.session_state.lang], icon="⚠️")
         else:
             st.success(content['untrue_getect_no'][st.session_state.lang] +
-                       str(round(st.session_state.untrue_detector_probability[0] * 100, 2)) +
+                       str(round(st.session_state.untrue_detector_probability * 100, 2)) +
                        content['untrue_no_proba'][st.session_state.lang], icon="✅")
     
         if st.session_state.bert_result[0]['label'] == 'LABEL_1':