Add dataset name mapping

app.py

@@ -32,8 +32,15 @@ The code is based on [clip-retrieval](https://github.com/rom1504/clip-retrieval)
 # From huggingface dataset
 from huggingface_hub import hf_hub_download, snapshot_download
 
+DATASET_NAME = {
+    "danbooru22": "booru22_000-300",
+    "DiffusionDB": "diffusiondb",
+}
+
 def load_faiss_index(dataset):
     index_dir = "data/faiss_index"
+    dataset = DATASET_NAME[dataset]
+
     hf_hub_download(
         repo_id="Eun02/text_image_faiss_index",
         subfolder=dataset,
@@ -106,7 +113,7 @@ with gr.Blocks() as demo:
     gr.Markdown(description)
 
     with gr.Row():
-        dataset = gr.Dropdown(label="dataset", choices=["booru22", "DiffusionDB"], value="DiffusionDB")
+        dataset = gr.Dropdown(label="dataset", choices=["danbooru22", "DiffusionDB"], value="DiffusionDB")
         top_k = gr.Slider(label="top k", minimum=1, maximum=20, value=8)
         show_score = gr.Checkbox(label="Show score", value=True)
     query_text = gr.Textbox(label="query text")