Update app.py

app.py

@@ -4,37 +4,38 @@ from sentence_transformers import SentenceTransformer
 import duckdb
 from huggingface_hub import get_token
 
-model = SentenceTransformer("Snowflake/snowflake-arctic-embed-m-v1.5")
-
-def similarity_search(
-    query: str, 
-    k: int = 5, 
-    dataset_name: str = "smol-blueprint-project/hf-blogs-text-embeddings", 
-    embedding_column: str = "embedding",
-):
-    # Use same model as used for indexing
-    query_vector = model.encode(query)
-    embedding_dim = model.get_sentence_embedding_dimension()
-    
-    sql = f"""
-        SELECT 
-            title,
-            author,
-            date,
-            local,
-            tags,
-            URL,
-            chunk,
-            array_cosine_distance(
-                {embedding_column}::float[{embedding_dim}], 
-                {query_vector.tolist()}::float[{embedding_dim}]
-            ) as distance
-        FROM 'hf://datasets/{dataset_name}/**/*.parquet'
-        ORDER BY distance
-        LIMIT {k}
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.models import StaticEmbedding
+import duckdb
+
+# Initialize a StaticEmbedding module
+static_embedding = StaticEmbedding.from_model2vec("minishlab/potion-base-8M")
+model = SentenceTransformer(modules=[static_embedding])
+
+dataset_name = "smol-blueprint/fineweb-bbc-news-text-embeddings"
+embedding_column = "embedding"
+
+duckdb.sql(
+    query=f"""
+    INSTALL vss;
+    LOAD vss;
+    CREATE TABLE embeddings AS 
+    SELECT *, {embedding_column}::float[{model.get_sentence_embedding_dimension()}] as embedding_float 
+    FROM 'hf://datasets/{dataset_name}/**/*.parquet';
+    CREATE INDEX my_hnsw_index ON embeddings USING HNSW (embedding_float) WITH (metric = 'cosine');
+"""
+)
+
+def similarity_search(query: str, k: int = 5):
+    embedding = model.encode(query).tolist()
+    return duckdb.sql(
+        query=f"""
+        SELECT url, chunk, array_cosine_distance(embedding_float, {embedding}::FLOAT[{model.get_sentence_embedding_dimension()}]) as distance 
+        FROM embeddings 
+        ORDER BY distance 
+        LIMIT {k};
     """
-    
-    return duckdb.sql(sql).to_df()
+    ).to_df()
 
 with gr.Blocks() as demo:
     gr.Markdown("""# Vector Search Hub Datasets
@@ -43,7 +44,7 @@ with gr.Blocks() as demo:
     query = gr.Textbox(label="Query")
     k = gr.Slider(1, 10, value=5, label="Number of results")
     btn = gr.Button("Search")
-    results = gr.Dataframe(headers=["title", "url", "content", "distance"])
+    results = gr.Dataframe(headers=["url", "chunk", "distance"])
     btn.click(fn=similarity_search, inputs=[query, k], outputs=[results])