Update app.py

app.py

@@ -13,7 +13,7 @@ from datasets import load_dataset
 @st.cache_data
 def read_data(dataset_repo='dhmeltzer/ELI5_embedded'):
     """Read the data from huggingface."""
-    return load_dataset(dataset_repo)
+    return load_dataset(dataset_repo)['train']
 
 @st.cache_data
 def load_faiss_index(path_to_faiss="./faiss_index.pickle"):