rcsb-embedding-model 0.0.6__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.6 → rcsb_embedding_model-0.0.7}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.6"
+version = "0.0.7"
 authors = [
     { name="Joan Segura", email="joan.segura@rcsb.org" },
 ]

{rcsb_embedding_model-0.0.6 → rcsb_embedding_model-0.0.7}/src/rcsb_embedding_model/cli/args_utils.py

@@ -1,6 +1,4 @@
 
-from enum import Enum
-
 
 def arg_devices(devices):
     if len(devices) == 1:

{rcsb_embedding_model-0.0.6 → rcsb_embedding_model-0.0.7}/src/rcsb_embedding_model/cli/inference.py

@@ -5,17 +5,29 @@ 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()
+app = typer.Typer(
+    add_completion=False
+)
 
 
-@app.command(name="residue-embedding")
+@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 | Chain Id (asym_i for cif files) | Output file name.'
+            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.'
@@ -37,14 +49,7 @@ def residue_embedding(
         )] = 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']),
-        output_path: Annotated[typer.FileText, typer.Option(
-            exists=True,
-            file_okay=False,
-            dir_okay=True,
-            resolve_path=True,
-            help='Output path to store predictions.'
-        )] = None
+        )] = tuple(['auto'])
 ):
     from rcsb_embedding_model.inference.esm_inference import predict
     predict(
@@ -60,14 +65,27 @@ def residue_embedding(
     )
 
 
-@app.command(name="structure-embedding")
+@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 | Chain Id (asym_i for cif files) | Output file name.'
+            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.'
@@ -89,19 +107,27 @@ def structure_embedding(
         )] = 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']),
-        output_path: Annotated[typer.FileText, typer.Option(
-            exists=True,
-            file_okay=False,
-            dir_okay=True,
-            resolve_path=True,
-            help='Output path to store predictions.'
-        )] = None
+        )] = tuple(['auto'])
 ):
-    pass
+    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")
+@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,
@@ -110,6 +136,13 @@ def chain_embedding(
             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,
@@ -124,14 +157,7 @@ def chain_embedding(
         )] = 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']),
-        output_path: Annotated[typer.FileText, typer.Option(
-            exists=True,
-            file_okay=False,
-            dir_okay=True,
-            resolve_path=True,
-            help='Output path to store predictions.'
-        )] = None
+        )] = tuple(['auto'])
 ):
     from rcsb_embedding_model.inference.chain_inference import predict
     predict(

{rcsb_embedding_model-0.0.6 → rcsb_embedding_model-0.0.7}/src/rcsb_embedding_model/dataset/esm_prot_from_csv.py

@@ -1,5 +1,4 @@
 import argparse
-import os
 
 import torch
 from biotite.structure import chain_iter
@@ -48,11 +47,11 @@ class EsmProtFromCsv(Dataset):
         return len(self.data)
 
     def __getitem__(self, idx):
-        structure_src = self.data.loc[idx, EsmProtFromCsv.STREAM_ATTR]
+        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(
-            structure_src if self.src_location == SrcLocation.local else stringio_from_url(structure_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
         )

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

@@ -0,0 +1,51 @@
+from torch.utils.data import DataLoader
+from lightning import Trainer
+from typer import FileText
+
+from rcsb_embedding_model.dataset.esm_prot_from_csv import EsmProtFromCsv
+from rcsb_embedding_model.modules.esm_module import EsmModule
+from rcsb_embedding_model.types.api_types import SrcFormat, Accelerator, Devices, OptionalPath, SrcLocation
+from rcsb_embedding_model.writer.batch_writer import DataFrameStorage
+
+
+def predict(
+        csv_file: FileText,
+        src_location: SrcLocation = SrcLocation.local,
+        src_format: SrcFormat = SrcFormat.mmcif,
+        batch_size: int = 1,
+        num_workers: int = 0,
+        num_nodes: int = 1,
+        accelerator: Accelerator = Accelerator.auto,
+        devices: Devices = 'auto',
+        out_path: OptionalPath = None,
+        out_df_id: str = None
+):
+
+    inference_set = EsmProtFromCsv(
+        csv_file=csv_file,
+        src_location=src_location,
+        src_format=src_format
+    )
+
+    inference_dataloader = DataLoader(
+        dataset=inference_set,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        collate_fn=lambda _: _
+    )
+
+    module = EsmModule()
+    inference_writer = DataFrameStorage(out_path, out_df_id) if out_path is not None and out_df_id 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

rcsb_embedding_model-0.0.7/src/rcsb_embedding_model/modules/structure_module.py

@@ -0,0 +1,27 @@
+from esm.sdk.api import SamplingConfig
+from lightning import LightningModule
+
+from rcsb_embedding_model.utils.data import collate_seq_embeddings
+from rcsb_embedding_model.utils.model import get_residue_model, get_aggregator_model
+
+
+class StructureModule(LightningModule):
+
+    def __init__(
+            self
+    ):
+        super().__init__()
+        self.esm3 = get_residue_model(self.device)
+        self.aggregator = get_aggregator_model(device=self.device)
+
+    def predict_step(self, prot_batch, batch_idx):
+        prot_embeddings = []
+        prot_names = []
+        for esm_prot, name in prot_batch:
+            embeddings = self.esm3.forward_and_sample(
+                self.esm3.encode(esm_prot), SamplingConfig(return_per_residue_embeddings=True)
+            ).per_residue_embedding
+            prot_embeddings.append(embeddings)
+            prot_names.append(name)
+        res_batch_embedding, res_batch_mask = collate_seq_embeddings(prot_embeddings)
+        return self.aggregator(res_batch_embedding, res_batch_mask), tuple(prot_names)

{rcsb_embedding_model-0.0.6 → rcsb_embedding_model-0.0.7}/src/rcsb_embedding_model/rcsb_structure_embedding.py

@@ -2,7 +2,6 @@ import torch
 from biotite.structure import get_residues, chain_iter, filter_amino_acids
 from esm.sdk.api import ESMProtein, SamplingConfig
 from esm.utils.structure.protein_chain import ProteinChain
-from huggingface_hub import hf_hub_download
 
 from rcsb_embedding_model.types.api_types import StreamSrc, SrcFormat
 from rcsb_embedding_model.utils.model import get_aggregator_model, get_residue_model
@@ -42,23 +41,23 @@ class RcsbStructureEmbedding:
 
     def structure_embedding(
             self,
-            structure_src: StreamSrc,
+            src_structure: StreamSrc,
             src_format: SrcFormat = SrcFormat.mmcif,
             chain_id: str = None,
             assembly_id: str = None
     ):
-        res_embedding = self.residue_embedding(structure_src, src_format, chain_id, assembly_id)
+        res_embedding = self.residue_embedding(src_structure, src_format, chain_id, assembly_id)
         return self.aggregator_embedding(res_embedding)
 
     def residue_embedding(
             self,
-            structure_src: StreamSrc,
+            src_structure: StreamSrc,
             src_format: SrcFormat = SrcFormat.mmcif,
             chain_id: str = None,
             assembly_id: str = None
     ):
         self.__check_residue_embedding()
-        structure = get_structure_from_src(structure_src, src_format, chain_id, assembly_id)
+        structure = get_structure_from_src(src_structure, src_format, chain_id, assembly_id)
         embedding_ch = []
         for atom_ch in chain_iter(structure):
             atom_res = atom_ch[filter_amino_acids(atom_ch)]

{rcsb_embedding_model-0.0.6 → rcsb_embedding_model-0.0.7}/src/rcsb_embedding_model/utils/structure_parser.py

@@ -4,19 +4,19 @@ from biotite.structure.io.pdbx import CIFFile, get_structure, get_assembly, Bina
 
 
 def get_structure_from_src(
-        structure_src,
+        src_structure,
         src_format="mmcif",
         chain_id=None,
         assembly_id=None
 ):
     if src_format == "pdb":
-        pdb_file = PDBFile.read(structure_src)
+        pdb_file = PDBFile.read(src_structure)
         structure = __get_pdb_structure(pdb_file, assembly_id)
     elif src_format == "mmcif":
-        cif_file = CIFFile.read(structure_src)
+        cif_file = CIFFile.read(src_structure)
         structure = __get_structure(cif_file, assembly_id)
     elif src_format == "binarycif":
-        cif_file = BinaryCIFFile.read(structure_src)
+        cif_file = BinaryCIFFile.read(src_structure)
         structure = __get_structure(cif_file, assembly_id)
     else:
         raise RuntimeError(f"Unknown file format {src_format}")

{rcsb_embedding_model-0.0.6 → rcsb_embedding_model-0.0.7}/tests/test_embedding_model.py

@@ -13,7 +13,7 @@ class TestEmbeddingModel(unittest.TestCase):
 
         model = RcsbStructureEmbedding()
         res_embedding = model.residue_embedding(
-            f"{self.__test_path}/resources/1acb.cif",
+            src_structure=f"{self.__test_path}/resources/1acb.cif",
             src_format=SrcFormat.mmcif,
             chain_id='A'
         )
@@ -33,7 +33,7 @@ class TestEmbeddingModel(unittest.TestCase):
 
         model = RcsbStructureEmbedding()
         res_embedding = model.residue_embedding(
-            f"{self.__test_path}/resources/1acb.cif",
+            src_structure=f"{self.__test_path}/resources/1acb.cif",
             src_format=SrcFormat.mmcif,
             chain_id='A'
         )

rcsb_embedding_model-0.0.6/PKG-INFO

@@ -1,117 +0,0 @@
-Metadata-Version: 2.4
-Name: rcsb-embedding-model
-Version: 0.0.6
-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: A Deep Learning Approach for 3D Structure Embeddings
-
-## Overview
-RCSB Embedding Model is a PyTorch-based neural network that transforms macromolecular 3D structures into vector embeddings.
-
-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.
-
----
-
-## Embedding Model
-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.
-
----
-
-## How to Use the Model
-This repository provides the tools to compute embeddings for 3D macromolecular structure data.
-
-### **Installation**
-`pip install rcsb-embedding-model`
-
-### **Requirements**
-Ensure you have the following dependencies installed:
-- `python >= 3.10`
-- `esm`
-- `torch`
-
-### **Generating Residue Embeddings**
-ESM3 embeddings for the 3D structures can be calculated as:
-
-```python
-from rcsb_embedding_model import RcsbStructureEmbedding
-
-mmcif_file = "<path_to_file>/<name>.cif"
-model = RcsbStructureEmbedding()
-res_embedding = model.residue_embedding(
-    structure_src=mmcif_file,
-    format="mmcif",
-    chain_id='A'
-)
-```
-
-### **Generating Protein Structure Embeddings**
-Protein 3D structure embedding can be calculated as:
-
-```python
-from rcsb_embedding_model import RcsbStructureEmbedding
-
-mmcif_file = "<path_to_file>/<name>.cif"
-model = RcsbStructureEmbedding()
-res_embedding = model.residue_embedding(
-    structure_src=mmcif_file,
-    format="mmcif",
-    chain_id='A'
-)
-structure_embedding = model.aggregator_embedding(
-    res_embedding
-)
-```
-
-### **Pretrained Model**
-You can download a pretrained Residue Embedding Aggregator model from [Hugging Face](https://huggingface.co/jseguramora/rcsb-embedding-model/resolve/main/rcsb-embedding-model.pt).
-
----
-
-## Questions & Issues
-For any questions or comments, please open an issue on this repository.
-
----
-
-## License
-This software is released under the BSD 3-Clause License. See the full license text below.
-
-### BSD 3-Clause License
-
-Copyright (c) 2024, RCSB Protein Data Bank, UC San Diego
-
-Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this list of conditions, and the following disclaimer.
-2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions, and the following disclaimer in the documentation and/or other materials provided with the distribution.
-3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-

rcsb_embedding_model-0.0.6/README.md

@@ -1,99 +0,0 @@
-# RCSB Embedding Model: A Deep Learning Approach for 3D Structure Embeddings
-
-## Overview
-RCSB Embedding Model is a PyTorch-based neural network that transforms macromolecular 3D structures into vector embeddings.
-
-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.
-
----
-
-## Embedding Model
-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.
-
----
-
-## How to Use the Model
-This repository provides the tools to compute embeddings for 3D macromolecular structure data.
-
-### **Installation**
-`pip install rcsb-embedding-model`
-
-### **Requirements**
-Ensure you have the following dependencies installed:
-- `python >= 3.10`
-- `esm`
-- `torch`
-
-### **Generating Residue Embeddings**
-ESM3 embeddings for the 3D structures can be calculated as:
-
-```python
-from rcsb_embedding_model import RcsbStructureEmbedding
-
-mmcif_file = "<path_to_file>/<name>.cif"
-model = RcsbStructureEmbedding()
-res_embedding = model.residue_embedding(
-    structure_src=mmcif_file,
-    format="mmcif",
-    chain_id='A'
-)
-```
-
-### **Generating Protein Structure Embeddings**
-Protein 3D structure embedding can be calculated as:
-
-```python
-from rcsb_embedding_model import RcsbStructureEmbedding
-
-mmcif_file = "<path_to_file>/<name>.cif"
-model = RcsbStructureEmbedding()
-res_embedding = model.residue_embedding(
-    structure_src=mmcif_file,
-    format="mmcif",
-    chain_id='A'
-)
-structure_embedding = model.aggregator_embedding(
-    res_embedding
-)
-```
-
-### **Pretrained Model**
-You can download a pretrained Residue Embedding Aggregator model from [Hugging Face](https://huggingface.co/jseguramora/rcsb-embedding-model/resolve/main/rcsb-embedding-model.pt).
-
----
-
-## Questions & Issues
-For any questions or comments, please open an issue on this repository.
-
----
-
-## License
-This software is released under the BSD 3-Clause License. See the full license text below.
-
-### BSD 3-Clause License
-
-Copyright (c) 2024, RCSB Protein Data Bank, UC San Diego
-
-Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this list of conditions, and the following disclaimer.
-2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions, and the following disclaimer in the documentation and/or other materials provided with the distribution.
-3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
-

rcsb_embedding_model-0.0.6/examples/esm_embeddings.py

@@ -1,77 +0,0 @@
-import argparse
-
-import torch
-from biotite.structure import chain_iter, get_residues, filter_amino_acids
-from biotite.structure.io.pdb import PDBFile
-from biotite.structure.io.pdbx import CIFFile, get_structure, BinaryCIFFile
-from esm.models.esm3 import ESM3, ESM3_OPEN_SMALL
-from esm.sdk.api import ESMProtein, SamplingConfig
-from esm.utils.structure.protein_chain import ProteinChain
-
-
-def get_structure_from_file(file_name, format="pdb", chain_id=None):
-    try:
-        if format == "pdb":
-            structure = PDBFile.read(file_name).get_structure(
-                model=1
-            )
-        elif format == "mmcif":
-            cif_file = CIFFile.read(file_name)
-            structure = get_structure(
-                cif_file,
-                model=1,
-                use_author_fields=False
-            )
-        elif format == "binarycif":
-            cif_file = BinaryCIFFile.read(file_name)
-            structure = get_structure(
-                cif_file,
-                model=1,
-                use_author_fields=False
-            )
-
-        if chain_id:
-            structure = structure[structure.chain_id == chain_id]
-        return structure
-    except:
-        return None
-
-
-esm3_model = ESM3.from_pretrained(ESM3_OPEN_SMALL)
-
-
-def compute_embeddings(structure):
-    embedding_ch = []
-    for atom_ch in chain_iter(structure):
-        atom_res = atom_ch[filter_amino_acids(atom_ch)]
-        if len(atom_res) == 0 or len(get_residues(atom_res)[0]) < 10:
-            continue
-        protein_chain = ProteinChain.from_atomarray(atom_ch)
-        protein = ESMProtein.from_protein_chain(protein_chain)
-        protein_tensor = esm3_model.encode(protein)
-        embedding_ch.append( esm3_model.forward_and_sample(
-            protein_tensor, SamplingConfig(return_per_residue_embeddings=True)
-        ).per_residue_embedding)
-    return torch.cat(
-        embedding_ch,
-        dim=0
-    )
-
-
-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()
-
-    structure = get_structure_from_file(
-        args.file,
-        "pdb" if not args.file_format else args.file_format,
-        chain_id=args.chain
-    )
-
-    embeddings = compute_embeddings(structure)
-
-    print(embeddings.shape)