update app to use indexes directly to get embeddings

app.py

@@ -11,11 +11,11 @@ import click
 
 
 def getRandID():
-    indx = random.randrange(0, 396503)
-    return indx_to_id_dict[indx], indx
+    indx = random.randrange(0, len(index_to_id_dict))
+    return index_to_id_dict[indx], indx
 
 
-def chooseImageIndex(indexType):
+def get_image_index(indexType):
     if indexType == "FlatIP(default)":
         return image_index_IP
     elif indexType == "FlatL2":
@@ -32,7 +32,7 @@ def chooseImageIndex(indexType):
         return image_index_LSH
 
 
-def chooseDNAIndex(indexType):
+def get_dna_index(indexType):
     if indexType == "FlatIP(default)":
         return dna_index_IP
     elif indexType == "FlatL2":
@@ -49,35 +49,36 @@ def chooseDNAIndex(indexType):
         return dna_index_LSH
 
 
-def searchEmbeddings(id, key_type, query_type, index_type):
-    # variable and index initialization
-    dim = 768
-    count = 0
-    num_neighbors = 10
-
-    index = faiss.IndexFlatIP(dim)
+def searchEmbeddings(id, key_type, query_type, index_type, num_results: int = 10):
+    image_index = get_image_index(index_type)
+    dna_index = get_dna_index(index_type)
 
     # get index
     if query_type == "Image":
-        index = chooseImageIndex(index_type)
+        query = image_index.reconstruct(id_to_index_dict[id])
     elif query_type == "DNA":
-        index = chooseDNAIndex(index_type)
+        query = dna_index.reconstruct(id_to_index_dict[id])
+    else:
+        raise ValueError(f"Invalid query type: {query_type}")
+    query = query.astype(np.float32)
+    query = np.expand_dims(query, axis=0)
 
     # search for query
     if key_type == "Image":
-        query = id_to_image_emb_dict[id]
+        index = image_index
     elif key_type == "DNA":
-        query = id_to_dna_emb_dict[id]
-    query = query.astype(np.float32)
-    D, I = index.search(query, num_neighbors)
+        index = dna_index
+    else:
+        raise ValueError(f"Invalid key type: {key_type}")
 
-    id_list = []
-    i = 1
+    _, I = index.search(query, num_results)
+
+    closest_ids = []
     for indx in I[0]:
-        id = indx_to_id_dict[indx]
-        id_list.append(id)
+        id = index_to_id_dict[indx]
+        closest_ids.append(id)
 
-    return id_list
+    return closest_ids
 
 
 with gr.Blocks() as demo:
@@ -102,14 +103,8 @@ with gr.Blocks() as demo:
     # 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:
-        indx_to_id_dict = pickle.load(f)
-
-    # initialize both possible dicts
-    with open("big_id_to_image_emb_dict.pickle", "rb") as f:
-        id_to_image_emb_dict = pickle.load(f)
-    # with open("big_id_to_dna_emb_dict.pickle", "rb") as f:
-    #     id_to_dna_emb_dict = pickle.load(f)
-    id_to_dna_emb_dict = None
+        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():
         with gr.Row():
@@ -124,12 +119,18 @@ with gr.Blocks() as demo:
         index_type = gr.Radio(
             choices=["FlatIP(default)", "FlatL2", "HNSWFlat", "IVFFlat", "LSH"], label="Index:", value="FlatIP(default)"
         )
+        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:")
         search_btn = gr.Button("Search")
         id_btn.click(fn=getRandID, inputs=[], outputs=[rand_id, rand_id_indx])
 
-    search_btn.click(fn=searchEmbeddings, inputs=[process_id, key_type, query_type, index_type], outputs=[process_id_list])
+    search_btn.click(
+        fn=searchEmbeddings,
+        inputs=[process_id, key_type, query_type, index_type, num_results],
+        outputs=[process_id_list],
+    )
 
 
 demo.launch()

big_id_to_image_emb_dict.pickle

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

big_indx_to_id_dict.pickle

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