Update app.py

app.py

@@ -16,7 +16,7 @@ dataset = load_dataset("Yoxas/statistical_literacyv2")
 
 data = dataset["train"]
 
-# Convert the string embeddings to numerical arrays
+# Convert the string embeddings to numerical arrays and ensure they are 2D
 def convert_and_ensure_2d_embeddings(example):
     # Clean the embedding string
     embedding_str = example['embedding']
@@ -31,7 +31,20 @@ def convert_and_ensure_2d_embeddings(example):
 # Apply the function to ensure embeddings are 2-dimensional and of type float32
 data = data.map(convert_and_ensure_2d_embeddings)
 
-data = data.add_faiss_index("embedding")
+# Flatten embeddings if they are nested 2D arrays
+def flatten_embeddings(example):
+    embedding = example['embedding']
+    if embedding.ndim == 2 and embedding.shape[0] == 1:
+        embedding = embedding.flatten()
+    return {'embedding': embedding}
+
+data = data.map(flatten_embeddings)
+
+# Ensure embeddings are in the correct shape for FAISS
+embeddings = np.array(data['embedding'].tolist(), dtype=np.float32)
+
+# Add FAISS index
+data.add_faiss_index_from_external_arrays(embeddings)
 
 model_id = "meta-llama/Meta-Llama-3-8B-Instruct"