Update app.py

app.py

@@ -1,5 +1,6 @@
 import gradio as gr
 from datasets import load_dataset
+
 import os
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
@@ -15,21 +16,19 @@ dataset = load_dataset("Yoxas/statistical_literacyv2")
 
 data = dataset["train"]
 
-# Check the structure of embeddings
-example_embedding = data[0]['embedding']
-print(f"Example embedding shape: {np.array(example_embedding).shape}")
-
-# Ensure embeddings are 2-dimensional and of type string
-def ensure_2d_embeddings(embeddings):
-    embeddings = np.array(embeddings, dtype=np.float32)
+# Convert the string embeddings to numerical arrays
+def convert_and_ensure_2d_embeddings(example):
+    # Convert the string to a numpy array
+    embeddings = np.fromstring(example['embedding'].strip("[]"), sep=' ', dtype=np.float32)
+    # Ensure the embeddings are 2-dimensional
     if embeddings.ndim == 1:
         embeddings = embeddings.reshape(1, -1)
-    return embeddings
+    return {'embedding': embeddings}
 
 # Apply the function to ensure embeddings are 2-dimensional and of type float32
-data = data.map(lambda example: {'embedding': ensure_2d_embeddings(example['embedding'])})
+data = data.map(convert_and_ensure_2d_embeddings)
 
-data = data.add_faiss_index("Abstract")
+data = data.add_faiss_index("embedding")
 
 model_id = "meta-llama/Meta-Llama-3-8B-Instruct"