Update pages/3_Earnings_Semantic_Search_🔎_.py

pages/3_Earnings_Semantic_Search_🔎_.py

@@ -5,6 +5,10 @@ st.set_page_config(page_title="Earnings Semantic Search", page_icon="🔎")
 st.sidebar.header("Semantic Search")
 st.markdown("## Earnings Semantic Search with SBert")
 
+def gen_sentiment(text):
+    '''Generate sentiment of given text'''
+    return sent_pipe(text)[0]['label']
+
 search_input = st.text_input(
         label='Enter Your Search Query',value= "What challenges did the business face?", key='search')
         
@@ -13,8 +17,9 @@ top_k = st.sidebar.slider("Number of Top Hits Generated",min_value=1,max_value=5
 window_size = st.sidebar.slider("Number of Sentences Generated in Search Response",min_value=1,max_value=5,value=3)
 
 if search_input:
+    
 
-    if any(st.session_state["sen_df"]) is not None and st.session_state["earnings_passages"] is not None:
+    if "sen_df" in st.session_state and "earnings_passages" in st.session_state:
     
         ## Save to a dataframe for ease of visualization
         sen_df = st.session_state['sen_df']
@@ -44,34 +49,39 @@ if search_input:
         score='cross-score'
         df = pd.DataFrame([(hit[score],passages[hit['corpus_id']]) for hit in hits[0:int(top_k)]],columns=['Score','Text'])
         df['Score'] = round(df['Score'],2)
+        df['Sentiment'] = df.Text.apply(gen_sentiment)
         
-        def gen_annotated_text(para,df=sen_df):
-            tag_list = []
-            print(df)
-            for i in sent_tokenize(para):
-                print(df.loc[df['text']==i, 'label'])
-                label = df.loc[df['text']==i, 'label'].values[0]
-                if label == 'Negative':
-                    tag_list.append((i,label,'#faa'))
-                elif label == 'Positive':
-                    tag_list.append((i,label,'#afa'))
+        def gen_annotated_text(df):
+            '''Generate annotated text'''
+            
+            tag_list=[]
+            for row in df.itertuples():
+                label = row[3]
+                text = row[2]
+                if label == 'Positive':
+                    tag_list.append((text,label,'#8fce00'))
+                elif label == 'Negative':
+                    tag_list.append((text,label,'#f44336'))
                 else:
-                    tag_list.append((i,label,'#fea'))
+                    tag_list.append((text,label,'#fff2cc'))
+                
             return tag_list  
         
-        text_to_annotate = [gen_annotated_text(para,sen_df) for para in df.Text.tolist()]
+        text_annotations = gen_annotated_text(df)
+
+        first, second = text_annotations[0], text_annotations[1]
         
-        first,second = text_to_annotate[0],text_to_annotate[-1]
         
-        with st.container():
-            annotated_text(*first)
+        with st.expander(label='Best Search Query Result', expanded=True):
+            annotated_text(first)
             
-        with st.container():
-            annotated_text(*second)
+        with st.expander(label='Alternative Search Query Result'):
+            annotated_text(second)
             
     else:
         
         st.write('Please ensure you have entered the YouTube URL or uploaded the Earnings Call file')
+        
 else:
 
     st.write('Please ensure you have entered the YouTube URL or uploaded the Earnings Call file')