Update app.py

app.py

@@ -39,7 +39,6 @@ desc2opt = {v: k for k, v in opt2desc.items()}
 optimization_options = list(opt2desc.values())
 
 
-
 def download_and_tokenize(
     ds_name,
     ds_config,
@@ -51,7 +50,6 @@ def download_and_tokenize(
     num2embed,
     progress=gr.Progress(track_tqdm=True),
 ):
-    
     num_samples = download_dataset(ds_name, ds_config, ds_split, num2skip, num2embed)
 
     opt_level = desc2opt[opt_desc]
@@ -69,8 +67,6 @@ def download_and_tokenize(
     )
 
     return f"Downloaded! It has {len(num_samples)} docs."
-    
-    
 
 
 def embed(
@@ -85,7 +81,6 @@ def embed(
     num2embed,
     progress=gr.Progress(track_tqdm=True),
 ):
-
     ds = load_tokenized_dataset(ds_name, ds_config, ds_split)
 
     opt_level = desc2opt[opt_desc]
@@ -154,7 +149,9 @@ with gr.Blocks(title="Bulk embeddings") as demo:
             value="wikipedia",
         )
         ds_config = gr.Textbox(
-            lines=1, label="Dataset config (leave blank to use default)", value="20220301.en"
+            lines=1,
+            label="Dataset config (leave blank to use default)",
+            value="20220301.en",
         )
 
         column_name = gr.Textbox(lines=1, label="Enter column to embed", value="text")
@@ -208,18 +205,20 @@ with gr.Blocks(title="Bulk embeddings") as demo:
         )
 
     with gr.Row():
-
         download_btn = gr.Button(value="Download and tokenize dataset!")
         embed_btn = gr.Button(value="Embed texts!")
 
         last = gr.Textbox(value="")
 
     download_btn.click(
-        fn=download,
+        fn=download_and_tokenize,
         inputs=[
             ds_name,
             ds_config,
+            column_name,
             ds_split,
+            model_choice,
+            opt_desc,
             num2skip,
             num2embed,
         ],
@@ -244,4 +243,4 @@ with gr.Blocks(title="Bulk embeddings") as demo:
 
 
 if __name__ == "__main__":
-    demo.queue(concurrency_count=20).launch(show_error=True, debug=True)
+    demo.queue(concurrency_count=20).launch(show_error=True, debug=True)