Update app.py

app.py

@@ -1,4 +1,5 @@
 import streamlit as st
+from dataset import load_dataset
 from sentence_transformers import SentenceTransformer, CrossEncoder, util
 import torch
 from huggingface_hub import hf_hub_download
@@ -33,6 +34,16 @@ def load_encoders():
 bi_encoder, cross_encoder = load_encoders()
 st.success('Encoders loaded!')
 st_model_load.text("")
+
+@st.cache_resource
+def load_wikipedia_dataset():
+    print("Loading wikipedia dataset...")
+    dataset = load_dataset("abokbot/wikipedia-first-paragraph")
+    return dataset
+    
+dataset = load_wikipedia_dataset()
+st.success('Datset loaded!')
+st_model_load.text("")
     
 if 'text' not in st.session_state:
     st.session_state.text = ""
@@ -49,7 +60,7 @@ def search():
     ##### Sematic Search #####
     # Encode the query using the bi-encoder and find potentially relevant passages
     question_embedding = bi_encoder.encode(query, convert_to_tensor=True)
-    hits = util.semantic_search(question_embedding, corpus_embeddings, top_k=top_k)
+    hits = util.semantic_search(question_embedding, wikipedia_embedding, top_k=top_k)
     hits = hits[0]  # Get the hits for the first query
 
     ##### Re-Ranking #####