update code to download dataset files from separate repo

.gitignore

@@ -0,0 +1,7 @@
+.build
+.data
+.singularity
+slurm/
+.env
+*.sif
+__pycache__/

app.py

@@ -1,13 +1,10 @@
-import gradio as gr
-import torch
-import numpy as np
-import h5py
-import faiss
-from PIL import Image
-import io
 import pickle
 import random
-import click
+
+import gradio as gr
+import numpy as np
+
+from data import load_indexes_local, load_indexes_hf, load_index_pickle
 
 
 def getRandID():
@@ -16,37 +13,17 @@ def getRandID():
 
 
 def get_image_index(indexType):
-    if indexType == "FlatIP(default)":
-        return image_index_IP
-    elif indexType == "FlatL2":
-        raise NotImplementedError
-        return image_index_L2
-    elif indexType == "HNSWFlat":
-        raise NotImplementedError
-        return image_index_HNSW
-    elif indexType == "IVFFlat":
-        raise NotImplementedError
-        return image_index_IVF
-    elif indexType == "LSH":
-        raise NotImplementedError
-        return image_index_LSH
+    try:
+        return image_indexes[indexType]
+    except KeyError:
+        raise KeyError(f"Tried to load an image index that is not supported: {indexType}")
 
 
 def get_dna_index(indexType):
-    if indexType == "FlatIP(default)":
-        return dna_index_IP
-    elif indexType == "FlatL2":
-        raise NotImplementedError
-        return dna_index_L2
-    elif indexType == "HNSWFlat":
-        raise NotImplementedError
-        return dna_index_HNSW
-    elif indexType == "IVFFlat":
-        raise NotImplementedError
-        return dna_index_IVF
-    elif indexType == "LSH":
-        raise NotImplementedError
-        return dna_index_LSH
+    try:
+        return dna_indexes[indexType]
+    except KeyError:
+        raise KeyError(f"Tried to load a DNA index that is not supported: {indexType}")
 
 
 def searchEmbeddings(id, key_type, query_type, index_type, num_results: int = 10):
@@ -86,24 +63,13 @@ with gr.Blocks() as demo:
     # for hf: change all file paths, indx_to_id_dict as well
 
     # load indexes
-    image_index_IP = faiss.read_index("bioscan_5m_image_IndexFlatIP.index")
-    # image_index_L2 = faiss.read_index("big_image_index_FlatL2.index")
-    # image_index_HNSW = faiss.read_index("big_image_index_HNSWFlat.index")
-    # image_index_IVF = faiss.read_index("big_image_index_IVFFlat.index")
-    # image_index_LSH = faiss.read_index("big_image_index_LSH.index")
-
-    dna_index_IP = faiss.read_index("bioscan_5m_dna_IndexFlatIP.index")
-    # dna_index_L2 = faiss.read_index("big_dna_index_FlatL2.index")
-    # dna_index_HNSW = faiss.read_index("big_dna_index_HNSWFlat.index")
-    # dna_index_IVF = faiss.read_index("big_dna_index_IVFFlat.index")
-    # dna_index_LSH = faiss.read_index("big_dna_index_LSH.index")
-
-    # with open("dataset_processid_list.pickle", "rb") as f:
-    #     dataset_processid_list = pickle.load(f)
-    # with open("processid_to_index.pickle", "rb") as f:
-    #     processid_to_index = pickle.load(f)
-    with open("big_indx_to_id_dict.pickle", "rb") as f:
-        index_to_id_dict = pickle.load(f)
+    image_indexes = load_indexes_hf(
+        {"FlatIP(default)": "bioscan_5m_image_IndexFlatIP.index"}, repo_name="bioscan-ml/bioscan-clibd"
+    )
+    dna_indexes = load_indexes_hf(
+        {"FlatIP(default)": "bioscan_5m_dna_IndexFlatIP.index"}, repo_name="bioscan-ml/bioscan-clibd"
+    )
+    index_to_id_dict = load_index_pickle("big_indx_to_id_dict.pickle", repo_name="bioscan-ml/bioscan-clibd")
     id_to_index_dict = {v: k for k, v in index_to_id_dict.items()}
 
     with gr.Column():
@@ -113,8 +79,8 @@ with gr.Blocks() as demo:
                 rand_id_indx = gr.Textbox(label="Index:")
                 id_btn = gr.Button("Get Random ID")
             with gr.Column():
-                key_type = gr.Radio(choices=["Image", "DNA"], label="Search From:", value="Image")
-                query_type = gr.Radio(choices=["Image", "DNA"], label="Search To:", value="Image")
+                query_type = gr.Radio(choices=["Image", "DNA"], label="Query:", value="Image")
+                key_type = gr.Radio(choices=["Image", "DNA"], label="Key:", value="Image")
 
         index_type = gr.Radio(
             choices=["FlatIP(default)", "FlatL2", "HNSWFlat", "IVFFlat", "LSH"], label="Index:", value="FlatIP(default)"
@@ -122,7 +88,7 @@ with gr.Blocks() as demo:
         num_results = gr.Number(label="Number of Results:", value=10, precision=0)
 
         process_id = gr.Textbox(label="ID:", info="Enter a sample ID to search for")
-        process_id_list = gr.Textbox(label="Closest 10 matches:")
+        process_id_list = gr.Textbox(label="Closest matches:")
         search_btn = gr.Button("Search")
         id_btn.click(fn=getRandID, inputs=[], outputs=[rand_id, rand_id_indx])
 

big_indx_to_id_dict.pickle

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:ee0a9044e054f640b704247a2fa2e74219180b78ded6ba07f551bfc222657fc5
-size 885457

bioscan_5m_dna_IndexFlatIP.index

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:83fe6599724756652689b76ef942ffaca2f8d5863ff3dd7fe7ac655199e0968d
-size 136009773

bioscan_5m_image_IndexFlatIP.index

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:01e5d74fd5194551e2b8e43aba8e41153efeb29589fa82a7839791d2e057c21d
-size 136009773

data.py

@@ -0,0 +1,34 @@
+import os
+import pickle
+from typing import Any
+
+import faiss
+from huggingface_hub import hf_hub_download
+
+
+def load_indexes_local(index_files: dict[str, str], *, data_folder: str, **kw) -> dict[str, Any]:
+    indexes = {}
+    for index_type, index_file in index_files.items():
+        indexes[index_type] = faiss.read_index(os.path.join(data_folder, index_file))
+
+    return indexes
+
+
+def load_indexes_hf(index_files: dict[str, str], *, repo_name: str, **kw) -> dict[str, Any]:
+    indexes = {}
+    for index_type, index_file in index_files.items():
+        indexes[index_type] = faiss.read_index(
+            hf_hub_download(repo_id=repo_name, filename=index_file, repo_type="dataset")
+        )
+
+    return indexes
+
+
+def load_index_pickle(index_file: str, repo_name: str) -> Any:
+    index_to_id_dict_file = hf_hub_download(
+        repo_id=repo_name,
+        filename=index_file,
+        repo_type="dataset",
+    )
+    with open(index_to_id_dict_file, "rb") as f:
+        return pickle.load(f)

prepare_index.py

@@ -1,3 +1,4 @@
+import pickle
 from pathlib import Path
 
 import click
@@ -9,55 +10,33 @@ ALL_INDEX_TYPES = ["IndexFlatIP", "IndexFlatL2", "IndexIVFFlat", "IndexHNSWFlat"
 EMBEDDING_SIZE = 768
 
 
-def process(input: Path, output: Path, key_type: str, index_type: str):
+def process(embedding_data, output: Path, key_type: str, index_type: str):
     # load embeddings
-    all_keys = h5py.File(input / "extracted_features_of_all_keys.hdf5", "r", libver="latest")[
-        f"encoded_{key_type}_feature"
-    ][:]
-    seen_test = h5py.File(input / "extracted_features_of_seen_test.hdf5", "r", libver="latest")[
-        f"encoded_{key_type}_feature"
-    ][:]
-    unseen_test = h5py.File(input / "extracted_features_of_unseen_test.hdf5", "r", libver="latest")[
-        f"encoded_{key_type}_feature"
-    ][:]
-    seen_val = h5py.File(input / "extracted_features_of_seen_val.hdf5", "r", libver="latest")[
-        f"encoded_{key_type}_feature"
-    ][:]
-    unseen_val = h5py.File(input / "extracted_features_of_unseen_val.hdf5", "r", libver="latest")[
-        f"encoded_{key_type}_feature"
-    ][:]
+    embeddings = embedding_data[f"encoded_{key_type}_feature"][:]
 
     # FlatIP and FlatL2
     if index_type == "IndexFlatIP":
-        test_index = faiss.IndexFlatIP(EMBEDDING_SIZE)
+        test_index = faiss.IndexFlatIP(embeddings.shape[-1])
     elif index_type == "IndexFlatL2":
-        test_index = faiss.IndexFlatL2(EMBEDDING_SIZE)
+        test_index = faiss.IndexFlatL2(embeddings.shape[-1])
     elif index_type == "IndexIVFFlat":
         # IVFFlat
-        quantizer = faiss.IndexFlatIP(EMBEDDING_SIZE)
-        test_index = faiss.IndexIVFFlat(quantizer, EMBEDDING_SIZE, 128)
-        test_index.train(all_keys)
-        test_index.train(seen_test)
-        test_index.train(unseen_test)
-        test_index.train(seen_val)
-        test_index.train(unseen_val)
+        quantizer = faiss.IndexFlatIP(embeddings.shape[-1])
+        test_index = faiss.IndexIVFFlat(quantizer, embeddings.shape[-1], 128)
+        test_index.train(embeddings)
     elif index_type == "IndexHNSWFlat":
         # HNSW
         # 16: connections for each vertex. efSearch: depth of search during search. efConstruction: depth of search during build
-        test_index = faiss.IndexHNSWFlat(EMBEDDING_SIZE, 16)
+        test_index = faiss.IndexHNSWFlat(embeddings.shape[-1])
         test_index.hnsw.efSearch = 32
         test_index.hnsw.efConstruction = 64
     elif index_type == "IndexLSH":
         # LSH
-        test_index = faiss.IndexLSH(EMBEDDING_SIZE, EMBEDDING_SIZE * 2)
+        test_index = faiss.IndexLSH(embeddings.shape[-1], embeddings.shape[-1] * 2)
     else:
         raise ValueError(f"Index type {index_type} is not supported")
 
-    test_index.add(all_keys)
-    test_index.add(seen_test)
-    test_index.add(unseen_test)
-    test_index.add(seen_val)
-    test_index.add(unseen_val)
+    test_index.add(embeddings)
 
     faiss.write_index(test_index, str(output / f"bioscan_5m_{key_type}_{index_type}.index"))
     print("Saved index to", output / f"bioscan_5m_{key_type}_{index_type}.index")
@@ -96,9 +75,15 @@ def main(input, output, key_type, index_type):
     else:
         index_types = [index_type]
 
+    embedding_data = h5py.File(input / "extracted_features_for_all_5m_data.hdf5", "r", libver="latest")
     for key_type in key_types:
         for index_type in index_types:
-            process(input, output, key_type, index_type)
+            process(embedding_data, output, key_type, index_type)
+
+    sample_ids = [raw_id.decode("utf-8") for raw_id in embedding_data["file_name_list"][:]]
+    index_to_id = {index: id for index, id in enumerate(sample_ids)}
+    with open(output / "big_indx_to_id_dict.pickle", "wb") as f:
+        pickle.dump(index_to_id, f)
 
 
 if __name__ == "__main__":