Update demo/modules/search.py

demo/modules/search.py

@@ -13,11 +13,25 @@ tmp_file_path = "/tmp/results.tsv"
 tmp_plot_path = "/tmp/histogram.svg"
 
 # Samples for input
-samples = [
-    ["Proteins with zinc bindings."],
-    ["Proteins locating at cell membrane."],
-    ["Protein that serves as an enzyme."]
-]
+samples = {
+    "sequence": [
+            ["MSATAEQNARNPKGKGGFARTVSQRKRKRLFLIGGALAVLAVAVGLMLTAFNQDIRFFRTPADLTEQDMTSGARFRLGGLVEEGSVSRTGSELRFTVTDTIKTVKVVFEGIPPDLFREGQGVVAEGRFGSDGLFRADNVLAKHDENYVPKDLADSLKKKGVWEGK"],
+            ["MITLDWEKANGLITTVVQDATTKQVLMVAYMNQESLAKTMATGETWFWSRSRKTLWHKGATSGNIQTVKTIAVDCDADTLLVTVDPAGPACHTGHISCFYRHYPEGKDLT"],
+            ["MDLKQYVSEVQDWPKPGVSFKDITTIMDNGEAYGYATDKIVEYAKDRDVDIVVGPEARGFIIGCPVAYSMGIGFAPVRKEGKLPREVIRYEYDLEYGTNVLTMHKDAIKPGQRVLITDDLLATGGTIEAAIKLVEKLGGIVVGIAFIIELKYLNGIEKIKDYDVMSLISYDE"]
+        ],
+
+    "structure": [
+            ["dddddddddddddddpdpppvcppvnvvvvvvvvvvvvvvvvvvvvvvvvvvqdpqdedeqvrddpcqqpvqhkhkykafwappqwdddpqkiwtwghnppgiaieieghdappqddhrfikifiaghdpvrhtygdhidtdddpddddvvnvvvcvvvvndpdd"],
+            ["dddadcpvpvqkakefeaeppprdtadiaiagpvqvvvcvvpqwhwgqdpvvrdidgqcpvpvqiwrwddwdaddnrryiytythtpahsdpvrhvhpppadvvgpddpd"],
+            ["dplvvqwdwdaqpphhpdtdthcvscvvppvslvvqlvvvlvvcvvqvaqeeeeepdqrcsnrvsscvvvvhyywykyfpppddaawdwdwdddppgitiiithlpseaaageyeyegaeqalqprvlrvvvrcvvnnyddaeyeyqeyevcrvncvsvvvhhydyvyydpd"]
+        ],
+
+    "text": [
+        ["Proteins with zinc bindings."],
+        ["Proteins locating at cell membrane."],
+        ["Protein that serves as an enzyme."]
+    ],
+}
 
 # Databases for different modalities
 now_db = {
@@ -60,12 +74,11 @@ def plot(scores) -> None:
 
 
 # Search from database
-def search(input: str, nprobe: int, topk: int, input_type: str, query_type: str, subsection_type: str):
+def search(input: str, nprobe: int, topk: int, input_type: str, query_type: str, subsection_type: str, db: str):
     input_modality = input_type.replace("sequence", "protein")
     with torch.no_grad():
         input_embedding = getattr(model, f"get_{input_modality}_repr")([input]).cpu().numpy()
 
-    db = now_db[query_type]
     if query_type == "text":
         index = all_index["text"][db][subsection_type]["index"]
         ids = all_index["text"][db][subsection_type]["ids"]
@@ -139,26 +152,6 @@ def search(input: str, nprobe: int, topk: int, input_type: str, query_type: str,
 def change_input_type(choice: str):
     # Change examples if input type is changed
     global samples
-    if choice == "text":
-        samples = [
-            ["Proteins with zinc bindings."],
-            ["Proteins locating at cell membrane."],
-            ["Protein that serves as an enzyme."]
-        ]
-
-    elif choice == "sequence":
-        samples = [
-            ["MSATAEQNARNPKGKGGFARTVSQRKRKRLFLIGGALAVLAVAVGLMLTAFNQDIRFFRTPADLTEQDMTSGARFRLGGLVEEGSVSRTGSELRFTVTDTIKTVKVVFEGIPPDLFREGQGVVAEGRFGSDGLFRADNVLAKHDENYVPKDLADSLKKKGVWEGK"],
-            ["MITLDWEKANGLITTVVQDATTKQVLMVAYMNQESLAKTMATGETWFWSRSRKTLWHKGATSGNIQTVKTIAVDCDADTLLVTVDPAGPACHTGHISCFYRHYPEGKDLT"],
-            ["MDLKQYVSEVQDWPKPGVSFKDITTIMDNGEAYGYATDKIVEYAKDRDVDIVVGPEARGFIIGCPVAYSMGIGFAPVRKEGKLPREVIRYEYDLEYGTNVLTMHKDAIKPGQRVLITDDLLATGGTIEAAIKLVEKLGGIVVGIAFIIELKYLNGIEKIKDYDVMSLISYDE"]
-        ]
-
-    elif choice == "structure":
-        samples = [
-            ["dddddddddddddddpdpppvcppvnvvvvvvvvvvvvvvvvvvvvvvvvvvqdpqdedeqvrddpcqqpvqhkhkykafwappqwdddpqkiwtwghnppgiaieieghdappqddhrfikifiaghdpvrhtygdhidtdddpddddvvnvvvcvvvvndpdd"],
-            ["dddadcpvpvqkakefeaeppprdtadiaiagpvqvvvcvvpqwhwgqdpvvrdidgqcpvpvqiwrwddwdaddnrryiytythtpahsdpvrhvhpppadvvgpddpd"],
-            ["dplvvqwdwdaqpphhpdtdthcvscvvppvslvvqlvvvlvvcvvqvaqeeeeepdqrcsnrvsscvvvvhyywykyfpppddaawdwdwdddppgitiiithlpseaaageyeyegaeqalqprvlrvvvrcvvnnyddaeyeyqeyevcrvncvsvvvhhydyvyydpd"]
-        ]
     
     # Set visibility of upload button
     if choice == "text":
@@ -166,12 +159,12 @@ def change_input_type(choice: str):
     else:
         visible = True
     
-    return gr.update(samples=samples), "", gr.update(visible=visible), gr.update(visible=visible)
+    return gr.update(samples=samples[choice]), "", gr.update(visible=visible), gr.update(visible=visible)
 
 
 # Load example from dataset
 def load_example(example_id):
-    return samples[example_id][0]
+    return example_id[0]
  
  
 # Change the visibility of subsection type
@@ -299,7 +292,7 @@ def build_search_module():
                 # Plot the distribution of scores
                 histogram = gr.Image(label="Histogram of matching scores", type="filepath", scale=1, visible=False)
             
-        search_btn.click(fn=search, inputs=[input, nprobe, topk, input_type, query_type, subsection_type],
+        search_btn.click(fn=search, inputs=[input, nprobe, topk, input_type, query_type, subsection_type, db_type],
                       outputs=[results, download_btn, histogram])
         
         clear_btn.click(fn=clear_results, outputs=[results, download_btn, histogram])