deduplicating on uniprot id

app.py

@@ -42,20 +42,19 @@ def limit_n_results(n):
 
 
 def aggregate_search_results(raw_results: List[dict], max_res: int) -> Dict[str, dict]:
-    aggregated_by_gene = collections.defaultdict(list)
+    aggregated_by_uniprot = collections.defaultdict(list)
     for raw_result in raw_results:
         entry = select_keys(
             raw_result,
-            keys=["pdb_name", "chain_id", "score", "organism", "uniprot_id"]
+            keys=["pdb_name", "chain_id", "score", "organism", "uniprot_id", "genes"],
         )
-        genes = raw_result["genes"]
-        if genes is not None:
-            gene_names = genes.split(" ")
-            for gene in gene_names:
-                aggregated_by_gene[gene].append(entry)
-                if len(aggregated_by_gene) >= max_res:
-                    return dict(aggregated_by_gene)
-    return dict(aggregated_by_gene)
+        uniprot_id = raw_result["uniprot_id"]
+
+        if uniprot_id is not None:
+            aggregated_by_uniprot[uniprot_id].append(entry)
+            if len(aggregated_by_uniprot) >= max_res:
+                return dict(aggregated_by_uniprot)
+    return dict(aggregated_by_uniprot)
 
 
 def select_keys(d: dict, keys: List[str]):
@@ -64,21 +63,28 @@ def select_keys(d: dict, keys: List[str]):
 
 def format_search_results(agg_search_results):
     formatted_search_results = {}
-    for gene, entries in agg_search_results.items():
+    for uniprot_id, entries in agg_search_results.items():
         entry = entries[0]
         organism = entry["organism"]
         score = entry["score"]
-        uniprot_id = entry["uniprot_id"]
-        key = f"Gene: {gene} | Organism: {organism} | Uniprot ID: {uniprot_id}"
+        genes = entry["genes"]
+        key = f"Uniprot ID: {uniprot_id} | Organism: {organism} | Gene Names: {genes}"
         formatted_search_results[key] = score
     return formatted_search_results
 
 
 def update_dropdown_menu(agg_search_res):
     choices = []
-    for gene, entries in agg_search_res.items():
+    for uniprot_id, entries in agg_search_res.items():
         for entry in entries:
-            choice = choice_sep.join([gene, entry["pdb_name"], entry["chain_id"]])
+            choice = choice_sep.join(
+                [
+                    uniprot_id,
+                    entry["pdb_name"],
+                    entry["chain_id"],
+                    entry["genes"] or "",
+                ]
+            )
             choices.append(choice)
 
     if choices:

get_index.py

@@ -18,7 +18,8 @@ def get_engines(index_repo: str, model_repo: str):
     )
     sys.path.append(str(local_arch_path))
     from protein_index import (  # pylint: disable=import-error,import-outside-toplevel
-        ProteinSearchEngine, ProteinIndexError
+        ProteinSearchEngine,
+        ProteinIndexError,
     )
 
     subindex_paths = glob(str(index_path / "*/"))