Change hugginface dataset

app.py

@@ -31,29 +31,29 @@ 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
 
-# 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()
+def load_faiss_index(dataset):
+    index_dir = "data/faiss_index"
+    hf_hub_download(
+        repo_id="Eun02/text_image_faiss_index",
+        subfolder=dataset,
+        filename="text.index",
+        repo_type="dataset",
+        local_dir=index_dir,
+    )
+
+    # Download text file
+    snapshot_download(
+        repo_id="Eun02/text_image_faiss_index", 
+        allow_patterns=f"{dataset}/*.parquet",
+        repo_type="dataset",
+        local_dir=index_dir,
+    )
+    index = faiss.read_index(f"{index_dir}/{dataset}/text.index")
+    text_list = pd.concat(
+        pd.read_parquet(file) for file in sorted(glob.glob(f"{index_dir}/{dataset}/metadata/*.parquet"))
+    )['caption'].tolist()
+
+    return index, text_list
 
 # Load CLIP model
 device = "cpu"
@@ -70,16 +70,18 @@ def get_emb(text, device="cpu"):
 @torch.inference_mode
 def search_text(dataset, top_k, show_score, query_text, device):
 
+    ind, text_list = load_faiss_index(dataset)
+
     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 = 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()
+        item_str = text_list[ind].strip()
         if item_str == "":
             continue
         result_str += f"{item_str}"