Update app.py

app.py

@@ -1,10 +1,3 @@
-# -*- coding: utf-8 -*-
-"""
-Created on Thu Oct 10 16:46:36 2024
-
-@author: ZNDX002
-"""
-
 from __future__ import annotations
 import sqlite3
 import gradio as gr
@@ -32,7 +25,6 @@ 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__(
@@ -107,7 +99,6 @@ custom_css = """
        height: 480px !important;
    }
 </style>
-
 """
 def spectrum_processing(s):
     """This is how one would typically design a desired pre- and post-
@@ -261,48 +252,40 @@ def retrieve_similarity_scores( table_name, target_mass,collision_energy, ms2_em
     filtered_smiles = cur.fetchall()
     similarity_scores = []
     
-    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)]
+    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))
 
     weighted_similarity_scores.sort(key=lambda x: x[1], reverse=True)
 
@@ -493,3 +476,4 @@ with gr.Blocks(theme=seafoam) as demo:
 demo.launch(share=True)
 
 
+