fix: avoid global usage

app.py

@@ -1,3 +1,5 @@
+from functools import lru_cache
+
 import duckdb
 import gradio as gr
 import polars as pl
@@ -5,7 +7,6 @@ from datasets import load_dataset
 from gradio_huggingfacehub_search import HuggingfaceHubSearch
 from model2vec import StaticModel
 
-global ds
 global df
 
 # Load a model from the HuggingFace hub (in this case the potion-base-8M model)
@@ -28,14 +29,13 @@ def get_iframe(hub_repo_id):
     return iframe
 
 
-def load_dataset_from_hub(hub_repo_id):
-    gr.Info("Loading dataset...")
-    global ds
+def load_dataset_from_hub(hub_repo_id: str):
+    gr.Info(message="Loading dataset...")
     ds = load_dataset(hub_repo_id)
 
 
-def get_columns(split: str):
-    global ds
+def get_columns(hub_repo_id: str, split: str):
+    ds = load_dataset(hub_repo_id)
     ds_split = ds[split]
     return gr.Dropdown(
         choices=ds_split.column_names,
@@ -45,33 +45,35 @@ def get_columns(split: str):
     )
 
 
-def get_splits():
-    global ds
+def get_splits(hub_repo_id: str):
+    ds = load_dataset(hub_repo_id)
     splits = list(ds.keys())
     return gr.Dropdown(
         choices=splits, value=splits[0], label="Select a split", visible=True
     )
 
 
-def vectorize_dataset(split: str, column: str):
+@lru_cache
+def vectorize_dataset(hub_repo_id: str, split: str, column: str):
     gr.Info("Vectorizing dataset...")
-    global df
-    global ds
+    ds = load_dataset(hub_repo_id)
     df = ds[split].to_polars()
     embeddings = model.encode(df[column].cast(str), max_length=512)
-    df = df.with_columns(pl.Series(embeddings).alias(f"{column}_embeddings"))
+    return embeddings
 
 
-def run_query(query: 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)
+    df = ds[split].to_polars()
+    df = df.with_columns(pl.Series(embeddings).alias("embeddings"))
     try:
-        global df
-
         vector = model.encode(query)
         df_results = duckdb.sql(
             query=f"""
             SELECT *
             FROM df
-            ORDER BY array_cosine_distance({column}_embeddings, {vector.tolist()}::FLOAT[256])
+            ORDER BY array_cosine_distance(embeddings, {vector.tolist()}::FLOAT[256])
             LIMIT 5
             """
         ).to_df()
@@ -134,6 +136,7 @@ with gr.Blocks() as demo:
         query_input = gr.Textbox(label="Query", visible=False)
 
     btn_run = gr.Button("Search", visible=False)
+
     results_output = gr.Dataframe(label="Results", visible=False)
 
     search_in.submit(get_iframe, inputs=search_in, outputs=search_out).then(
@@ -143,23 +146,23 @@ with gr.Blocks() as demo:
     ).then(
         fn=hide_components,
         outputs=[split_dropdown, column_dropdown, query_input, btn_run, results_output],
-    ).then(fn=get_splits, outputs=split_dropdown).then(
-        fn=get_columns, inputs=split_dropdown, outputs=column_dropdown
+    ).then(fn=get_splits, inputs=search_in, outputs=split_dropdown).then(
+        fn=get_columns, inputs=[search_in, split_dropdown], outputs=column_dropdown
     )
 
     split_dropdown.change(
-        fn=get_columns, inputs=split_dropdown, outputs=column_dropdown
+        fn=get_columns, inputs=[search_in, split_dropdown], outputs=column_dropdown
     )
 
     column_dropdown.change(
         fn=partial_hide_components,
         outputs=[query_input, btn_run, results_output],
-    ).then(fn=vectorize_dataset, inputs=[split_dropdown, column_dropdown]).then(
-        fn=show_components, outputs=[query_input, btn_run]
-    )
+    ).then(fn=show_components, outputs=[query_input, btn_run])
 
     btn_run.click(
-        fn=run_query, inputs=[query_input, column_dropdown], outputs=results_output
+        fn=run_query,
+        inputs=[search_in, query_input, split_dropdown, column_dropdown],
+        outputs=results_output,
     )
 
 demo.launch()