Add prefix and Optimize load index

app.py

@@ -4,6 +4,7 @@ import pandas as pd
 import faiss
 import clip
 import torch
+from huggingface_hub import hf_hub_download, snapshot_download
 
 title = r"""
 <h1 align="center" id="space-title"> 🔍 Search Similar Text/Image in the Dataset</h1>
@@ -11,9 +12,8 @@ title = r"""
 
 description = r"""
 
-In this demo, we use subset of [danbooru22](https://huggingface.co/datasets/animelover/danbooru2022) or [DiffusionDB](https://huggingface.co/datasets/poloclub/diffusiondb) instead of [LAION](https://laion.ai/blog/laion-400-open-dataset/) because LAION is currently not available. 
-<br>
-This demo currently supports text search only.
+Find text or images similar to your query text with this demo. Currently, it supports text search only.<br>
+In this demo, we use a subset of [danbooru22](https://huggingface.co/datasets/animelover/danbooru2022) or [DiffusionDB](https://huggingface.co/datasets/poloclub/diffusiondb) instead of [LAION](https://laion.ai/blog/laion-400-open-dataset/) because LAION is currently not available. 
 <br>
 The content will be updated to include image search once LAION is available.
 
@@ -21,7 +21,6 @@ The code is based on [clip-retrieval](https://github.com/rom1504/clip-retrieval)
 
 """
 
-#In this demo, we use  because LAION is currently not available. 
 # From local file
 # INDEX_DIR = "dataset/diffusiondb/text_index_folder"
 # IND = faiss.read_index(f"{INDEX_DIR}/text.index")
@@ -29,13 +28,9 @@ The code is based on [clip-retrieval](https://github.com/rom1504/clip-retrieval)
 #     pd.read_parquet(file) for file in glob.glob(f"{INDEX_DIR}/metadata/*.parquet")
 # )['caption'].tolist()
 
-# From huggingface dataset
-from huggingface_hub import hf_hub_download, snapshot_download
-
-DATASET_NAME = {
-    "danbooru22": "booru22_000-300",
-    "DiffusionDB": "diffusiondb",
-}
+def download_all_index(dataset_dict):
+    for k in dataset_dict:
+        load_faiss_index(k)
 
 def load_faiss_index(dataset):
     index_dir = "data/faiss_index"
@@ -63,9 +58,13 @@ def load_faiss_index(dataset):
 
     return index, text_list
 
-# Load CLIP model
-device = "cpu"
-CLIP_MODEL, _ = clip.load("ViT-B/32", device=device)
+def change_index(dataset):
+    global INDEX, TEXT_LIST, PREV_DATASET
+    if PREV_DATASET != dataset:
+        gr.Info("Load index...")
+        INDEX, TEXT_LIST = load_faiss_index(dataset)
+        PREV_DATASET = dataset
+        gr.Info("Done!!")
 
 @torch.inference_mode
 def get_emb(text, device="cpu"):
@@ -76,55 +75,88 @@ def get_emb(text, device="cpu"):
     return text_embeddings
 
 @torch.inference_mode
-def search_text(dataset, top_k, show_score, query_text, device):
-
-    ind, text_list = load_faiss_index(dataset)
-
+def search_text(top_k, show_score, numbering_prefix, output_file, query_text):
     if query_text is None or query_text == "":
         raise gr.Error("Query text is missing")
     
     text_embeddings = get_emb(query_text, device)
-    scores, retrieved_texts = ind.search(text_embeddings, top_k)
+    scores, retrieved_texts = INDEX.search(text_embeddings, top_k)
     scores, retrieved_texts = scores[0], retrieved_texts[0]
 
-    result_str = ""
+    result_list = [] 
     for score, ind in zip(scores, retrieved_texts):
-        item_str = text_list[ind].strip()
+        item_str = TEXT_LIST[ind].strip()
         if item_str == "":
             continue
-        result_str += f"{item_str}"
+        if (item_str, score) not in result_list:
+            result_list.append((item_str, score))
+
+    # Postprocessing text
+    result_str = ""
+    for count, (item_str, score) in enumerate(result_list):
+        if numbering_prefix: 
+            item_str = f"######################  {count+1}  ######################\n {item_str}"
         if show_score:
-            result_str += f", {score:0.2f}"
-        result_str += "\n"
-    
+            item_str += f", {score:0.2f}"
+        result_str += f"{item_str}\n"
+            
     # file_name = query_text.replace(" ", "_")
     # if show_score:
     #     file_name += "_score"
-    file_name = "output"
-    output_path = f"./{file_name}.txt"
-    with open(output_path, "w") as f:
-        f.writelines(result_str)
+    output_path = None
+    if output_file:
+        file_name = "output"
+        output_path = f"./{file_name}.txt"
+        with open(output_path, "w") as f:
+            f.writelines(result_str)
     
     return result_str, output_path
 
 
+# Load CLIP model
+device = "cpu"
+CLIP_MODEL, _ = clip.load("ViT-B/32", device=device)
+
+# Dataset
+DATASET_NAME = {
+    "danbooru22": "booru22_000-300",
+    "DiffusionDB": "diffusiondb",
+}
+
+DEFAULT_DATASET = "danbooru22"
+PREV_DATASET = "danbooru22"
+
+# Download needed index
+download_all_index(DATASET_NAME)
+
+# Load default index
+INDEX, TEXT_LIST = load_faiss_index(DEFAULT_DATASET)
+
+
 with gr.Blocks() as demo:
     gr.Markdown(title)
     gr.Markdown(description)
 
     with gr.Row():
-        dataset = gr.Dropdown(label="dataset", choices=["danbooru22", "DiffusionDB"], value="DiffusionDB")
+        dataset = gr.Dropdown(label="dataset", choices=["danbooru22", "DiffusionDB"], value=DEFAULT_DATASET)
         top_k = gr.Slider(label="top k", minimum=1, maximum=20, value=8)
-        show_score = gr.Checkbox(label="Show score", value=True)
+        with gr.Column():
+            show_score = gr.Checkbox(label="Show score", value=False)
+            numbering_prefix = gr.Checkbox(label="Add numbering prefix", value=True)
+            output_file = gr.Checkbox(label="Return text file", value=True)
     query_text = gr.Textbox(label="query text")
     btn = gr.Button()
-    with gr.Row():
-        result_text = gr.Textbox(label="retrieved text", interactive=False)
-        result_file = gr.File(label="output file")
+    result_text = gr.Textbox(label="retrieved text", interactive=False)
+    result_file = gr.File(label="output file", visible=True)
+    
+    #dataset.change(change_index, dataset, None)
 
     btn.click(
+        fn=change_index,
+        inputs=[dataset],
+    ).success(
         fn=search_text,
-        inputs=[dataset, top_k, show_score, query_text],
+        inputs=[top_k, show_score, numbering_prefix, output_file, query_text],
         outputs=[result_text, result_file],
     )