Update app.py

app.py

@@ -1,6 +1,5 @@
 import gradio as gr
 from datasets import load_dataset
-
 import os
 import spaces
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer, BitsAndBytesConfig
@@ -22,7 +21,10 @@ print(f"Example embedding shape: {np.array(example_embedding).shape}")
 
 # Ensure embeddings are 2-dimensional
 def ensure_2d_embeddings(embeddings):
-    return [np.atleast_2d(embedding) for embedding in embeddings]
+    embeddings = np.array(embeddings)
+    if embeddings.ndim == 1:
+        embeddings = embeddings.reshape(1, -1)
+    return embeddings
 
 # Apply the function to ensure embeddings are 2-dimensional
 data = data.map(lambda example: {'embedding': ensure_2d_embeddings(example['embedding'])})
@@ -57,7 +59,7 @@ def search(query: str, k: int = 3):
     """a function that embeds a new query and returns the most probable results"""
     embedded_query = ST.encode(query)  # embed new query
     scores, retrieved_examples = data.get_nearest_examples(  # retrieve results
-        "Abstract_Embeddings", embedded_query,  # compare our new embedded query with the dataset embeddings
+        "embedding", embedded_query,  # compare our new embedded query with the dataset embeddings
         k=k  # get only top k results
     )
     return scores, retrieved_examples