Update app.py

app.py

@@ -8,7 +8,7 @@ embedding_path = "abokbot/wikipedia-embedding"
 
 st.header("Wikipedia Search Engine app")
 
-st_model_load = st.text('Loading wikipedia embedding...')
+st_model_load = st.text('Loading encoders, embeddings and dataset (takes about 5min)')
 
 @st.cache_resource
 def load_embedding():
@@ -19,7 +19,6 @@ def load_embedding():
     return wikipedia_embedding
 
 wikipedia_embedding = load_embedding()
-st.success('Embedding loaded!')
 st_model_load.text("")
 
 @st.cache_resource
@@ -29,6 +28,7 @@ def load_encoders():
     bi_encoder.max_seq_length = 256     #Truncate long passages to 256 tokens
     top_k = 32  
     cross_encoder = CrossEncoder('cross-encoder/ms-marco-TinyBERT-L-2-v2')
+    print("Encoders loaded!")
     return bi_encoder, cross_encoder
 
 bi_encoder, cross_encoder = load_encoders()
@@ -39,10 +39,11 @@ st_model_load.text("")
 def load_wikipedia_dataset():
     print("Loading wikipedia dataset...")
     dataset = load_dataset("abokbot/wikipedia-first-paragraph")["train"]
+    print("Dataset loaded!")
     return dataset
     
 dataset = load_wikipedia_dataset()
-st.success('Datset loaded!')
+st.success('Loading done')
 st_model_load.text("")
     
 if 'text' not in st.session_state:
@@ -57,7 +58,10 @@ st_text_area = st.text_area(
 def search():
     st.session_state.text = st_text_area
     query = st_text_area
+    print("Input question:", query)
+    
     ##### Sematic Search #####
+    print("Semantic Search")
     # Encode the query using the bi-encoder and find potentially relevant passages
     top_k = 32
     question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
@@ -66,6 +70,7 @@ def search():
 
     ##### Re-Ranking #####
     # Now, score all retrieved passages with the cross_encoder
+    print("Re-Ranking")
     cross_inp = [[query, dataset[hit['corpus_id']]["text"]] for hit in hits]
     cross_scores = cross_encoder.predict(cross_inp)
 
@@ -99,7 +104,7 @@ if 'results' not in st.session_state:
 if len(st.session_state.results) > 0:
     with st.container():
         st.subheader("Search results")
-        for result in st.session_state.questions:
+        for result in st.session_state.results:
             for k,v in result.items():
                 st.markdown("score: " + results["score"])
                 st.markdown("title: " + results["title"])