feat: use SOTA model

app.py

@@ -9,7 +9,7 @@ global ds
 global df
 
 # Load a model from the HuggingFace hub (in this case the potion-base-8M model)
-model_name = "minishlab/M2V_multilingual_output"
+model_name = "minishlab/potion-base-8M"
 model = StaticModel.from_pretrained(model_name)
 
 
@@ -53,7 +53,7 @@ def vectorize_dataset(split: str, column: str):
     global df
     global ds
     df = ds[split].to_polars()
-    embeddings = model.encode(df[column])
+    embeddings = model.encode(df[column], max_length=512 * 4)
     df = df.with_columns(pl.Series(embeddings).alias("embeddings"))
 
 
@@ -64,7 +64,7 @@ def run_query(query: str):
         query=f"""
         SELECT *
         FROM df
-        ORDER BY array_distance(embeddings, {vector.tolist()}::FLOAT[256])
+        ORDER BY array_cosine_distance(embeddings, {vector.tolist()}::FLOAT[256])
         LIMIT 5
         """
     ).to_df()
@@ -91,18 +91,16 @@ with gr.Blocks() as demo:
             )
     with gr.Row():
         search_out = gr.HTML(label="Search Results")
-    search_in.submit(get_iframe, inputs=search_in, outputs=search_out)
-
-    btn_load_dataset = gr.Button("Load Dataset")
 
     with gr.Row(variant="panel"):
         split_dropdown = gr.Dropdown(label="Select a split")
         column_dropdown = gr.Dropdown(label="Select a column")
     with gr.Row(variant="panel"):
         query_input = gr.Textbox(label="Query")
-
-    btn_load_dataset.click(
-        load_dataset_from_hub, inputs=search_in, show_progress=True
+    search_in.submit(get_iframe, inputs=search_in, outputs=search_out).then(
+        fn=load_dataset_from_hub,
+        inputs=search_in,
+        show_progress=True,
     ).then(fn=get_splits, outputs=split_dropdown).then(
         fn=get_columns, inputs=split_dropdown, outputs=column_dropdown
     )

demo.py

@@ -4,20 +4,20 @@ from datasets import load_dataset
 from model2vec import StaticModel
 
 # Load a model from the HuggingFace hub (in this case the potion-base-8M model)
-model_name = "minishlab/M2V_multilingual_output"
+model_name = "minishlab/potion-base-8M"
 model = StaticModel.from_pretrained(model_name)
 
 # Make embeddings
 ds = load_dataset("fka/awesome-chatgpt-prompts")
 df = ds["train"].to_polars()
-embeddings = model.encode(df["prompt"])
+embeddings = model.encode(df["act"])
 df = df.with_columns(pl.Series(embeddings).alias("embeddings"))
-vector = model.encode("vector search", show_progress_bar=True)
+vector = model.encode("An Ethereum Developer", show_progress_bar=True)
 duckdb.sql(
     query=f"""
     SELECT *
     FROM df
-    ORDER BY array_distance(embeddings, {vector.tolist()}::FLOAT[256])
-    LIMIT 1
+    ORDER BY array_cosine_distance(embeddings, {vector.tolist()}::FLOAT[256])
+    LIMIT 10
     """
 ).show()