Update main.py

main.py

@@ -24,6 +24,10 @@ def load_models():
                                      model=BertForSequenceClassification.from_pretrained("usmiva/bert-deepfake-bg", num_labels=2),
                                      tokenizer=AutoTokenizer.from_pretrained("usmiva/bert-deepfake-bg"))
 
+    st.session_state.emotions = pipeline(task="text-classification",
+                                     model=BertForSequenceClassification.from_pretrained("TRACES/emotions", use_auth_token=os.environ['ACCESS_TOKEN'],  num_labels=2),
+                                     tokenizer=AutoTokenizer.from_pretrained("TRACES/emotions", use_auth_token=os.environ['ACCESS_TOKEN']))
+
 
 def load_content():
     with open('resource/page_content.json', encoding='utf8') as json_file:
@@ -44,7 +48,8 @@ if 'lang' not in st.session_state:
 if all([
     'bert_gpt_result' not in st.session_state,
     'untrue_detector_result' not in st.session_state,
-    'bert_disinfo_result' not in st.session_state
+    'bert_disinfo_result' not in st.session_state,
+    'emotions_result' not in st.session_state
     ]):
     st.session_state.bert_gpt_result = [{'label': '', 'score': 1}]
 
@@ -97,6 +102,7 @@ if st.session_state.agree:
             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_disinfo_result = st.session_state.bert_disinfo(user_input)
+            st.session_state.bert_disinfo_result = st.session_state.emotions(user_input)
 
             
 
@@ -127,6 +133,20 @@ if st.session_state.agree:
                        str(round(st.session_state.bert_disinfo_result[0]['score'] * 100, 2)) +
                        content['bert_no_2'][st.session_state.lang], icon="✅")
 
+        if st.session_state.emotions[0]['score'] < 0.97:
+            st.warning(content['emotions_label_1'][st.session_state.lang] + 
+                       str(st.session_state.emotions[0]['label']) + 
+                       content['emotions_label_2'][st.session_state.lang] 
+                       str(round(st.session_state.emotions[0]['score'] * 100, 2)) + 
+                       content['emotions_label_3'][st.session_state.lang] +
+                       content['emotions_label_4'][st.session_state.lang], icon = "⚠️")
+        else:
+            st.success(content['emotions_label_1'][st.session_state.lang] + 
+                       str(st.session_state.emotions[0]['label']) + 
+                       content['emotions_label_2'][st.session_state.lang] 
+                       str(round(st.session_state.emotions[0]['score'] * 100, 2)) + 
+                       content['emotions_label_3'][st.session_state.lang], icon="✅")
+
     
         st.info(content['disinformation_definition'][st.session_state.lang], icon="ℹ️")