Update app.py

app.py

@@ -16,7 +16,15 @@ ST = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
 dataset = load_dataset("Yoxas/statistical_literacyv2")
 
 data = dataset["train"]
-data = data.add_faiss_index("Abstract_Embeddings") # column name that has the embeddings of the dataset
+
+# Check the shape of the embeddings
+print(data["Abstract_Embeddings"].shape)
+
+# Reshape the embeddings if necessary
+if len(data["Abstract_Embeddings"].shape) == 1:
+    data["Abstract_Embeddings"] = data["Abstract_Embeddings"].reshape(-1, 1)
+
+data = data.add_faiss_index("Abstract_Embeddings")
 
 
 model_id = "meta-llama/Meta-Llama-3-8B-Instruct"