Update pages/3_Earnings_Semantic_Search_🔎_.py

pages/3_Earnings_Semantic_Search_🔎_.py

@@ -8,13 +8,17 @@ st.markdown("## Earnings Semantic Search with SBert")
 def gen_sentiment(text):
     '''Generate sentiment of given text'''
     return sent_pipe(text)[0]['label']
-    
-bi_enc_options = ["all-mpnet-base-v2","multi-qa-mpnet-base-dot-v1","intfloat/e5-base","all-MiniLM-L12-v2"]
+
+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(
         label='Enter Your Search Query',value= "What key challenges did the business face?", key='search')
         
-sbert_model_name = st.sidebar.selectbox("Embedding Model", options=bi_enc_options, key='sbox')
+sbert_model_name = st.sidebar.selectbox("Embedding Model", options=list(bi_enc_dict.keys()), key='sbox')
         
 top_k = 2
 
@@ -32,31 +36,34 @@ try:
             passages = chunk_long_text(st.session_state['earnings_passages'],150,window_size=window_size)
             
             with st.spinner(
-                text=f"Loading {sbert_model_name} encoder..."
+                text=f"Loading {bi_enc_dict[sbert_model_name]} encoder model..."
             ):
-                sbert = load_sbert(sbert_model_name)
+                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
+            # 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)
+            # ##### 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]
+            # # 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)
+            # # 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)