Update app.py

app.py

@@ -32,6 +32,7 @@ from gradio.themes.base import Base
 from gradio.themes.utils import colors, fonts, sizes
 from huggingface_hub import hf_hub_download
 import time
+import concurrent.futures
 
 class Seafoam(Base):
     def __init__(
@@ -260,40 +261,48 @@ def retrieve_similarity_scores( table_name, target_mass,collision_energy, ms2_em
     filtered_smiles = cur.fetchall()
     similarity_scores = []
     
-    for smile in filtered_smiles:
-        query = f"""
-            SELECT low_energy_embedding, median_energy_embedding, high_energy_embedding
-            FROM {table_name}
-            WHERE SMILES = ?
-        """
-        cur.execute(query, (smile[0],))
-        row = cur.fetchone()
-        if row is None:
-            return None
-        low_energy_embedding_db = np.array(pickle.loads(row[0]), dtype=np.float64)
-        median_energy_embedding_db = np.array(pickle.loads(row[1]), dtype=np.float64)
-        high_energy_embedding_db = np.array(pickle.loads(row[2]), dtype=np.float64)
-        low_energy_embedding_db,median_energy_embedding_db,high_energy_embedding_db = torch.tensor(low_energy_embedding_db).float(),torch.tensor(median_energy_embedding_db).float(),torch.tensor(high_energy_embedding_db).float()
-        low_similarity =(ms2_embedding_low @ low_energy_embedding_db.t()).item()
-        median_similarity = (ms2_embedding_median @  median_energy_embedding_db.t()).item()
-        high_similarity = (ms2_embedding_high @ high_energy_embedding_db.t()).item()
-        '''
-        low_similarity = calculate_cosine_similarity(ms2_embedding_low, low_energy_embedding_db)
-        median_similarity = calculate_cosine_similarity(ms2_embedding_median, median_energy_embedding_db)
-        high_similarity = calculate_cosine_similarity(ms2_embedding_high, high_energy_embedding_db)'''
-        similarity_scores.append((smile, low_similarity, median_similarity, high_similarity))
-
-    weighted_similarity_scores = []
-    for smile, low_similarity, median_similarity, high_similarity in similarity_scores:
-        if collision_energy <=15:
-            weighted_similarity = 0.4 * low_similarity + 0.3 * median_similarity + 0.3 * high_similarity
-            weighted_similarity_scores.append((smile, weighted_similarity))
-        elif collision_energy >15 and collision_energy <= 25:
-            weighted_similarity = 0.3 * low_similarity + 0.4 * median_similarity + 0.3 * high_similarity
-            weighted_similarity_scores.append((smile, weighted_similarity))
-        elif collision_energy > 25:
-            weighted_similarity = 0.2 * low_similarity + 0.3 * median_similarity + 0.5 * high_similarity
-            weighted_similarity_scores.append((smile, weighted_similarity))
+    query = f"""
+    SELECT SMILES, low_energy_embedding, median_energy_embedding, high_energy_embedding
+    FROM {table_name}
+    WHERE SMILES IN ({','.join(['?']*len(filtered_smiles))})
+    """
+    cur.execute(query, tuple(s[0] for s in filtered_smiles))
+    rows = cur.fetchall()
+    
+    def decode_row(row):
+        return (
+            row[0],  # SMILES
+            np.array(pickle.loads(row[1]), dtype=np.float32),
+            np.array(pickle.loads(row[2]), dtype=np.float32),
+            np.array(pickle.loads(row[3]), dtype=np.float32),
+        )
+    
+    with concurrent.futures.ThreadPoolExecutor() as executor:
+        results = list(executor.map(decode_row, rows))
+    
+    ms2_embedding_low_np = ms2_embedding_low.numpy()
+    ms2_embedding_median_np = ms2_embedding_median.numpy()
+    ms2_embedding_high_np = ms2_embedding_high.numpy()
+    
+    similarity_scores = [
+        (
+            smile,
+            np.dot(ms2_embedding_low_np, low_embedding),
+            np.dot(ms2_embedding_median_np, median_embedding),
+            np.dot(ms2_embedding_high_np, high_embedding),
+        )
+        for smile, low_embedding, median_embedding, high_embedding in results
+    ]
+    
+    collision_weights = np.array([
+        [0.4, 0.3, 0.3] if collision_energy <= 15 else
+        [0.3, 0.4, 0.3] if collision_energy <= 25 else
+        [0.2, 0.3, 0.5]
+    ])
+    
+    similarity_array = np.array([[low, median, high] for _, low, median, high in similarity_scores])
+    weighted_similarities = similarity_array @ collision_weights.T  # 矩阵乘法
+    weighted_similarity_scores = [(smile, weighted) for (smile, _low, _med, _high), weighted in zip(similarity_scores, weighted_similarities)]
 
     weighted_similarity_scores.sort(key=lambda x: x[1], reverse=True)