init commit

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.9-slim
+FROM python:3.12-slim
 
 WORKDIR /app
 
@@ -18,4 +18,4 @@ EXPOSE 8501
 
 HEALTHCHECK CMD curl --fail http://localhost:8501/_stcore/health
 
-ENTRYPOINT ["streamlit", "run", "src/streamlit_app.py", "--server.port=8501", "--server.address=0.0.0.0"]
+ENTRYPOINT ["streamlit", "run", "src/app.py", "--server.port=8501", "--server.address=0.0.0.0"]

requirements.txt

@@ -1,3 +1,9 @@
-altair
-pandas
-streamlit
+matchms==0.27.0
+pandas==2.2.3
+matplotlib==3.7.2
+numba==0.59.1
+numpy==1.26.4
+rdkit==2024.9.6
+seaborn==0.13.2
+streamlit==1.44.1
+torch==2.2.0

src/app.py

@@ -0,0 +1,372 @@
+import os
+import tempfile
+from typing import Sequence
+
+import torch
+from torch.utils.data import DataLoader
+import numpy as np
+import pandas as pd
+import seaborn as sns
+sns.set_style("whitegrid")
+sns.set_palette("deep")
+import streamlit as st
+import matplotlib.pyplot as plt
+from matplotlib.container import StemContainer
+from matchms import Spectrum
+from rdkit import Chem
+from rdkit.Chem import Draw
+
+from type import TokenizerConfig
+from data import Tokenizer, TestDataset
+from model import SiameseModel
+from tester import ModelTester
+from utils import top_k_indices, cosine_similarity, read_raw_spectra
+
+torch.classes.__path__ = [os.path.join(torch.__path__[0], torch.classes.__file__)]
+
+PAGE_SIZE = 5
+BATCH_SIZE = 64
+LOADER_BATCH_SIZE = 32
+CANDIDATE_PAGE = [2, 5, 10, 20]
+SHOW_PROGRESS_BAR = False
+
+device = torch.device("cpu")
+tokenizer_config = TokenizerConfig(
+    max_len=100,
+    show_progress_bar=SHOW_PROGRESS_BAR
+)
+tokenizer = Tokenizer(100, SHOW_PROGRESS_BAR)
+model = SiameseModel(
+    embedding_dim=512,
+    n_head=16,
+    n_layer=4,
+    dim_feedward=512,
+    dim_target=512,
+    feedward_activation="selu"
+)
+model_state = torch.load("model.ckpt", map_location=device)
+model.load_state_dict(model_state)
+tester = ModelTester(model, device, SHOW_PROGRESS_BAR)
+
+def custom_stemcontainer(stem_container: StemContainer):
+    stem_container.markerline.set_marker("")
+    stem_container.baseline.set_color("none")
+    stem_container.baseline.set_alpha(0.5)
+
+def draw_mol(smiles: str):
+    mol = Chem.MolFromSmiles(smiles)
+    image = Draw.MolToImage(mol)
+    return image
+
+def plot_pair(q: Spectrum, r: Spectrum):
+    q_peaks = q.peaks.to_numpy
+    r_peaks = r.peaks.to_numpy
+    fig, ax = plt.subplots(1, 1, figsize=(5, 2.7), dpi=300)
+    ax.text(0.8, 0.8, "query", transform=ax.transAxes)
+    ax.text(0.8, 0.2, "reference", transform=ax.transAxes)
+    container1 = ax.stem(q_peaks[:, 0], q_peaks[:, 1])
+    custom_stemcontainer(container1)
+    container2 = ax.stem(r_peaks[:, 0], -r_peaks[:, 1])
+    custom_stemcontainer(container2)
+    return fig
+
+def generate_result():
+    ref_smiles = st.session_state.ref_smiles
+    match_indices = st.session_state.match_indices
+    df = pd.DataFrame(columns=["ID", "Smiles"])
+    for i, index in enumerate(match_indices):
+        df.loc[len(df)] = [i + 1, ref_smiles[index]]
+    st.session_state.result = df.to_csv(index=False).encode("utf8")
+
+def get_smiles(spectra: Sequence[Spectrum]):
+    smiles_seq = [
+        s.get("smiles", "")
+        for s in spectra
+    ]
+    return np.array(smiles_seq)
+
+def batch_match(
+    progress_bar,
+    query_embedding,
+    ref_embedding
+):
+    length = len(query_embedding)
+    start_seq, end_seq = gen_start_end_seq(length)
+    indices = []
+
+    progress = 0
+    for start, end in zip(start_seq, end_seq):
+        batch_embedding = query_embedding[start:end]
+        cosine_scores = cosine_similarity(batch_embedding, ref_embedding)
+        batch_indices = top_k_indices(cosine_scores, 1)
+        indices.append(batch_indices)
+        if progress + BATCH_SIZE >= length:
+            progress = length - 1
+        else:
+            progress += BATCH_SIZE
+        progress_bar.progress((progress + 1) / length)
+
+    return np.concatenate(indices, axis=0)[:, 0]
+
+
+def init_session_state():
+    if "query_path" not in st.session_state:
+        st.session_state.query_path = None
+    
+    if "ref_path" not in st.session_state:
+        st.session_state.ref_path = None
+
+    if "data_len" not in st.session_state:
+        st.session_state.data_len = None
+
+    if "query_embedding" not in st.session_state:
+        st.session_state.query_embedding = None
+    
+    if "ref_embedding" not in st.session_state:
+        st.session_state.ref_embedding = None
+    
+    if "query_smiles" not in st.session_state:
+        st.session_state.query_smiles = None
+    
+    if "ref_smiles" not in st.session_state:
+        st.session_state.ref_smiles = None
+
+    if "query_spectra" not in st.session_state:
+        st.session_state.query_spectra = None
+    
+    if "ref_spectra" not in st.session_state:
+        st.session_state.ref_spectra = None
+    
+    if "match_indices" not in st.session_state:
+        st.session_state.match_indices = None
+
+    if "current_page" not in st.session_state:
+        st.session_state.current_page = None
+    
+    if "last_page" not in st.session_state:
+        st.session_state.last_page = None
+    
+    if "page_size" not in st.session_state:
+        st.session_state.page_size = PAGE_SIZE
+
+def previous_page():
+    current_page = st.session_state.current_page
+    if current_page != 1:
+        st.session_state.current_page -= 1
+
+def next_page():
+    current_page = st.session_state.current_page
+    last_page = st.session_state.last_page
+    if current_page != last_page:
+        st.session_state.current_page += 1
+
+def select_page():
+    st.session_state.current_page = int(st.session_state.page_selector)
+
+def set_page_size():
+    st.session_state.current_page = 1
+    page_size = int(st.session_state.page_size_selector)
+    st.session_state.page_size = page_size
+    cal_page_num(st.session_state.data_len, page_size)
+
+def cal_page_num(
+    length: int,
+    page_size: int
+):
+    page_num, rest = divmod(length, page_size)
+    if rest != 0:
+        page_num += 1
+    st.session_state.last_page = page_num
+
+def gen_start_end_seq(
+    length: int,
+):
+    start_seq = range(0, length, BATCH_SIZE)
+    end_seq = range(BATCH_SIZE, length + BATCH_SIZE, BATCH_SIZE)
+    return start_seq, end_seq
+
+def embedding(
+    progress_bar, 
+    tester: ModelTester, 
+    tokenizer: Tokenizer,
+    spectra: Sequence[Spectrum], 
+):
+    sequences = tokenizer.tokenize_sequence(spectra)
+    start_seq, end_seq = gen_start_end_seq(len(spectra))
+    progress = 0
+    embedding = []
+    for start, end in zip(start_seq, end_seq):
+        test_dataset = TestDataset(sequences[start:end])
+        test_dataloader = DataLoader(
+            test_dataset,
+            LOADER_BATCH_SIZE,
+            False
+        )
+        step_embedding = tester.test(test_dataloader)
+        if progress + BATCH_SIZE >= len(spectra):
+            progress = len(spectra) - 1
+        else:
+            progress += BATCH_SIZE
+
+        embedding.append(step_embedding)
+        progress_bar.progress((progress + 1) / len(spectra))
+
+    embedding = np.concatenate(embedding, axis=0)
+    return embedding
+
+def main():
+    st.set_page_config(layout="wide")
+    st.title("SpecEmbedding")
+    tab1, tab2, tab3 = st.tabs(["upload query file", "upload reference/library file", "library match"])
+
+    with tab1:
+        st.header("Upload query spectra file(positive mode)")
+        query_file = st.file_uploader(
+            "upload the query spectra file", 
+            type=["msp", "mgf", "mzxml"], 
+            key="query_file",
+            accept_multiple_files=False
+        )
+        query_embedding_btn = st.button("Embedding", "query_embedding_btn")
+        query_status_box = st.empty()
+        if query_embedding_btn:
+            if query_file is not None:
+                with tempfile.NamedTemporaryFile(delete=True, suffix="." + query_file.name.split(".")[-1]) as tmp_file:
+                    tmp_file.write(query_file.getvalue())
+                    query_spectra = read_raw_spectra(tmp_file.name)
+
+                progress_bar = st.progress(0, text="Embedding...")
+                st.session_state.data_len = len(query_spectra)
+                st.session_state.query_spectra = query_spectra
+                st.session_state.query_smiles = get_smiles(query_spectra)
+                query_embedding = embedding(
+                    progress_bar,
+                    tester, 
+                    tokenizer, 
+                    query_spectra, 
+                )
+                st.session_state.query_embedding = query_embedding
+                query_status_box.success("Embedding Success ✅")
+            else:
+                query_status_box.error("Please upload the spectra file")
+
+    with tab2:
+        st.header("Upload reference/library spectra file(positive mode)")
+        ref_file = st.file_uploader(
+            "upload the reference/library spectra file", 
+            type=["msp", "mgf", "mzxml"], 
+            key="ref_file", 
+            accept_multiple_files=False
+        )
+        ref_embedding_btn = st.button("Embedding", "ref_embedding_btn")
+        ref_status_box = st.empty()
+        if ref_embedding_btn:
+            if ref_file is not None:
+                progress_bar = st.progress(0, text="Embedding...")
+                with tempfile.NamedTemporaryFile(delete=True, suffix="." + ref_file.name.split(".")[-1]) as tmp_file:
+                    tmp_file.write(ref_file.getvalue())
+                    ref_spectra = read_raw_spectra(tmp_file.name)
+                
+                st.session_state.ref_spectra = ref_spectra
+                st.session_state.ref_smiles = get_smiles(ref_spectra)
+                ref_embedding = embedding(
+                    progress_bar,
+                    tester,
+                    tokenizer,
+                    ref_spectra,
+                )
+                st.session_state.ref_embedding = ref_embedding
+                ref_status_box.success("Embedding Success ✅")
+            else:
+                ref_status_box.error("Please upload the spectra file")
+
+    with tab3:
+        st.header("Start to match")
+        launch_btn = st.button("Launch", key="launch_btn")
+        match_status_box = st.empty()
+        if launch_btn:
+            query_embedding = st.session_state.query_embedding
+            ref_embedding = st.session_state.ref_embedding
+            if query_embedding is None:
+                match_status_box.error("No query embedding")
+            elif ref_embedding is None:
+                match_status_box.error("No reference embedding")
+            else:
+                progress_bar = st.progress(0, "Match...")
+                match_indices = batch_match(progress_bar, query_embedding, ref_embedding)
+                st.session_state.match_indices = match_indices
+                st.session_state.current_page = 1
+                generate_result()
+                cal_page_num(st.session_state.data_len, st.session_state.page_size)
+                match_status_box.success("match success")
+
+        if st.session_state.match_indices is not None:
+            st.subheader(f"Match Result")
+            current_page = st.session_state.current_page
+            last_page = st.session_state.last_page
+
+            ref_smiles = st.session_state.ref_smiles
+            query_spectra = st.session_state.query_spectra
+            ref_spectra = st.session_state.ref_spectra
+            page_size = st.session_state.page_size
+
+            indices = st.session_state.match_indices
+            start = (current_page - 1) * page_size
+            end = start + page_size
+        
+            if current_page == last_page:
+                end = indices.shape[0]
+
+            col1, col2, _ = st.columns([1, 1, 5])
+
+            col1.selectbox(
+                "page size",
+                CANDIDATE_PAGE,
+                key="page_size_selector",
+                disabled=False,
+                label_visibility="collapsed",
+                index=CANDIDATE_PAGE.index(page_size),
+                on_change=set_page_size,
+            )
+
+            col2.download_button(
+                label="download result",
+                data=st.session_state.result,
+                file_name="data.csv",
+                mime="text/csv"
+            )
+
+            pre_btn, current, next_btn, page_selector, _ =  st.columns([1, 1, 1, 1, 2])
+            pre_btn.button("previous page", key="pre_btn", on_click=previous_page)
+            current.subheader(f"current page: {current_page}")
+            next_btn.button("next page", key="next_btn", on_click=next_page)
+            page_selector.selectbox(
+                label="target page",
+                key="page_selector",
+                options=range(1, last_page + 1),
+                disabled=False,
+                index=current_page - 1,
+                label_visibility="collapsed",
+                on_change=select_page,
+            )
+
+            col1, col2, col3, col4 = st.columns([1, 4, 6, 4])
+            col1.subheader("Index")
+            col2.subheader("Smiles")
+            col3.subheader("MS/MS Spectra Pair")
+            col4.subheader("Molecular Structure")
+
+            for i in range(start, end):
+                query_index = i
+                ref_index = indices[i]
+                id_label, smiles_label, pair_viewer, mol_viewer = st.columns([2, 4, 6, 4])
+                id_label.subheader(i + 1)
+                smiles_label.text(ref_smiles[ref_index])
+                pair_fig = plot_pair(query_spectra[query_index], ref_spectra[ref_index])
+                pair_viewer.pyplot(pair_fig, use_container_width=True)
+                mol_image = draw_mol(ref_smiles[ref_index])
+                mol_viewer.image(mol_image, use_container_width=True)
+
+if __name__ == "__main__":
+    init_session_state()
+    main()

src/data.py

@@ -0,0 +1,118 @@
+from typing import Sequence
+from collections.abc import Sequence
+
+import numpy as np
+from tqdm import tqdm
+from matchms import Spectrum
+from torch.utils.data import Dataset
+
+from type import Peak, MetaData, TokenSequence
+
+SpecialToken = {
+    "PAD": 0,
+}
+
+class TestDataset(Dataset):
+    def __init__(self, sequences: list[TokenSequence]) -> None:
+        super(TestDataset, self).__init__()
+        self._sequences = sequences
+        self.length = len(sequences)
+
+    def __len__(self):
+        return self.length
+
+    def __getitem__(self, index: int):
+        sequence = self._sequences[index]
+        return sequence["mz"], sequence["intensity"], sequence["mask"]
+
+class Tokenizer:
+    def __init__(self, max_len: int, show_progress_bar: bool = True) -> None:
+        """
+            Tokenization of mass spectrometry data
+
+            Parameters:
+            ---
+            -   max_len: Maximum number of peaks to extract
+            -   show_progress_bar: Whether to display a progress bar
+        """
+        self.max_len = max_len
+        self.show_progress_bar = show_progress_bar
+
+    def tokenize(self, s: Spectrum):
+        """
+            Tokenization of mass spectrometry data
+        """
+        metadata = self.get_metadata(s)
+        mz = []
+        intensity = []
+        for peak in metadata["peaks"]:
+            mz.append(peak["mz"])
+            intensity.append(peak["intensity"])
+
+        mz = np.array(mz)
+        intensity = np.array(intensity)
+        mask = np.zeros((self.max_len, ), dtype=bool)
+        if len(mz) < self.max_len:
+            mask[len(mz):] = True
+            mz = np.pad(
+                mz, (0, self.max_len - len(mz)),
+                mode='constant', constant_values=SpecialToken["PAD"]
+            )
+
+            intensity = np.pad(
+                intensity, (0, self.max_len - len(intensity)),
+                mode='constant', constant_values=SpecialToken["PAD"]
+            )
+
+        return TokenSequence(
+            mz=np.array(mz, np.float32),
+            intensity=np.array(intensity, np.float32),
+            mask=mask,
+            smiles=metadata["smiles"]
+        )
+
+    def tokenize_sequence(self, spectra: Sequence[Spectrum]):
+        sequences: list[TokenSequence] = []
+        pbar = spectra
+        if self.show_progress_bar:
+            pbar = tqdm(spectra, total=len(spectra), desc="tokenization")
+        for s in pbar:
+            sequences.append(self.tokenize(s))
+
+        return sequences
+
+    def get_metadata(self, s: Spectrum):
+        """
+            get the metadata from spectrum
+
+            -   smiles
+            -   precursor_mz
+            -   peaks
+        """
+        precursor_mz = s.get("precursor_mz")
+        smiles = s.get("smiles")
+        peaks = np.array(s.peaks.to_numpy, np.float32)
+        intensity = peaks[:, 1]
+        argmaxsort_index = np.sort(
+            np.argsort(intensity)[::-1][:self.max_len - 1]
+        )
+        peaks = peaks[argmaxsort_index]
+        peaks[:, 1] = peaks[:, 1] / max(peaks[:, 1])
+        packaged_peaks: list[Peak] = [
+            Peak(
+                mz=np.array(precursor_mz, np.float32),
+                intensity=2
+            )
+        ]
+        for mz, intensity in peaks:
+            packaged_peaks.append(
+                Peak(
+                    mz=mz,
+                    intensity=intensity
+                )
+            )
+        metadata = MetaData(
+            smiles=smiles,
+            peaks=packaged_peaks
+        )
+        return metadata

src/model.ckpt

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

src/model.py

@@ -0,0 +1,274 @@
+import math
+from typing import Literal, Union, Iterable, Tuple
+
+import torch
+import torch.nn as nn
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+
+LAMBDA_MIN = math.pow(10, -3.0)
+LAMBDA_MAX = math.pow(10, 3.0)
+
+class MultiFeedForwardModule(nn.Module):
+    def __init__(
+        self,
+        input_size: int,
+        hidden_size: Union[int, Iterable[int]],
+        output_size: int,
+        *,
+        activation: Literal['relu', 'selu', 'gelu'] = 'relu',
+        dropout: float = 0.1,
+        dropout_last_layer: bool = True
+    ):
+        super(MultiFeedForwardModule, self).__init__()
+        if activation == 'relu':
+            self._activation = nn.ReLU()
+        elif activation == 'selu':
+            self._activation = nn.SELU()
+        elif activation == 'gelu':
+            self._activation = nn.GELU()
+        else:
+            raise ValueError('activation must be relu or selu')
+
+        if not hasattr(hidden_size, '__iter__'):
+            if hidden_size is None:
+                hidden_size = [output_size]
+            else:
+                hidden_size = [hidden_size]
+
+        self._layers = []
+        layer_dims = [input_size] + hidden_size + [output_size]
+
+        for i in range(1, len(layer_dims) - 1):
+            self._layers.append(nn.Linear(layer_dims[i - 1], layer_dims[i]))
+            self._layers.append(self._activation)
+            self._layers.append(nn.Dropout(dropout))
+
+        self._layers.append(nn.Linear(layer_dims[-2], layer_dims[-1]))
+
+        if dropout_last_layer:
+            self._layers.append(nn.Dropout(dropout))
+        self._layers = nn.Sequential(*self._layers)
+
+    def forward(self, x):
+        return self._layers(x)
+
+
+class SinusodialMz(nn.Module):
+    def __init__(self, embedding_dim: int, *, lambda_params: Tuple[float, float] = (LAMBDA_MIN, LAMBDA_MAX)) -> None:
+        super(SinusodialMz, self).__init__()
+        self.lambda_min, self.lambda_max = lambda_params
+        self.lambda_div_value = self.lambda_max / self.lambda_min
+        self.x = torch.arange(0, embedding_dim, 2)
+        self.x = (
+            2 * math.pi *
+            (
+                self.lambda_min *
+                self.lambda_div_value ** (self.x / (embedding_dim - 2))
+            ) ** -1
+        )
+
+    def forward(self, mz: torch.Tensor):
+        self.x = self.x.to(mz.device)
+        x = torch.einsum('bl,d->bld', mz, self.x)
+        sin_embedding = torch.sin(x)
+        cos_embedding = torch.cos(x)
+        b, l, d = sin_embedding.shape
+        x = torch.zeros(b, l, 2 * d, dtype=mz.dtype, device=mz.device)
+        x[:, :, ::2] = sin_embedding
+        x[:, :, 1::2] = cos_embedding
+        return x
+
+
+class SinusodialMzEmbedding(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        *,
+        lambda_params: Tuple[float, float] = (LAMBDA_MIN, LAMBDA_MAX),
+        feedward_activation: Literal['relu', 'selu', 'gelu'] = 'relu',
+        dropout: float = 0.1,
+        dropout_last_layer: bool = True
+    ):
+        super(SinusodialMzEmbedding, self).__init__()
+        if embedding_dim % 2 != 0:
+            raise ValueError('embedding_dim must be even')
+        self.embedding = SinusodialMz(
+            embedding_dim, lambda_params=lambda_params)
+        self.feedward_layers = MultiFeedForwardModule(
+            embedding_dim, embedding_dim, embedding_dim,
+            activation=feedward_activation, dropout=dropout, dropout_last_layer=dropout_last_layer
+        )
+
+    def forward(self, mz: torch.Tensor):
+        x = self.embedding(mz)
+        x = self.feedward_layers(x)
+        return x
+
+
+class PeaksEmbedding(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        *,
+        lambda_params: Tuple[float, float] = (LAMBDA_MIN, LAMBDA_MAX),
+        feedward_activation: Literal['relu', 'selu', 'gelu'] = 'relu',
+        dropout: float = 0.1,
+        dropout_last_layer: bool = False
+    ) -> None:
+        super(PeaksEmbedding, self).__init__()
+        self.mz_embedding = SinusodialMzEmbedding(
+            embedding_dim,
+            lambda_params=lambda_params,
+            feedward_activation=feedward_activation,
+            dropout=dropout,
+            dropout_last_layer=dropout_last_layer
+        )
+        self.intensity_embedding = MultiFeedForwardModule(
+            embedding_dim + 1, embedding_dim, embedding_dim,
+            activation=feedward_activation,
+            dropout=dropout,
+            dropout_last_layer=dropout_last_layer
+        )
+
+    def forward(self, mz: torch.Tensor, intensity: torch.Tensor):
+        mz_tensor = self.mz_embedding(mz)
+        intensity_tensor = torch.unsqueeze(intensity, dim=-1)
+        x = self.intensity_embedding(
+            torch.cat([mz_tensor, intensity_tensor], dim=-1))
+        return x
+
+
+class SiameseModel(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        n_head: int,
+        n_layer: int,
+        dim_feedward: int,
+        dim_target: int,
+        *,
+        lambda_params: Tuple[float, float] = (LAMBDA_MIN, LAMBDA_MAX),
+        feedward_activation: Literal['relu', 'selu', 'gelu'] = 'relu',
+        dropout: float = 0.1,
+        dropout_last_layer: bool = False,
+        norm_first: bool = True
+    ) -> None:
+        super(SiameseModel, self).__init__()
+        if embedding_dim % n_head != 0:
+            raise ValueError('embedding must be divisible by n_head')
+
+        self.embedding = PeaksEmbedding(
+            embedding_dim,
+            lambda_params=lambda_params,
+            feedward_activation=feedward_activation,
+            dropout=dropout,
+            dropout_last_layer=dropout_last_layer
+        )
+
+        if feedward_activation == 'selu':
+            # transformer encoder activation
+            # only gelu or relu
+            self.activation = 'gelu'
+        else:
+            self.activation = feedward_activation
+
+        if feedward_activation == 'relu':
+            self._activation = nn.ReLU()
+        elif feedward_activation == 'selu':
+            self._activation = nn.SELU()
+        elif feedward_activation == 'gelu':
+            self._activation = nn.GELU()
+        else:
+            raise ValueError('activation must be relu or selu or gelu')
+
+        encoder_layer = TransformerEncoderLayer(
+            embedding_dim,
+            n_head,
+            dim_feedforward=dim_feedward,
+            dropout=dropout,
+            activation=self.activation,
+            batch_first=True,
+            norm_first=norm_first
+        )
+        self._encoder = TransformerEncoder(
+            encoder_layer,
+            n_layer,
+            enable_nested_tensor=False
+        )
+
+        self._decoder = MultiFeedForwardModule(
+            embedding_dim,
+            dim_feedward,
+            dim_target,
+            activation=feedward_activation,
+            dropout=dropout,
+            dropout_last_layer=dropout_last_layer
+        )
+
+    def forward(self, mz: torch.Tensor, intensity: torch.Tensor, mask: torch.Tensor):
+        x = self.embedding(mz, intensity)
+        x = self._encoder(x, src_key_padding_mask=mask)
+        # mean pooling or cls position vector
+        x = torch.mean(x, dim=1)
+        x = self._activation(self._decoder(x))
+        return x
+
+
+# class MambaSiameseModel(nn.Module):
+#     def __init__(
+#         self,
+#         embedding_dim: int,
+#         n_layer: int,
+#         dim_feedward: int,
+#         dim_target: int,
+#         *,
+#         lambda_params: Tuple[float, float] = (LAMBDA_MIN, LAMBDA_MAX),
+#         feedward_activation: Literal['relu', 'selu', 'gelu'] = 'relu',
+#         dropout: float = 0.1,
+#         dropout_last_layer: bool = False,
+#     ):
+#         super(MambaSiameseModel, self).__init__()
+
+#         self.embedding = PeaksEmbedding(
+#             embedding_dim,
+#             lambda_params=lambda_params,
+#             feedward_activation=feedward_activation,
+#             dropout=dropout,
+#             dropout_last_layer=dropout_last_layer
+#         )
+
+#         if feedward_activation == 'relu':
+#             self._activation = nn.ReLU()
+#         elif feedward_activation == 'selu':
+#             self._activation = nn.SELU()
+#         elif feedward_activation == 'gelu':
+#             self._activation = nn.GELU()
+#         else:
+#             raise ValueError('activation must be relu or selu or gelu')
+
+#         self._encoder = nn.Sequential(*[
+#             Mamba2(
+#                 d_model=embedding_dim,
+#                 d_state=64,
+#                 d_conv=4,
+#                 expand=2
+#             )
+#             for _ in range(n_layer)
+#         ])
+
+#         self._decoder = MultiFeedForwardModule(
+#             embedding_dim,
+#             dim_feedward,
+#             dim_target,
+#             activation=feedward_activation,
+#             dropout=dropout,
+#             dropout_last_layer=dropout_last_layer
+#         )
+
+#     def forward(self, mz: torch.Tensor, intensity: torch.Tensor, mask: torch.Tensor):
+#         x = self.embedding(mz, intensity)
+#         x = self._encoder(x)
+#         # mean pooling or cls position vector
+#         x = torch.mean(x, dim=1)
+#         x = self._activation(self._decoder(x))
+#         return x

src/streamlit_app.py

@@ -1,40 +0,0 @@
-import altair as alt
-import numpy as np
-import pandas as pd
-import streamlit as st
-
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))

src/tester.py

@@ -0,0 +1,30 @@
+import torch
+from torch.nn import Module
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+import numpy as np
+
+class ModelTester:
+    def __init__(
+        self,
+        model: Module,
+        device: torch.device,
+        show_prgress_bar: bool = True
+    ) -> None:
+        self.model = model
+        self.device = device
+        self.show_prgress_bar = show_prgress_bar
+
+    def test(self, dataloader: DataLoader):
+        self.model.eval()
+        result = []
+        with torch.no_grad():
+            pbar = dataloader
+            if self.show_prgress_bar:
+                pbar = tqdm(dataloader, total=len(
+                    dataloader), desc="embedding")
+            for x in pbar:
+                x = [d.to(self.device) for d in x]
+                pred: torch.Tensor = self.model(*x)
+                result.append(pred.cpu().numpy())
+        return np.concatenate(result, axis=0)

src/type.py

@@ -0,0 +1,34 @@
+from typing import TypedDict, Sequence, Callable, Optional
+
+import torch
+from torch import nn
+from torch import device
+import numpy as np
+import numpy.typing as npt
+
+BatchType = Sequence[torch.Tensor]
+StepTrain = Callable[[nn.Module, nn.Module, device,
+                      BatchType, Optional[Callable[..., int]]], Sequence[torch.Tensor]]
+StepVal = Callable[[nn.Module, nn.Module, device,
+                    BatchType, Optional[Callable[..., int]]], Sequence[torch.Tensor]]
+
+class Peak(TypedDict):
+    mz: str
+    intensity: npt.NDArray
+
+
+class MetaData(TypedDict):
+    peaks: Sequence[Peak]
+    smiles: str
+
+
+class TokenSequence(TypedDict):
+    mz: npt.NDArray[np.int32]
+    intensity: npt.NDArray[np.float32]
+    mask: npt.NDArray[np.bool_]
+    smiles: str
+
+
+class TokenizerConfig(TypedDict):
+    max_len: int
+    show_progress_bar: bool

src/utils.py

@@ -0,0 +1,41 @@
+from pathlib import Path
+
+import numpy as np
+import numpy.typing as npt
+from numba import prange, njit
+from matchms.importing import load_from_mgf, load_from_msp, load_from_mzxml
+from matchms.filtering import default_filters, normalize_intensities
+
+def read_raw_spectra(path: str):
+    suffix = Path(path).suffix
+    if suffix == ".mgf":
+        spectra = list(load_from_mgf(path))
+    elif suffix == ".msp":
+        spectra = list(load_from_msp(path))
+    elif suffix == ".mzxml":
+        spectra = list(load_from_mzxml(path))
+    else:
+        raise ValueError(f"Not support the {suffix} format")
+    
+    spectra = [default_filters(s) for s in spectra]
+    spectra = [normalize_intensities(s) for s in spectra]
+    return spectra
+
+@njit
+def cosine_similarity(A: npt.NDArray, B: npt.NDArray):
+    norm_A = np.sqrt(np.sum(A ** 2, axis=1)) + 1e-8
+    norm_B = np.sqrt(np.sum(B ** 2, axis=1)) + 1e-8
+    normalize_A = A / norm_A[:, np.newaxis]
+    normalize_B = B / norm_B[:, np.newaxis]
+    scores = np.dot(normalize_A, normalize_B.T)
+    return scores
+
+@njit(parallel=True)
+def top_k_indices(score, top_k):
+    rows, cols = score.shape
+    indices = np.empty((rows, top_k), dtype=np.int64)
+    for i in prange(rows):
+        row = score[i]
+        sorted_idx = np.argsort(row)[::-1]
+        indices[i] = sorted_idx[:top_k]
+    return indices