Update app.py

app.py

@@ -5,13 +5,13 @@ import gradio as gr
 import polars as pl
 from datasets import load_dataset
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
-from model2vec import StaticModel
+from sentence_transformers import SentenceTransformer
 
 global df
 
-# Load a model from the HuggingFace hub (in this case the potion-base-8M model)
-model_name = "minishlab/potion-base-8M"
-model = StaticModel.from_pretrained(model_name)
+# Load the static embeddings model from HuggingFace hub
+model_name = "sentence-transformers/static-retrieval-mrl-en-v1"
+model = SentenceTransformer(model_name, device="cpu")
 
 
 def get_iframe(hub_repo_id):
@@ -58,7 +58,7 @@ def vectorize_dataset(hub_repo_id: str, split: str, column: str):
     gr.Info("Vectorizing dataset...")
     ds = load_dataset(hub_repo_id)
     df = ds[split].to_polars()
-    embeddings = model.encode(df[column].cast(str), max_length=512)
+    embeddings = model.encode(df[column].cast(str).to_list(), show_progress_bar=True)
     return embeddings
 
 
@@ -73,7 +73,7 @@ def run_query(hub_repo_id: str, query: str, split: str, column: str):
             query=f"""
             SELECT *
             FROM df
-            ORDER BY array_cosine_distance(embeddings, {vector.tolist()}::FLOAT[256])
+            ORDER BY array_cosine_distance(embeddings, {vector.tolist()}::FLOAT[1024])
             LIMIT 5
             """
         ).to_df()
@@ -165,4 +165,4 @@ with gr.Blocks() as demo:
         outputs=results_output,
     )
 
-demo.launch()
+demo.launch()