Update pages/3_Earnings_Semantic_Search_🔎_.py

pages/3_Earnings_Semantic_Search_🔎_.py

@@ -9,10 +9,24 @@ def gen_sentiment(text):
     '''Generate sentiment of given text'''
     return sent_pipe(text)[0]['label']
 
+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((text,label,'#000000'))
+                    
+                return tag_list
+
 bi_enc_dict = {'mpnet-base-v2':"all-mpnet-base-v2",
-               'e5-base':'intfloat/e5-base',
               'instructor-base': 'hkunlp/instructor-base',
-              'mpnet-base-dot-v1':'multi-qa-mpnet-base-dot-v1',
               'setfit-finance': 'nickmuchi/setfit-finetuned-financial-text-classification'}
 
 search_input = st.text_input(
@@ -20,9 +34,9 @@ search_input = st.text_input(
         
 sbert_model_name = st.sidebar.selectbox("Embedding Model", options=list(bi_enc_dict.keys()), key='sbox')
         
-top_k = 2
-
-window_size = st.sidebar.slider("Number of Sentences Generated in Search Response",min_value=1,max_value=7,value=3)
+chunk_size = st.sidebar.slider("Number of Words per Chunk of Text",min_value=100,max_value=250,value=200)
+overlap_size = st.sidebar.slider("Number of Overlap Words in Search Response",min_value=30,max_value=100,value=50)
+chain_type = st.sidebar.radio("Langchain Chain Type",options = ['Normal','Refined'])
 
 try:
 
@@ -32,69 +46,36 @@ try:
         
             ## Save to a dataframe for ease of visualization
             sen_df = st.session_state['sen_df']
-                
-            passages = chunk_long_text(st.session_state['earnings_passages'],150,window_size=window_size)
-            
+
+            title = st.session_state['title']
+                            
             with st.spinner(
-                text=f"Loading {bi_enc_dict[sbert_model_name]} encoder model..."
+                text=f"Loading {bi_enc_dict[sbert_model_name]} embedding model and Generating Response..."
             ):
-                sbert = load_sbert(bi_enc_dict[sbert_model_name])
-                
-            
-            ##### Sematic Search #####
-            # Encode the query using the bi-encoder and find potentially relevant passages
-            # corpus_embeddings = sbert.encode(passages, convert_to_tensor=True, show_progress_bar=True)
-            # question_embedding = sbert.encode(search_input, convert_to_tensor=True)
-            # question_embedding = question_embedding.cpu()
-            # hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
-            # hits = hits[0]  # Get the hits for the first query
-            
-            # ##### Re-Ranking #####
-            # # Now, score all retrieved passages with the cross_encoder
-            # cross_inp = [[search_input, passages[hit['corpus_id']]] for hit in hits]
-            # cross_scores = cross_encoder.predict(cross_inp)
-            
-            # # Sort results by the cross-encoder scores
-            # for idx in range(len(cross_scores)):
-            #     hits[idx]['cross-score'] = cross_scores[idx]
-            
-            # # Output of top-3 hits from re-ranker
-            # hits = sorted(hits, key=lambda x: x['cross-score'], reverse=True)
-            embedding_model = bi_enc_dict[sbert_model_name]
-            
-            hits = embed_text(search_input,passages,embedding_model)
 
-            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(df):
-                '''Generate annotated text'''
+                result = embed_text(search_input,st.session_state['earnings_passages'],title,
+                                    bi_enc_dict[sbert_model_name],
+                                   emb_tokenizer,chain_type=chain_type)
+
+            references = [doc.page_content for doc in result['input_documents']]
+
+            answer = result['output_text']
+
+            sentiment_label = gen_sentiment(answer)
                 
-                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((text,label,'#000000'))
-                    
-                return tag_list  
+            ##### Sematic Search #####
             
-            text_annotations = gen_annotated_text(df)
-    
-            first, second = text_annotations[0], text_annotations[1]
+            df = pd.DataFrame([(num,res,lab) for num, res, lab in zip(1,answer,sentiment_label)],columns=['Index','Text','Sentiment'])
+              
             
+            text_annotations = gen_annotated_text(df)[0]            
             
-            with st.expander(label='Best Search Query Result', expanded=True):
-                annotated_text(first)
+            with st.expander(label='Query Result', expanded=True):
+                annotated_text(text_annotations)
                 
-            with st.expander(label='Alternative Search Query Result'):
-                annotated_text(second)
+            with st.expander(label='References from Corpus used to Generate Result'):
+                for ref in references:
+                    st.write(ref)
                 
         else: