Revert "update code to download dataset files from separate repo"

This reverts commit 07356cd41d62a22354471068093aa528e5724930.

.gitignore

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

app.py

@@ -1,10 +1,13 @@
-import pickle
-import random
-
 import gradio as gr
+import torch
 import numpy as np
-
-from data import load_indexes_local, load_indexes_hf, load_index_pickle
+import h5py
+import faiss
+from PIL import Image
+import io
+import pickle
+import random
+import click
 
 
 def getRandID():
@@ -13,17 +16,37 @@ def getRandID():
 
 
 def get_image_index(indexType):
-    try:
-        return image_indexes[indexType]
-    except KeyError:
-        raise KeyError(f"Tried to load an image index that is not supported: {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
 
 
 def get_dna_index(indexType):
-    try:
-        return dna_indexes[indexType]
-    except KeyError:
-        raise KeyError(f"Tried to load a DNA index that is not supported: {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
 
 
 def searchEmbeddings(id, key_type, query_type, index_type, num_results: int = 10):
@@ -63,13 +86,24 @@ with gr.Blocks() as demo:
     # for hf: change all file paths, indx_to_id_dict as well
 
     # load indexes
-    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")
+    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)
     id_to_index_dict = {v: k for k, v in index_to_id_dict.items()}
 
     with gr.Column():
@@ -79,8 +113,8 @@ with gr.Blocks() as demo:
                 rand_id_indx = gr.Textbox(label="Index:")
                 id_btn = gr.Button("Get Random ID")
             with gr.Column():
-                query_type = gr.Radio(choices=["Image", "DNA"], label="Query:", value="Image")
-                key_type = gr.Radio(choices=["Image", "DNA"], label="Key:", value="Image")
+                key_type = gr.Radio(choices=["Image", "DNA"], label="Search From:", value="Image")
+                query_type = gr.Radio(choices=["Image", "DNA"], label="Search To:", value="Image")
 
         index_type = gr.Radio(
             choices=["FlatIP(default)", "FlatL2", "HNSWFlat", "IVFFlat", "LSH"], label="Index:", value="FlatIP(default)"
@@ -88,7 +122,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 matches:")
+        process_id_list = gr.Textbox(label="Closest 10 matches:")
         search_btn = gr.Button("Search")
         id_btn.click(fn=getRandID, inputs=[], outputs=[rand_id, rand_id_indx])
 

big_indx_to_id_dict.pickle

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

bioscan_5m_dna_IndexFlatIP.index

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

bioscan_5m_image_IndexFlatIP.index

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

data.py

@@ -1,34 +0,0 @@
-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,4 +1,3 @@
-import pickle
 from pathlib import Path
 
 import click
@@ -10,33 +9,55 @@ ALL_INDEX_TYPES = ["IndexFlatIP", "IndexFlatL2", "IndexIVFFlat", "IndexHNSWFlat"
 EMBEDDING_SIZE = 768
 
 
-def process(embedding_data, output: Path, key_type: str, index_type: str):
+def process(input: Path, output: Path, key_type: str, index_type: str):
     # load embeddings
-    embeddings = embedding_data[f"encoded_{key_type}_feature"][:]
+    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"
+    ][:]
 
     # FlatIP and FlatL2
     if index_type == "IndexFlatIP":
-        test_index = faiss.IndexFlatIP(embeddings.shape[-1])
+        test_index = faiss.IndexFlatIP(EMBEDDING_SIZE)
     elif index_type == "IndexFlatL2":
-        test_index = faiss.IndexFlatL2(embeddings.shape[-1])
+        test_index = faiss.IndexFlatL2(EMBEDDING_SIZE)
     elif index_type == "IndexIVFFlat":
         # IVFFlat
-        quantizer = faiss.IndexFlatIP(embeddings.shape[-1])
-        test_index = faiss.IndexIVFFlat(quantizer, embeddings.shape[-1], 128)
-        test_index.train(embeddings)
+        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)
     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(embeddings.shape[-1])
+        test_index = faiss.IndexHNSWFlat(EMBEDDING_SIZE, 16)
         test_index.hnsw.efSearch = 32
         test_index.hnsw.efConstruction = 64
     elif index_type == "IndexLSH":
         # LSH
-        test_index = faiss.IndexLSH(embeddings.shape[-1], embeddings.shape[-1] * 2)
+        test_index = faiss.IndexLSH(EMBEDDING_SIZE, EMBEDDING_SIZE * 2)
     else:
         raise ValueError(f"Index type {index_type} is not supported")
 
-    test_index.add(embeddings)
+    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)
 
     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")
@@ -75,15 +96,9 @@ 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(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)
+            process(input, output, key_type, index_type)
 
 
 if __name__ == "__main__":