wiki download option

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 
-from utils import load_hf_dataset, get_model_and_tokenizer, batch_embed
+from utils import load_hf_dataset, get_model_and_tokenizer, batch_embed, download_wikipedia
 
 
 # TODO: add instructor models
@@ -48,7 +48,11 @@ def download(
 ):
     if progress is not None:
         progress(0.5, "Loading dataset...")
-    ds = load_hf_dataset(ds_name, ds_config, ds_split)
+    
+    if ds_name == "wikipedia":
+        ds = download_wikipedia(ds_name, ds_config)
+    else:
+        ds = load_hf_dataset(ds_name, ds_config, ds_split)
 
     return f"Downloaded! It has {len(ds)} docs."