rcsb-embedding-model 0.0.5__tar.gz → 0.0.7__tar.gz

rcsb_embedding_model-0.0.7/PKG-INFO

@@ -0,0 +1,126 @@
+Metadata-Version: 2.4
+Name: rcsb-embedding-model
+Version: 0.0.7
+Summary: Protein Embedding Model for Structure Search
+Project-URL: Homepage, https://github.com/rcsb/rcsb-embedding-model
+Project-URL: Issues, https://github.com/rcsb/rcsb-embedding-model/issues
+Author-email: Joan Segura <joan.segura@rcsb.org>
+License-Expression: BSD-3-Clause
+License-File: LICENSE.md
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.10
+Requires-Dist: esm>=3.2.0
+Requires-Dist: lightning>=2.5.0
+Requires-Dist: torch>=2.2.0
+Requires-Dist: typer>=0.15.0
+Description-Content-Type: text/markdown
+
+# RCSB Embedding Model
+
+**Version** 0.0.7
+
+
+## Overview
+
+RCSB Embedding Model is a neural network architecture designed to encode macromolecular 3D structures into fixed-length vector embeddings for efficient large-scale structure similarity search.
+
+Preprint: [Multi-scale structural similarity embedding search across entire proteomes](https://www.biorxiv.org/content/10.1101/2025.02.28.640875v1).
+
+A web-based implementation using this model for structure similarity search is available at [rcsb-embedding-search](http://embedding-search.rcsb.org).
+
+If you are interested in training the model with a new dataset, visit the [rcsb-embedding-search repository](https://github.com/bioinsilico/rcsb-embedding-search), which provides scripts and documentation for training.
+
+
+## Features
+
+- **Residue-level embeddings** computed using the ESM3 protein language model  
+- **Structure-level embeddings** aggregated via a transformer-based aggregator network  
+- **Command-line interface** implemented with Typer for high-throughput inference workflows  
+- **Python API** for interactive embedding computation and integration into analysis pipelines  
+- **High-performance inference** leveraging PyTorch Lightning, with multi-node and multi-GPU support  
+
+---
+
+## Installation
+
+    pip install rcsb-embedding-model
+
+**Requirements:**
+
+- Python ≥ 3.10  
+- ESM ≥ 3.2.0  
+- PyTorch ≥ 2.2.0  
+- Lightning ≥ 2.5.0  
+- Typer ≥ 0.15.0  
+
+---
+
+## Quick Start
+
+### CLI
+
+    # 1. Compute residue embeddings: Calculate residue level embeddings of protein structures using ESM3.
+    inference residue-embedding --src-file data/structures.csv --output-path results/residue_embeddings --src-format mmcif --batch-size 8 --devices auto
+
+    # 2. Compute structure embeddings: Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas data-frame file.
+    inference structure-embedding --src-file results/residue_embeddings.csv --output-path results/structure_embeddings --out-df-id embeddings.pkl --batch-size 4 --devices 0 --devives 1
+
+    # 3. Compute chain embeddings: Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files.
+    inference chain-embedding --src-file results/residue_embeddings.csv --output-path results/chain_embeddings --batch-size 4
+
+### Python API
+
+    from rcsb_embedding_model import RcsbStructureEmbedding
+
+    model = RcsbStructureEmbedding()
+
+    # Compute per-residue embeddings
+    res_emb = model.residue_embedding(
+        src_structure="examples/1abc.cif",
+        src_format="mmcif",
+        chain_id="A"
+    )
+
+    # Aggregate to structure-level embedding
+    struct_emb = model.aggregator_embedding(res_emb)
+
+See the examples directory for complete scripts.
+
+---
+
+## Model Architecture
+
+The embedding model is trained to predict structural similarity by approximating TM-scores using cosine distances between embeddings. It consists of two main components:
+
+- **Protein Language Model (PLM)**: Computes residue-level embeddings from a given 3D structure.
+- **Residue Embedding Aggregator**: A transformer-based neural network that aggregates these residue-level embeddings into a single vector.
+
+![Embedding model architecture](assets/embedding-model-architecture.png)
+
+### **Protein Language Model (PLM)**
+Residue-wise embeddings of protein structures are computed using the [ESM3](https://www.evolutionaryscale.ai/) generative protein language model.
+
+### **Residue Embedding Aggregator**
+The aggregation component consists of six transformer encoder layers, each with a 3,072-neuron feedforward layer and ReLU activations. After processing through these layers, a summation pooling operation is applied, followed by 12 fully connected residual layers that refine the embeddings into a single 1,536-dimensional vector.
+
+---
+
+## Development
+
+    git clone https://github.com/rcsb/rcsb-embedding-model.git
+    cd rcsb-embedding-model
+    pip install -e .
+    pytest
+
+---
+
+## Citation
+
+Segura, J., Bittrich, S., et al. (2024). *Multi-scale structural similarity embedding search across entire proteomes*. bioRxiv. (Preprint: https://www.biorxiv.org/content/10.1101/2024.03.07.XXXXX)
+
+---
+
+## License
+
+This project is licensed under the BSD 3-Clause License. See [LICENSE.md](LICENSE.md) for details.

rcsb_embedding_model-0.0.7/README.md

@@ -0,0 +1,108 @@
+# RCSB Embedding Model
+
+**Version** 0.0.7
+
+
+## Overview
+
+RCSB Embedding Model is a neural network architecture designed to encode macromolecular 3D structures into fixed-length vector embeddings for efficient large-scale structure similarity search.
+
+Preprint: [Multi-scale structural similarity embedding search across entire proteomes](https://www.biorxiv.org/content/10.1101/2025.02.28.640875v1).
+
+A web-based implementation using this model for structure similarity search is available at [rcsb-embedding-search](http://embedding-search.rcsb.org).
+
+If you are interested in training the model with a new dataset, visit the [rcsb-embedding-search repository](https://github.com/bioinsilico/rcsb-embedding-search), which provides scripts and documentation for training.
+
+
+## Features
+
+- **Residue-level embeddings** computed using the ESM3 protein language model  
+- **Structure-level embeddings** aggregated via a transformer-based aggregator network  
+- **Command-line interface** implemented with Typer for high-throughput inference workflows  
+- **Python API** for interactive embedding computation and integration into analysis pipelines  
+- **High-performance inference** leveraging PyTorch Lightning, with multi-node and multi-GPU support  
+
+---
+
+## Installation
+
+    pip install rcsb-embedding-model
+
+**Requirements:**
+
+- Python ≥ 3.10  
+- ESM ≥ 3.2.0  
+- PyTorch ≥ 2.2.0  
+- Lightning ≥ 2.5.0  
+- Typer ≥ 0.15.0  
+
+---
+
+## Quick Start
+
+### CLI
+
+    # 1. Compute residue embeddings: Calculate residue level embeddings of protein structures using ESM3.
+    inference residue-embedding --src-file data/structures.csv --output-path results/residue_embeddings --src-format mmcif --batch-size 8 --devices auto
+
+    # 2. Compute structure embeddings: Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas data-frame file.
+    inference structure-embedding --src-file results/residue_embeddings.csv --output-path results/structure_embeddings --out-df-id embeddings.pkl --batch-size 4 --devices 0 --devives 1
+
+    # 3. Compute chain embeddings: Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files.
+    inference chain-embedding --src-file results/residue_embeddings.csv --output-path results/chain_embeddings --batch-size 4
+
+### Python API
+
+    from rcsb_embedding_model import RcsbStructureEmbedding
+
+    model = RcsbStructureEmbedding()
+
+    # Compute per-residue embeddings
+    res_emb = model.residue_embedding(
+        src_structure="examples/1abc.cif",
+        src_format="mmcif",
+        chain_id="A"
+    )
+
+    # Aggregate to structure-level embedding
+    struct_emb = model.aggregator_embedding(res_emb)
+
+See the examples directory for complete scripts.
+
+---
+
+## Model Architecture
+
+The embedding model is trained to predict structural similarity by approximating TM-scores using cosine distances between embeddings. It consists of two main components:
+
+- **Protein Language Model (PLM)**: Computes residue-level embeddings from a given 3D structure.
+- **Residue Embedding Aggregator**: A transformer-based neural network that aggregates these residue-level embeddings into a single vector.
+
+![Embedding model architecture](assets/embedding-model-architecture.png)
+
+### **Protein Language Model (PLM)**
+Residue-wise embeddings of protein structures are computed using the [ESM3](https://www.evolutionaryscale.ai/) generative protein language model.
+
+### **Residue Embedding Aggregator**
+The aggregation component consists of six transformer encoder layers, each with a 3,072-neuron feedforward layer and ReLU activations. After processing through these layers, a summation pooling operation is applied, followed by 12 fully connected residual layers that refine the embeddings into a single 1,536-dimensional vector.
+
+---
+
+## Development
+
+    git clone https://github.com/rcsb/rcsb-embedding-model.git
+    cd rcsb-embedding-model
+    pip install -e .
+    pytest
+
+---
+
+## Citation
+
+Segura, J., Bittrich, S., et al. (2024). *Multi-scale structural similarity embedding search across entire proteomes*. bioRxiv. (Preprint: https://www.biorxiv.org/content/10.1101/2024.03.07.XXXXX)
+
+---
+
+## License
+
+This project is licensed under the BSD 3-Clause License. See [LICENSE.md](LICENSE.md) for details.

rcsb_embedding_model-0.0.7/examples/esm_embeddings.py

@@ -0,0 +1,23 @@
+import argparse
+
+from rcsb_embedding_model import RcsbStructureEmbedding
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--file', type=str, required=True)
+    parser.add_argument('--file_format', type=str)
+    parser.add_argument('--chain', type=str)
+    args = parser.parse_args()
+
+    model = RcsbStructureEmbedding()
+    res_embedding = model.residue_embedding(
+        src_structure=args.file,
+        src_format=args.file_format,
+        chain_id=args.chain
+    )
+    structure_embedding = model.aggregator_embedding(
+        res_embedding
+    )
+
+    print(res_embedding.shape, structure_embedding.shape)

{rcsb_embedding_model-0.0.5 → rcsb_embedding_model-0.0.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.5"
+version = "0.0.7"
 authors = [
     { name="Joan Segura", email="joan.segura@rcsb.org" },
 ]
@@ -15,8 +15,11 @@ license = "BSD-3-Clause"
 license-files = ["LICEN[CS]E*"]
 dependencies=[
     "esm >= 3.2.0",
-    "torch >= 2.2.0"
+    "torch >= 2.2.0",
+    "lightning >= 2.5.0",
+    "typer >= 0.15.0"
 ]
+
 [project.urls]
 Homepage = "https://github.com/rcsb/rcsb-embedding-model"
 Issues = "https://github.com/rcsb/rcsb-embedding-model/issues"
@@ -26,4 +29,7 @@ requires = [
     "hatchling >= 1.14.1"
 ]
 
-build-backend = "hatchling.build"
+build-backend = "hatchling.build"
+
+[project.scripts]
+inference = "rcsb_embedding_model.cli.inference:app"

rcsb_embedding_model-0.0.7/src/rcsb_embedding_model/cli/args_utils.py

@@ -0,0 +1,9 @@
+
+
+def arg_devices(devices):
+    if len(devices) == 1:
+        return devices[0] if devices[0] == "auto" else int(devices[0])
+    return [int(x) for x in devices]
+
+
+

rcsb_embedding_model-0.0.7/src/rcsb_embedding_model/cli/inference.py

@@ -0,0 +1,175 @@
+from typing import Annotated, List
+
+import typer
+
+from rcsb_embedding_model.cli.args_utils import arg_devices
+from rcsb_embedding_model.types.api_types import SrcFormat, Accelerator, SrcLocation
+
+app = typer.Typer(
+    add_completion=False
+)
+
+
+@app.command(
+    name="residue-embedding",
+    help="Calculate residue level embeddings of protein structures using ESM3."
+)
+def residue_embedding(
+        src_file: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            resolve_path=True,
+            help='CSV file 3 columns: Structure File Path | Chain Id (asym_i for cif files) | Output file name.'
+        )],
+        output_path: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Output path to store predictions.'
+        )],
+        src_location: Annotated[SrcLocation, typer.Option(
+            help='Source input location.'
+        )] = SrcLocation.local,
+        src_format: Annotated[SrcFormat, typer.Option(
+            help='Structure file format.'
+        )] = SrcFormat.mmcif,
+        batch_size: Annotated[int, typer.Option(
+            help='Number of samples processed together in one iteration.'
+        )] = 1,
+        num_workers: Annotated[int, typer.Option(
+            help='Number of subprocesses to use for data loading.'
+        )] = 0,
+        num_nodes: Annotated[int, typer.Option(
+            help='Number of nodes to use for inference.'
+        )] = 1,
+        accelerator: Annotated[Accelerator, typer.Option(
+            help='Device used for inference.'
+        )] = Accelerator.auto,
+        devices: Annotated[List[str], typer.Option(
+            help='The devices to use. Can be set to a positive number or "auto". Repeat this argument to indicate multiple indices of devices. "auto" for automatic selection based on the chosen accelerator.'
+        )] = tuple(['auto'])
+):
+    from rcsb_embedding_model.inference.esm_inference import predict
+    predict(
+        csv_file=src_file,
+        src_location=src_location,
+        src_format=src_format,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        num_nodes=num_nodes,
+        accelerator=accelerator,
+        devices=arg_devices(devices),
+        out_path=output_path
+    )
+
+
+@app.command(
+    name="structure-embedding",
+    help="Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas data-frame file."
+)
+def structure_embedding(
+        src_file: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            resolve_path=True,
+            help='CSV file 3 columns: Structure File Path | Chain Id (asym_i for cif files) | Output file name.'
+        )],
+        output_path: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Output path to store predictions.'
+        )],
+        out_df_id: Annotated[str, typer.Option(
+            help='File name to store predicted embeddings.'
+        )],
+        src_location: Annotated[SrcLocation, typer.Option(
+            help='Source input location.'
+        )] = SrcLocation.local,
+        src_format: Annotated[SrcFormat, typer.Option(
+            help='Structure file format.'
+        )] = SrcFormat.mmcif,
+        batch_size: Annotated[int, typer.Option(
+            help='Number of samples processed together in one iteration.'
+        )] = 1,
+        num_workers: Annotated[int, typer.Option(
+            help='Number of subprocesses to use for data loading.'
+        )] = 0,
+        num_nodes: Annotated[int, typer.Option(
+            help='Number of nodes to use for inference.'
+        )] = 1,
+        accelerator: Annotated[Accelerator, typer.Option(
+            help='Device used for inference.'
+        )] = Accelerator.auto,
+        devices: Annotated[List[str], typer.Option(
+            help='The devices to use. Can be set to a positive number or "auto". Repeat this argument to indicate multiple indices of devices. "auto" for automatic selection based on the chosen accelerator.'
+        )] = tuple(['auto'])
+):
+    from rcsb_embedding_model.inference.structure_inference import predict
+    predict(
+        csv_file=src_file,
+        src_location=src_location,
+        src_format=src_format,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        num_nodes=num_nodes,
+        accelerator=accelerator,
+        devices=arg_devices(devices),
+        out_path=output_path,
+        out_df_id=out_df_id
+    )
+
+
+@app.command(
+    name="chain-embedding",
+    help="Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files."
+)
+def chain_embedding(
+        src_file: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            resolve_path=True,
+            help='CSV file 2 columns: Residue Embedding Tensor File | Output file name.'
+        )],
+        output_path: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Output path to store predictions.'
+        )],
+        batch_size: Annotated[int, typer.Option(
+            help='Number of samples processed together in one iteration.'
+        )] = 1,
+        num_workers: Annotated[int, typer.Option(
+            help='Number of subprocesses to use for data loading.'
+        )] = 0,
+        num_nodes: Annotated[int, typer.Option(
+            help='Number of nodes to use for inference.'
+        )] = 1,
+        accelerator: Annotated[Accelerator, typer.Option(
+            help='Device used for inference.'
+        )] = Accelerator.auto,
+        devices: Annotated[List[str], typer.Option(
+            help='The devices to use. Can be set to a positive number or "auto". Repeat this argument to indicate multiple indices of devices. "auto" for automatic selection based on the chosen accelerator.'
+        )] = tuple(['auto'])
+):
+    from rcsb_embedding_model.inference.chain_inference import predict
+    predict(
+        csv_file=src_file,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        num_nodes=num_nodes,
+        accelerator=accelerator,
+        devices=arg_devices(devices),
+        out_path=output_path
+    )
+
+
+if __name__ == "__main__":
+    app()

rcsb_embedding_model-0.0.7/src/rcsb_embedding_model/dataset/esm_prot_from_csv.py

@@ -0,0 +1,90 @@
+import argparse
+
+import torch
+from biotite.structure import chain_iter
+from esm.models.esm3 import ESM3
+from esm.sdk.api import ESMProtein, SamplingConfig
+from esm.utils.constants.models import ESM3_OPEN_SMALL
+from esm.utils.structure.protein_chain import ProteinChain
+from torch.utils.data import Dataset, DataLoader
+import pandas as pd
+
+from rcsb_embedding_model.types.api_types import SrcFormat, SrcLocation
+from rcsb_embedding_model.utils.data import stringio_from_url
+from rcsb_embedding_model.utils.structure_parser import get_structure_from_src
+
+
+class EsmProtFromCsv(Dataset):
+
+    MIN_RES = 10
+    STREAM_ATTR = 'stream'
+    CH_ATTR = 'chain_id'
+    NAME_ATTR = 'name'
+
+    COLUMNS = [STREAM_ATTR, CH_ATTR, NAME_ATTR]
+
+    def __init__(
+        self,
+        csv_file,
+        src_location=SrcLocation.local,
+        src_format=SrcFormat.mmcif,
+    ):
+        super().__init__()
+        self.src_location = src_location
+        self.src_format = src_format
+        self.data = pd.DataFrame()
+        self.__load_stream(csv_file)
+
+    def __load_stream(self, stream_list):
+        self.data = pd.read_csv(
+            stream_list,
+            header=None,
+            index_col=None,
+            names=EsmProtFromCsv.COLUMNS
+        )
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        src_structure = self.data.loc[idx, EsmProtFromCsv.STREAM_ATTR]
+        chain_id = self.data.loc[idx, EsmProtFromCsv.CH_ATTR]
+        name = self.data.loc[idx, EsmProtFromCsv.NAME_ATTR]
+        structure = get_structure_from_src(
+            src_structure=src_structure if self.src_location == SrcLocation.local else stringio_from_url(src_structure),
+            src_format=self.src_format,
+            chain_id=chain_id
+        )
+        for atom_ch in chain_iter(structure):
+            protein_chain = ProteinChain.from_atomarray(atom_ch)
+            return ESMProtein.from_protein_chain(protein_chain), name
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--file_list', type=argparse.FileType('r'), required=True)
+    args = parser.parse_args()
+
+    dataset = EsmProtFromCsv(
+        args.file_list
+    )
+
+    esm3 = ESM3.from_pretrained(
+        ESM3_OPEN_SMALL,
+        torch.device("cpu")
+    )
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=2,
+        collate_fn=lambda _: _
+    )
+
+    for _batch in dataloader:
+        for esm_prot, name in _batch:
+            protein_tensor = esm3.encode(esm_prot)
+            embeddings = esm3.forward_and_sample(
+                protein_tensor, SamplingConfig(return_per_residue_embeddings=True)
+            ).per_residue_embedding
+            print(name, embeddings.shape)

rcsb_embedding_model-0.0.7/src/rcsb_embedding_model/dataset/residue_embedding_from_csv.py

@@ -0,0 +1,32 @@
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+
+class ResidueEmbeddingFromCSV(Dataset):
+
+    STREAM_ATTR = 'stream'
+    NAME_ATTR = 'name'
+
+    COLUMNS = [STREAM_ATTR, NAME_ATTR]
+
+    def __init__(self, csv_file):
+        super().__init__()
+        self.data = pd.DataFrame()
+        self.__load_stream(csv_file)
+
+    def __load_stream(self, csv_file):
+        self.data = pd.read_csv(
+            csv_file,
+            header=None,
+            index_col=None,
+            names=ResidueEmbeddingFromCSV.COLUMNS
+        )
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        embedding_src = self.data.loc[idx, ResidueEmbeddingFromCSV.STREAM_ATTR]
+        name = self.data.loc[idx, ResidueEmbeddingFromCSV.NAME_ATTR]
+        return torch.load(embedding_src, map_location=torch.device('cpu')), name

rcsb_embedding_model-0.0.7/src/rcsb_embedding_model/inference/chain_inference.py

@@ -0,0 +1,50 @@
+from torch.utils.data import DataLoader
+from lightning import Trainer
+from typer import FileText
+
+from rcsb_embedding_model.dataset.residue_embedding_from_csv import ResidueEmbeddingFromCSV
+from rcsb_embedding_model.modules.chain_module import ChainModule
+from rcsb_embedding_model.types.api_types import Accelerator, Devices, OptionalPath
+from rcsb_embedding_model.utils.data import collate_seq_embeddings
+from rcsb_embedding_model.writer.batch_writer import CsvBatchWriter
+
+
+def predict(
+        csv_file: FileText,
+        batch_size: int = 1,
+        num_workers: int = 0,
+        num_nodes: int = 1,
+        accelerator: Accelerator = Accelerator.auto,
+        devices: Devices = 'auto',
+        out_path: OptionalPath = None
+):
+    inference_set = ResidueEmbeddingFromCSV(
+        csv_file=csv_file
+    )
+
+    inference_dataloader = DataLoader(
+        dataset=inference_set,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        collate_fn=lambda emb: (
+            collate_seq_embeddings([x for x, z in emb]),
+            tuple([z for x, z in emb])
+        )
+    )
+
+    module = ChainModule()
+
+    inference_writer = CsvBatchWriter(out_path) if out_path is not None else None
+    trainer = Trainer(
+        callbacks=[inference_writer] if inference_writer is not None else None,
+        num_nodes=num_nodes,
+        accelerator=accelerator,
+        devices=devices
+    )
+
+    prediction = trainer.predict(
+        module,
+        inference_dataloader
+    )
+
+    return prediction