rcsb-embedding-model 0.0.44__tar.gz

rcsb_embedding_model-0.0.44/.dockerignore

@@ -0,0 +1,36 @@
+# Only list files that could be under version control or could be created in CI/CD.
+
+# Note: Patterns are absolute, like `.prettierignore` but unlike `.gitignore`.
+
+# Gitignored files
+**/.*
+**/[~#$]*
+**/*[~#$]
+
+# Gitignored directories
+node_modules/
+
+# Directories
+/dist/
+/tests/
+/assets/
+/examples/
+
+# Files in the root directory
+/*.md
+/*.txt
+/CITATION.cff
+/compose.yaml
+/justfile
+/mkdocs.yaml
+
+# Keep README.md (needed for build)
+!/README.md
+
+# Keep .dockerignore and .gitignore
+!/.dockerignore
+!/.gitignore
+
+# Keep legal files
+!/LICENSE.*
+!/NOTICE.txt

rcsb_embedding_model-0.0.44/.github/workflows/_workflow-docker.yaml

@@ -0,0 +1,17 @@
+# Docker build and push workflow
+
+name: Run CI/CD Docker Workflow
+
+on:
+  workflow_call:
+
+jobs:
+  run-workflow:
+    if: github.event_name == 'release'
+    name: "Run automated docker workflow"
+    uses: rcsb/devops-cicd-github-actions/.github/workflows/workflow-docker.yaml@master
+    with:
+      dockerfile_location: "Dockerfile"  # The location of the Dockerfile relative to the root of the repository. Defaults to "Dockerfile".
+      repo_project: "rcsb"  # REQUIRED. The name of the project or organization in the remote Docker image repository.
+      docker_image_name: "rcsb-embedding-model"  # REQUIRED. The name of the Docker image to create.
+      docker_build_context: "."  # The path location of the docker build context, relative to the project root. Defaults to the project root.

rcsb_embedding_model-0.0.44/.github/workflows/publish.yaml

@@ -0,0 +1,91 @@
+name: CI Pipeline
+
+on:
+  push:
+    branches: [master]
+  pull_request:
+    branches: [master]
+  release:
+    types: [published]
+
+jobs:
+  hatch-test:
+    name: Test on Python ${{ matrix.python-version }}
+    runs-on: ["self-hosted", "buildchain"]
+    timeout-minutes: 20
+    strategy:
+      matrix:
+        python-version: ["3.11"]
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Install build dependencies
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y build-essential pkg-config libzstd-dev
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install Hatch and HaggingFace
+        run: pip install hatch huggingface_hub[cli]
+
+      - name: Run tests
+        run: hatch test
+
+  hatch-build:
+    name: Build to PyPI
+    needs: hatch-test
+    runs-on: ubuntu-latest
+    if: github.event_name == 'release'
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.10"
+
+      - name: Install Hatch
+        run: pip install hatch
+
+      - name: Build distribution
+        run: hatch build
+
+      - name: Store the distribution packages
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+
+  publish-to-pypi:
+    name: >-
+      Publish Python 🐍 distribution 📦 to PyPI
+    if: github.event_name == 'release'
+    needs:
+      - hatch-build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/rcsb-embedding-model
+    permissions:
+      id-token: write
+
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+      - name: Publish distribution 📦 to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+
+  push-image:
+    needs:
+      - hatch-test
+    name: Push image to harbor
+    uses: ./.github/workflows/_workflow-docker.yaml

rcsb_embedding_model-0.0.44/.gitignore

@@ -0,0 +1,6 @@
+/.idea
+/rcsb-embedding-model.iml
+/dist/
+/.pypi.rc
+__pycache__
+/tests/resources/tmp

rcsb_embedding_model-0.0.44/Dockerfile

@@ -0,0 +1,8 @@
+FROM python:3.11
+
+WORKDIR /app
+COPY . /app/
+
+RUN pip install --no-cache-dir -e .
+
+ENTRYPOINT ["inference"]

rcsb_embedding_model-0.0.44/LICENSE.md

@@ -0,0 +1,4 @@
+# Cambrian Non-Commercial License Agreement
+
+This project is licensed under the EvolutionaryScale Cambrian Non-Commercial License Agreement.
+See: https://www.evolutionaryscale.ai/policies/cambrian-non-commercial-license-agreement

rcsb_embedding_model-0.0.44/PKG-INFO

@@ -0,0 +1,136 @@
+Metadata-Version: 2.4
+Name: rcsb-embedding-model
+Version: 0.0.44
+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: # Cambrian Non-Commercial License Agreement
+        
+        This project is licensed under the EvolutionaryScale Cambrian Non-Commercial License Agreement.
+        See: https://www.evolutionaryscale.ai/policies/cambrian-non-commercial-license-agreement
+License-File: LICENSE.md
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Requires-Python: >=3.11
+Requires-Dist: biotite>=1.5.0
+Requires-Dist: esm>=3.2.0
+Requires-Dist: hf-xet>=1.1.10
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: huggingface-hub>=0.30.2
+Requires-Dist: importlib-metadata>=8.7.0
+Requires-Dist: lightning>=2.5.0
+Requires-Dist: typer>=0.15.0
+Description-Content-Type: text/markdown
+
+# RCSB Embedding Model
+
+**Version** 0.0.44
+
+
+## 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
+- Lightning ≥ 2.5.0 
+- Typer ≥ 0.15.0
+
+---
+
+## Quick Start
+
+### CLI
+
+    # 1. Compute residue embeddings: Calculate residue level embeddings of protein structures using ESM3. Predictions are stored as torch tensor files.
+    inference residue-embedding --src-file data/structures.csv --output-path results/residue_embeddings --structure-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 DataFrame file.
+    inference structure-embedding --src-file data/structures.csv --output-path results/residue_embeddings --out-df-name df-res-embeddings --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. Predictions a re stored as csv files.
+    inference chain-embedding --src-file data/structures.csv --output-path results/chain_embeddings --batch-size 4
+
+    # 4. Compute assembly embeddings: Calculate assembly embeddings from residue level embeddings stored as torch tensor files. Predictions are stored as csv files.
+    inference assembly-embedding --src-file data/structures.csv --res-embedding-location results/residue_embeddings --output-path results/assembly_embeddings
+
+### 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 and tests directories for more use cases.
+
+---
+
+## 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/2025.02.28.640875v1)
+
+---
+
+## License
+
+This project uses the EvolutionaryScale ESM-3 model and is distributed under the
+[Cambrian Non-Commercial License Agreement](https://www.evolutionaryscale.ai/policies/cambrian-non-commercial-license-agreement).

rcsb_embedding_model-0.0.44/README.md

@@ -0,0 +1,111 @@
+# RCSB Embedding Model
+
+**Version** 0.0.44
+
+
+## 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
+- Lightning ≥ 2.5.0 
+- Typer ≥ 0.15.0
+
+---
+
+## Quick Start
+
+### CLI
+
+    # 1. Compute residue embeddings: Calculate residue level embeddings of protein structures using ESM3. Predictions are stored as torch tensor files.
+    inference residue-embedding --src-file data/structures.csv --output-path results/residue_embeddings --structure-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 DataFrame file.
+    inference structure-embedding --src-file data/structures.csv --output-path results/residue_embeddings --out-df-name df-res-embeddings --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. Predictions a re stored as csv files.
+    inference chain-embedding --src-file data/structures.csv --output-path results/chain_embeddings --batch-size 4
+
+    # 4. Compute assembly embeddings: Calculate assembly embeddings from residue level embeddings stored as torch tensor files. Predictions are stored as csv files.
+    inference assembly-embedding --src-file data/structures.csv --res-embedding-location results/residue_embeddings --output-path results/assembly_embeddings
+
+### 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 and tests directories for more use cases.
+
+---
+
+## 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/2025.02.28.640875v1)
+
+---
+
+## License
+
+This project uses the EvolutionaryScale ESM-3 model and is distributed under the
+[Cambrian Non-Commercial License Agreement](https://www.evolutionaryscale.ai/policies/cambrian-non-commercial-license-agreement).

rcsb_embedding_model-0.0.44/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.44/pyproject.toml

@@ -0,0 +1,39 @@
+[project]
+name = "rcsb-embedding-model"
+version = "0.0.44"
+authors = [
+    { name="Joan Segura", email="joan.segura@rcsb.org" },
+]
+description = "Protein Embedding Model for Structure Search"
+readme = "README.md"
+requires-python = ">=3.11"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent",
+]
+license = {file = "LICENSE.md"}
+license-files = ["LICEN[CS]E*"]
+dependencies=[
+    "importlib-metadata >= 8.7.0",
+    "esm >= 3.2.0",
+    "biotite >= 1.5.0",
+    "lightning >= 2.5.0",
+    "typer >= 0.15.0",
+    "hf-xet >= 1.1.10",
+    "huggingface-hub >= 0.30.2",
+    "httpx >= 0.28.1"
+]
+
+[project.urls]
+Homepage = "https://github.com/rcsb/rcsb-embedding-model"
+Issues = "https://github.com/rcsb/rcsb-embedding-model/issues"
+
+[build-system]
+requires = [
+    "hatchling >= 1.14.1"
+]
+
+build-backend = "hatchling.build"
+
+[project.scripts]
+inference = "rcsb_embedding_model.cli.inference:app"

rcsb_embedding_model-0.0.44/src/rcsb_embedding_model/__init__.py

@@ -0,0 +1,9 @@
+from importlib_metadata import version, PackageNotFoundError
+from rcsb_embedding_model.rcsb_structure_embedding import RcsbStructureEmbedding
+
+try:
+    __version__ = version("rcsb-embedding-model")
+except PackageNotFoundError:
+    __version__ = "0.0.0"
+
+__all__ = ["RcsbStructureEmbedding", "__version__"]

rcsb_embedding_model-0.0.44/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]
+
+
+