TRACES/traces-tool

1import json2import os3 4import streamlit as st5import pickle6 7from transformers import AutoTokenizer, BertForSequenceClassification, pipeline8from sklearn.feature_extraction.text import TfidfVectorizer9 10 11def load_models():12    st.session_state.loaded = True13 14    # with open('models/tfidf_vectorizer_untrue_inform_detection_tfidf_bg_0.96_F1_score_3Y_N_Q1_082023.pkl', 'rb') as f:15    #     st.session_state.tfidf_vectorizer_untrue_inf = pickle.load(f)16 17    # with open('models/SVM_model_untrue_inform_detection_tfidf_bg_0.96_F1_score_3Y_N_Q1_082023.pkl', 'rb') as f:18    #     st.session_state.untrue_detector = pickle.load(f)19 20    st.session_state.bert_disinfo = pipeline(task="text-classification",21                                     model=BertForSequenceClassification.from_pretrained("usmiva/bert-desinform-bg", num_labels=2),22                                     tokenizer=AutoTokenizer.from_pretrained("usmiva/bert-desinform-bg"))23    st.session_state.bert_gpt = pipeline(task="text-classification",24                                     model=BertForSequenceClassification.from_pretrained("usmiva/bert-deepfake-bg", num_labels=2),25                                     tokenizer=AutoTokenizer.from_pretrained("usmiva/bert-deepfake-bg"))26 27    st.session_state.emotions = pipeline(task="text-classification",28                                     model=BertForSequenceClassification.from_pretrained("TRACES/emotions", use_auth_token=os.environ['ACCESS_TOKEN2'],  num_labels=11),29                                     tokenizer=AutoTokenizer.from_pretrained("usmiva/bert-web-bg"))30 31 32 33def load_content():34    with open('resource/page_content.json', encoding='utf8') as json_file:35        return json.load(json_file)36 37 38def switch_lang(lang):39    if 'lang' in st.session_state:40        if lang == 'bg':41            st.session_state.lang = 'bg'42        else:43            st.session_state.lang = 'en'44 45 46if 'lang' not in st.session_state:47    st.session_state.lang = 'bg'48 49if all([50    'bert_gpt_result' not in st.session_state,51    # 'untrue_detector_result' not in st.session_state,52    'bert_disinfo_result' not in st.session_state,53    'emotions_result' not in st.session_state54    ]):55    st.session_state.bert_gpt_result = [{'label': '', 'score': 1}]56 57    # st.session_state.untrue_detector_result = ''58    # st.session_state.untrue_detector_probability = 159    60    st.session_state.bert_disinfo_result = [{'label': '', 'score': 1}]61        62    st.session_state.emotions_result = [{'label': '', 'score': 1}]    63 64content = load_content()65if 'loaded' not in st.session_state:66    load_models()67 68#######################################################################################################################69 70st.title(content['title'][st.session_state.lang])71 72col1, col2, col3 = st.columns([1, 1, 10])73with col1:74    st.button(75        label='EN',76        key='en',77        on_click=switch_lang,78        args=['en']79    )80with col2:81    st.button(82        label='BG',83        key='bg',84        on_click=switch_lang,85        args=['bg']86    )87 88if 'agree' not in st.session_state:89    st.session_state.agree = False90 91if st.session_state.agree:92    tab_tool, tab_terms = st.tabs([content['tab_tool'][st.session_state.lang], content['tab_terms'][st.session_state.lang]])93 94    with tab_tool:95        user_input = st.text_area(content['textbox_title'][st.session_state.lang],96                                  content['text_placeholder'][st.session_state.lang]).strip('\n')97    98        if st.button(content['analyze_button'][st.session_state.lang]):99            st.session_state.bert_gpt_result = st.session_state.bert_gpt(user_input)100            101            # user_tfidf_untrue_inf = st.session_state.tfidf_vectorizer_untrue_inf.transform([user_input])102            # st.session_state.untrue_detector_result = st.session_state.untrue_detector.predict(user_tfidf_untrue_inf)[0]103            # st.session_state.untrue_detector_probability = st.session_state.untrue_detector.predict_proba(user_tfidf_untrue_inf)[0]104            # st.session_state.untrue_detector_probability = max(st.session_state.untrue_detector_probability[0], st.session_state.untrue_detector_probability[1]) 105 106            st.session_state.bert_disinfo_result = st.session_state.bert_disinfo(user_input)107            108            st.session_state.emotions_result = st.session_state.emotions(user_input)109 110            111 112        if st.session_state.bert_gpt_result[0]['label'] == 'LABEL_1':113            st.warning(content['bert_gpt'][st.session_state.lang] +114                       str(round(st.session_state.bert_gpt_result[0]['score'] * 100, 2)) +115                       content['bert_gpt_prob'][st.session_state.lang], icon = "⚠️")116        else:117            st.success(content['bert_human'][st.session_state.lang] +118                       str(round(st.session_state.bert_gpt_result[0]['score'] * 100, 2)) +119                       content['bert_human_prob'][st.session_state.lang], icon="✅")120        121        # if st.session_state.untrue_detector_result == 0:122        #     st.warning(content['untrue_getect_yes'][st.session_state.lang] +123        #                str(round(st.session_state.untrue_detector_probability * 100, 2)) +124        #                content['untrue_yes_proba'][st.session_state.lang], icon="⚠️")125        # else:126        #     st.success(content['untrue_getect_no'][st.session_state.lang] +127        #                str(round(st.session_state.untrue_detector_probability * 100, 2)) +128        #                content['untrue_no_proba'][st.session_state.lang], icon="✅")129    130        if st.session_state.bert_disinfo_result[0]['label'] == 'LABEL_1':131            st.warning(content['bert_yes_1'][st.session_state.lang] +132                       str(round(st.session_state.bert_disinfo_result[0]['score'] * 100, 2)) +133                       content['bert_yes_2'][st.session_state.lang], icon = "⚠️")134        else:135            st.success(content['bert_no_1'][st.session_state.lang] +136                       str(round(st.session_state.bert_disinfo_result[0]['score'] * 100, 2)) +137                       content['bert_no_2'][st.session_state.lang], icon="✅")138 139        if st.session_state.emotions_result[0]['score'] < 0.97:140            st.warning(content['emotions_label_1'][st.session_state.lang] + 141                       str(st.session_state.emotions_result[0]['label']) + 142                       content['emotions_label_2'][st.session_state.lang] +143                       str(round(st.session_state.emotions_result[0]['score'] * 100, 2)) + 144                       content['emotions_label_3'][st.session_state.lang] +145                       content['emotions_label_4'][st.session_state.lang], icon = "⚠️")146        else:147            st.info(content['emotions_label_1'][st.session_state.lang] + 148                       str(st.session_state.emotions_result[0]['label']) + 149                       content['emotions_label_2'][st.session_state.lang] +150                       str(round(st.session_state.emotions_result[0]['score'] * 100, 2)) + 151                       content['emotions_label_3'][st.session_state.lang]+152                       content['emotions_label_5'][st.session_state.lang])153 154    155        st.info(content['disinformation_definition'][st.session_state.lang], icon="ℹ️")156 157    with tab_terms:158        st.write(content['disclaimer'][st.session_state.lang])159 160else:161    st.write(content['disclaimer_title'][st.session_state.lang])162    st.write(content['disclaimer'][st.session_state.lang])163    if st.button(content['disclaimer_agree_text'][st.session_state.lang]):164        st.session_state.agree = True165        st.experimental_rerun()166