Add app file

app.py

@@ -0,0 +1,121 @@
+import glob
+import gradio as gr
+import pandas as pd
+import faiss
+import clip
+import torch
+from datasets import load_dataset
+
+title = r"""
+<h1 align="center" id="space-title"> 🔍 Search Similar Text/Image in the Dataset</h1>
+"""
+
+description = r"""
+
+In this demo, we use [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.
+<br>
+The content will be updated to include image search once LAION is available.
+
+The code is based on [clip-retrieval](https://github.com/rom1504/clip-retrieval) and [autofaiss](https://github.com/criteo/autofaiss)
+
+"""
+
+# From local file
+# INDEX_DIR = "dataset/diffusiondb/text_index_folder"
+# IND = faiss.read_index(f"{INDEX_DIR}/text.index")
+# TEXT_LIST = pd.concat(
+#     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
+
+# Download index file
+hf_hub_download(
+    repo_id="Eun02/diffusiondb_faiss_text_index",
+    filename="text.index",
+    repo_type="dataset",
+    local_dir="./",
+)
+
+# Download text file
+snapshot_download(
+    repo_id="Eun02/diffusiondb_faiss_text_index", 
+    allow_patterns="*.parquet",
+    repo_type="dataset",
+    local_dir="./",
+)
+
+# Load index and text data
+#root_path = "dataset/diffusiondb/text_index_folder"
+root_path = "."
+IND = faiss.read_index(f"{root_path}/text.index")
+TEXT_LIST = pd.concat(
+    pd.read_parquet(file) for file in sorted(glob.glob(f"{root_path}/metadata/*.parquet"))
+)['caption'].tolist()
+
+# Load CLIP model
+device = "cpu"
+CLIP_MODEL, _ = clip.load("ViT-B/32", device=device)
+
+@torch.inference_mode
+def get_emb(text, device="cpu"):
+    text_tokens = clip.tokenize([text], truncate=True)
+    text_features = CLIP_MODEL.encode_text(text_tokens.to(device))
+    text_features /= text_features.norm(dim=-1, keepdim=True)
+    text_embeddings = text_features.cpu().numpy().astype('float32')
+    return text_embeddings
+
+@torch.inference_mode
+def search_text(dataset, top_k, show_score, query_text, device):
+
+    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 = scores[0], retrieved_texts[0]
+
+    result_str = ""
+    for score, ind in zip(scores, retrieved_texts):
+        item_str = TEXT_LIST[ind].strip()
+        if item_str == "":
+            continue
+        result_str += f"{item_str}"
+        if show_score:
+            result_str += f", {score:0.2f}"
+        result_str += "\n"
+    
+    file_name = query_text.replace(" ", "_")
+    if show_score:
+        file_name += "_score"
+    output_path = f"./{file_name}.txt"
+    with open(output_path, "w") as f:
+        f.writelines(result_str)
+    
+    return result_str, output_path
+
+
+with gr.Blocks() as demo:
+    gr.Markdown(title)
+    gr.Markdown(description)
+
+    with gr.Row():
+        dataset = gr.Dropdown(label="dataset", choices=["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")
+    btn = gr.Button()
+    with gr.Row():
+        result_text = gr.Textbox(label="retrieved text", interactive=False)
+        result_file = gr.File(label="output file")
+
+    btn.click(
+        fn=search_text,
+        inputs=[dataset, top_k, show_score, query_text],
+        outputs=[result_text, result_file],
+    )
+
+demo.launch()