use num_proc for loading

app.py

@@ -1,3 +1,4 @@
+import os
 from functools import lru_cache
 
 import duckdb
@@ -14,7 +15,7 @@ model_name = "sentence-transformers/static-retrieval-mrl-en-v1"
 model = SentenceTransformer(
     model_name,
     device="cpu",
-    tokenizer_kwargs={"model_max_length": 512},
+    tokenizer_kwargs={"model_max_length": 512}, # arbitrary for this model, here to keep things fast
 )
 
 
@@ -35,7 +36,7 @@ def get_iframe(hub_repo_id):
 
 def load_dataset_from_hub(hub_repo_id: str):
     gr.Info(message="Loading dataset...")
-    ds = load_dataset(hub_repo_id)
+    ds = load_dataset(hub_repo_id, num_proc=os.cpu_count())
 
 
 def get_columns(hub_repo_id: str, split: str):
@@ -50,7 +51,7 @@ def get_columns(hub_repo_id: str, split: str):
 
 
 def get_splits(hub_repo_id: str):
-    ds = load_dataset(hub_repo_id)
+    ds = load_dataset(hub_repo_id, num_proc=os.cpu_count())
     splits = list(ds.keys())
     return gr.Dropdown(
         choices=splits, value=splits[0], label="Select a split", visible=True
@@ -60,7 +61,7 @@ def get_splits(hub_repo_id: str):
 @lru_cache
 def vectorize_dataset(hub_repo_id: str, split: str, column: str):
     gr.Info("Vectorizing dataset...")
-    ds = load_dataset(hub_repo_id)
+    ds = load_dataset(hub_repo_id, num_proc=os.cpu_count())
     df = ds[split].to_polars()
     embeddings = model.encode(df[column].cast(str).to_list(), show_progress_bar=True, batch_size=128)
     return embeddings
@@ -68,7 +69,7 @@ def vectorize_dataset(hub_repo_id: str, split: str, column: str):
 
 def run_query(hub_repo_id: str, query: str, split: str, column: str):
     embeddings = vectorize_dataset(hub_repo_id, split, column)
-    ds = load_dataset(hub_repo_id)
+    ds = load_dataset(hub_repo_id, num_proc=os.cpu_count())
     df = ds[split].to_polars()
     df = df.with_columns(pl.Series(embeddings).alias("embeddings"))
     try: