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

{rcsb_embedding_model-0.0.5 → rcsb_embedding_model-0.0.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.5
+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
@@ -11,7 +11,9 @@ 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

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

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.5"
+version = "0.0.6"
 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.6/src/rcsb_embedding_model/cli/args_utils.py

@@ -0,0 +1,11 @@
+
+from enum import Enum
+
+
+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.6/src/rcsb_embedding_model/cli/inference.py

@@ -0,0 +1,149 @@
+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()
+
+
+@app.command(name="residue-embedding")
+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.'
+        )],
+        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']),
+        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
+):
+    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")
+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.'
+        )],
+        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']),
+        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
+):
+    pass
+
+
+@app.command(name="chain-embedding")
+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.'
+        )],
+        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']),
+        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
+):
+    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.6/src/rcsb_embedding_model/dataset/esm_prot_from_csv.py

@@ -0,0 +1,91 @@
+import argparse
+import os
+
+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):
+        structure_src = 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_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.6/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.6/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

rcsb_embedding_model-0.0.6/src/rcsb_embedding_model/inference/esm_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.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 TensorBatchWriter
+
+
+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
+):
+
+    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 = TensorBatchWriter(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

rcsb_embedding_model-0.0.6/src/rcsb_embedding_model/modules/chain_module.py

@@ -0,0 +1,16 @@
+from lightning import LightningModule
+
+from rcsb_embedding_model.utils.model import get_aggregator_model
+
+
+class ChainModule(LightningModule):
+
+    def __init__(
+            self
+    ):
+        super().__init__()
+        self.model = get_aggregator_model(device=self.device)
+
+    def predict_step(self, batch, batch_idx):
+        (x, x_mask), dom_id = batch
+        return self.model(x, x_mask), dom_id

rcsb_embedding_model-0.0.6/src/rcsb_embedding_model/modules/esm_module.py

@@ -0,0 +1,24 @@
+from esm.sdk.api import SamplingConfig
+from lightning import LightningModule
+
+from rcsb_embedding_model.utils.model import get_residue_model
+
+
+class EsmModule(LightningModule):
+
+    def __init__(
+            self
+    ):
+        super().__init__()
+        self.esm3 = get_residue_model(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)
+        return tuple(prot_embeddings), tuple(prot_names)

rcsb_embedding_model-0.0.6/src/rcsb_embedding_model/rcsb_structure_embedding.py

@@ -0,0 +1,128 @@
+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
+from rcsb_embedding_model.utils.structure_parser import get_structure_from_src
+
+
+class RcsbStructureEmbedding:
+
+    MIN_RES = 10
+
+    def __init__(self):
+        self.__residue_embedding = None
+        self.__aggregator_embedding = None
+
+    def load_models(
+            self,
+            device: torch.device = None
+    ):
+        self.load_residue_embedding(device)
+        self.load_aggregator_embedding(device)
+
+    def load_residue_embedding(
+            self,
+            device: torch.device = None
+    ):
+        if not device:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.__residue_embedding = _load_res_model(device)
+
+    def load_aggregator_embedding(
+            self,
+            device: torch.device = None
+    ):
+        if not device:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.__aggregator_embedding = _load_model(device)
+
+    def structure_embedding(
+            self,
+            structure_src: 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)
+        return self.aggregator_embedding(res_embedding)
+
+    def residue_embedding(
+            self,
+            structure_src: 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)
+        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]) < RcsbStructureEmbedding.MIN_RES:
+                continue
+            protein_chain = ProteinChain.from_atomarray(atom_ch)
+            protein = ESMProtein.from_protein_chain(protein_chain)
+            protein_tensor = self.__residue_embedding.encode(protein)
+            embedding_ch.append(self.__residue_embedding.forward_and_sample(
+                protein_tensor, SamplingConfig(return_per_residue_embeddings=True)
+            ).per_residue_embedding)
+        return torch.cat(
+            embedding_ch,
+            dim=0
+        )
+
+    def sequence_embedding(
+            self,
+            sequence: str
+    ):
+        self.__check_residue_embedding()
+
+        if sequence.startswith(">"):
+            sequence = "".join(line.strip() for line in sequence.splitlines() if not line.startswith(">"))
+
+        if len(sequence) < RcsbStructureEmbedding.MIN_RES:
+            raise ValueError(f"Sequence too short for embedding (min {RcsbStructureEmbedding.MIN_RES} residues)")
+
+        protein = ESMProtein(sequence=sequence)
+        protein_tensor = self.__residue_embedding.encode(protein)
+
+        result = self.__residue_embedding.forward_and_sample(
+            protein_tensor,
+            SamplingConfig(return_per_residue_embeddings=True)
+        )
+
+        return result.per_residue_embedding
+
+    def aggregator_embedding(
+            self,
+            residue_embedding: torch.Tensor
+    ):
+        self.__check_aggregator_embedding()
+        return self.__aggregator_embedding(residue_embedding)
+
+    def __check_residue_embedding(self):
+        if self.__residue_embedding is None:
+            self.load_residue_embedding()
+
+    def __check_aggregator_embedding(self):
+        if self.__aggregator_embedding is None:
+            self.load_aggregator_embedding()
+
+
+def _load_model(device=None):
+    if not device:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    aggregator_model = get_aggregator_model(device=device)
+    aggregator_model.to(device)
+    aggregator_model.eval()
+    return aggregator_model
+
+
+def _load_res_model(device=None):
+    if not device:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return get_residue_model(device)

rcsb_embedding_model-0.0.6/src/rcsb_embedding_model/types/api_types.py

@@ -0,0 +1,29 @@
+from enum import Enum
+from os import PathLike
+from typing import NewType, Union, IO, Tuple, List, Optional
+
+StreamSrc = NewType('StreamSrc', Union[PathLike, IO])
+StreamTuple = NewType('StreamTuple', Tuple[StreamSrc, str, str])
+
+Devices = NewType('Devices', Union[int, List[int], "auto"])
+
+OptionalPath = NewType('OptionalPath', Optional[PathLike])
+
+
+class SrcFormat(str, Enum):
+    pdb = "pdb"
+    mmcif = "mmcif"
+    bciff = "binarycif"
+
+
+class Accelerator(str, Enum):
+    cpu = "cpu"
+    gpu = "gpu"
+    tpu = "tpu"
+    hpu = "hpu"
+    auto = "auto"
+
+
+class SrcLocation(str, Enum):
+    local = "local"
+    remote = "remote"

rcsb_embedding_model-0.0.6/src/rcsb_embedding_model/utils/data.py

@@ -0,0 +1,47 @@
+from io import StringIO
+
+import requests
+import torch
+
+
+def collate_seq_embeddings(batch_list):
+    """
+    Pads the tensors in a batch to the same size.
+
+    Args:
+        batch_list (list of torch.Tensor): A list of samples, where each sample is a tensor of shape (sequence_length, embedding_dim).
+
+    Returns:
+        tuple: A tuple containing:
+            - padded_batch (torch.Tensor): A tensor of shape (batch_size, max_seq_length, embedding_dim), where each sample is padded to the max sequence length.
+            - mask_batch (torch.Tensor): A tensor of shape (batch_size, max_seq_length) where padded positions are marked as False.
+    """
+    if batch_list[0] is None:
+        return None
+    device = batch_list[0].device  # Get the device of the input tensors
+    max_len = max(sample.size(0) for sample in batch_list)  # Determine the maximum sequence length
+    dim = batch_list[0].size(1)  # Determine the embedding dimension
+    batch_size = len(batch_list)  # Determine the batch size
+
+    # Initialize tensors for the padded batch and masks on the same device as the input tensors
+    padded_batch = torch.zeros((batch_size, max_len, dim), dtype=batch_list[0].dtype, device=device)
+    mask_batch = torch.ones((batch_size, max_len), dtype=torch.bool, device=device)
+
+    for i, sample in enumerate(batch_list):
+        seq_len = sample.size(0)  # Get the length of the current sequence
+        padded_batch[i, :seq_len] = sample  # Pad the sequence with zeros
+        mask_batch[i, :seq_len] = False  # Set mask positions for the actual data to False
+
+    return padded_batch, mask_batch
+
+
+def stringio_from_url(url):
+    try:
+        response = requests.get(url)
+        response.raise_for_status()
+        return StringIO(response.text)
+    except requests.exceptions.RequestException as e:
+        print(f"Error fetching URL: {e}")
+        return None
+
+

rcsb_embedding_model-0.0.6/src/rcsb_embedding_model/utils/model.py

@@ -0,0 +1,29 @@
+import torch
+from esm.models.esm3 import ESM3
+from esm.utils.constants.models import ESM3_OPEN_SMALL
+from huggingface_hub import hf_hub_download
+
+from rcsb_embedding_model.model.residue_embedding_aggregator import ResidueEmbeddingAggregator
+
+REPO_ID = "rcsb/rcsb-embedding-model"
+FILE_NAME = "rcsb-embedding-model.pt"
+REVISION = "410606e40b1bb7968ce318c41009355c3ac32503"
+
+
+def get_aggregator_model(device=None):
+    model_path = hf_hub_download(
+        repo_id=REPO_ID,
+        filename=FILE_NAME,
+        revision=REVISION
+    )
+    weights = torch.load(model_path, weights_only=True, map_location=device)
+    aggregator_model = ResidueEmbeddingAggregator()
+    aggregator_model.load_state_dict(weights)
+    return aggregator_model
+
+
+def get_residue_model(device=None):
+    return ESM3.from_pretrained(
+        ESM3_OPEN_SMALL,
+        device
+    )

rcsb_embedding_model-0.0.6/src/rcsb_embedding_model/utils/structure_parser.py

@@ -0,0 +1,51 @@
+
+from biotite.structure.io.pdb import PDBFile, get_structure as get_pdb_structure, get_assembly as get_pdb_assembly
+from biotite.structure.io.pdbx import CIFFile, get_structure, get_assembly, BinaryCIFFile
+
+
+def get_structure_from_src(
+        structure_src,
+        src_format="mmcif",
+        chain_id=None,
+        assembly_id=None
+):
+    if src_format == "pdb":
+        pdb_file = PDBFile.read(structure_src)
+        structure = __get_pdb_structure(pdb_file, assembly_id)
+    elif src_format == "mmcif":
+        cif_file = CIFFile.read(structure_src)
+        structure = __get_structure(cif_file, assembly_id)
+    elif src_format == "binarycif":
+        cif_file = BinaryCIFFile.read(structure_src)
+        structure = __get_structure(cif_file, assembly_id)
+    else:
+        raise RuntimeError(f"Unknown file format {src_format}")
+
+    if chain_id is not None:
+        structure = structure[structure.chain_id == chain_id]
+
+    return structure
+
+
+def __get_pdb_structure(pdb_file, assembly_id=None):
+    return get_pdb_structure(
+        pdb_file,
+        model=1
+    ) if assembly_id is None else get_pdb_assembly(
+        pdb_file,
+        assembly_id=assembly_id,
+        model=1
+    )
+
+
+def __get_structure(cif_file, assembly_id=None):
+    return get_structure(
+        cif_file,
+        model=1,
+        use_author_fields=False
+    ) if assembly_id is None else get_assembly(
+        cif_file,
+        assembly_id=assembly_id,
+        model=1,
+        use_author_fields=False
+    )

rcsb_embedding_model-0.0.6/src/rcsb_embedding_model/writer/batch_writer.py

@@ -0,0 +1,113 @@
+
+from abc import abstractmethod
+from collections import deque
+from abc import ABC
+
+import torch
+import pandas as pd
+
+from lightning.pytorch.callbacks import BasePredictionWriter
+
+
+class CoreBatchWriter(BasePredictionWriter, ABC):
+    def __init__(
+            self,
+            output_path,
+            postfix,
+            write_interval="batch"
+    ):
+        super().__init__(write_interval)
+        self.out_path = output_path
+        self.postfix = postfix
+
+    def write_on_batch_end(
+            self,
+            trainer,
+            pl_module,
+            prediction,
+            batch_indices,
+            batch,
+            batch_idx,
+            dataloader_idx
+    ):
+        if prediction is None:
+            return
+        embeddings, dom_ids = prediction
+        deque(map(
+            self._write_embedding,
+            embeddings,
+            dom_ids
+        ))
+
+    def file_name(self, dom_id):
+        return f'{self.out_path}/{dom_id}.{self.postfix}'
+
+    @abstractmethod
+    def _write_embedding(self, embedding, dom_id):
+        pass
+
+
+class CsvBatchWriter(CoreBatchWriter, ABC):
+    def __init__(
+            self,
+            output_path,
+            postfix="csv",
+            write_interval="batch"
+    ):
+        super().__init__(output_path, postfix, write_interval)
+
+    def _write_embedding(self, embedding, dom_id):
+        pd.DataFrame(embedding.to('cpu').numpy()).to_csv(
+            self.file_name(dom_id),
+            index=False,
+            header=False
+        )
+
+
+class TensorBatchWriter(CoreBatchWriter, ABC):
+    def __init__(
+            self,
+            output_path,
+            postfix="pt",
+            write_interval="batch",
+            device="cpu"
+    ):
+        super().__init__(output_path, postfix, write_interval)
+        self.device = device
+
+    def _write_embedding(self, embedding, dom_id):
+        torch.save(
+            embedding.to(self.device),
+            self.file_name(dom_id)
+        )
+
+
+class DataFrameStorage(CoreBatchWriter, ABC):
+    def __init__(
+            self,
+            output_path,
+            df_id,
+            postfix="pkl",
+            write_interval="batch"
+    ):
+        super().__init__(output_path, postfix, write_interval)
+        self.df_id = df_id
+        self.embedding = pd.DataFrame(
+            data={},
+            columns=['id', 'embedding'],
+        )
+
+    def _write_embedding(self, embedding, dom_id):
+        self.embedding = pd.concat([
+            self.embedding,
+            pd.DataFrame(
+                data={'id': dom_id, 'embedding': [embedding.to('cpu').numpy()]},
+                columns=['id', 'embedding'],
+            )
+        ], ignore_index=True)
+
+    def on_predict_end(self, trainer, pl_module):
+        self.embedding.to_pickle(
+            f"{self.out_path}/{self.df_id}.pkl.gz",
+            compression='gzip'
+        )

rcsb_embedding_model-0.0.5/tests/test_model.py → rcsb_embedding_model-0.0.6/tests/test_embedding_model.py

@@ -2,6 +2,7 @@ import os
 import unittest
 
 from rcsb_embedding_model import RcsbStructureEmbedding
+from rcsb_embedding_model.types.api_types import SrcFormat
 
 
 class TestEmbeddingModel(unittest.TestCase):
@@ -13,7 +14,7 @@ class TestEmbeddingModel(unittest.TestCase):
         model = RcsbStructureEmbedding()
         res_embedding = model.residue_embedding(
             f"{self.__test_path}/resources/1acb.cif",
-            format="mmcif",
+            src_format=SrcFormat.mmcif,
             chain_id='A'
         )
         self.assertEqual(list(res_embedding.shape), [243, 1536])
@@ -33,7 +34,7 @@ class TestEmbeddingModel(unittest.TestCase):
         model = RcsbStructureEmbedding()
         res_embedding = model.residue_embedding(
             f"{self.__test_path}/resources/1acb.cif",
-            format="mmcif",
+            src_format=SrcFormat.mmcif,
             chain_id='A'
         )
         structure_embedding = model.aggregator_embedding(

rcsb_embedding_model-0.0.5/src/rcsb_embedding_model/rcsb_structure_embedding.py

@@ -1,171 +0,0 @@
-import torch
-from biotite.structure import get_residues, chain_iter, filter_amino_acids
-from biotite.structure.io.pdb import PDBFile, get_structure as get_pdb_structure, get_assembly as get_pdb_assembly
-from biotite.structure.io.pdbx import CIFFile, get_structure, get_assembly, BinaryCIFFile
-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 huggingface_hub import hf_hub_download
-
-from rcsb_embedding_model.model.residue_embedding_aggregator import ResidueEmbeddingAggregator
-
-
-class RcsbStructureEmbedding:
-
-    MIN_RES = 10
-    REPO_ID = "rcsb/rcsb-embedding-model"
-    FILE_NAME = "rcsb-embedding-model.pt"
-    VERSION = "410606e40b1bb7968ce318c41009355c3ac32503"
-
-    def __init__(self):
-        self.__residue_embedding = None
-        self.__aggregator_embedding = None
-
-    def load_models(self, device=None):
-        self.load_residue_embedding(device)
-        self.load_aggregator_embedding(device)
-
-    def load_residue_embedding(self, device=None):
-        if not device:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.__residue_embedding = _load_res_model(device)
-
-    def load_aggregator_embedding(self, device=None):
-        if not device:
-            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.__aggregator_embedding = _load_model(
-            _download_model(
-                RcsbStructureEmbedding.REPO_ID,
-                RcsbStructureEmbedding.FILE_NAME,
-                RcsbStructureEmbedding.VERSION
-            ),
-            device
-        )
-
-    def structure_embedding(self, structure_src, format="pdb", chain_id=None, assembly_id=None):
-        res_embedding = self.residue_embedding(structure_src, format, chain_id, assembly_id)
-        return self.aggregator_embedding(res_embedding)
-
-    def residue_embedding(self, structure_src, format="pdb", chain_id=None, assembly_id=None):
-        self.__check_residue_embedding()
-        structure = _get_structure_from_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)]
-            if len(atom_res) == 0 or len(get_residues(atom_res)[0]) < RcsbStructureEmbedding.MIN_RES:
-                continue
-            protein_chain = ProteinChain.from_atomarray(atom_ch)
-            protein = ESMProtein.from_protein_chain(protein_chain)
-            protein_tensor = self.__residue_embedding.encode(protein)
-            embedding_ch.append(self.__residue_embedding.forward_and_sample(
-                protein_tensor, SamplingConfig(return_per_residue_embeddings=True)
-            ).per_residue_embedding)
-        return torch.cat(
-            embedding_ch,
-            dim=0
-        )
-
-    def sequence_embedding(self, sequence):
-        self.__check_residue_embedding()
-
-        if sequence.startswith(">"):
-            sequence = "".join(line.strip() for line in sequence.splitlines() if not line.startswith(">"))
-
-        if len(sequence) < RcsbStructureEmbedding.MIN_RES:
-            raise ValueError(f"Sequence too short for embedding (min {RcsbStructureEmbedding.MIN_RES} residues)")
-
-        protein = ESMProtein(sequence=sequence)
-        protein_tensor = self.__residue_embedding.encode(protein)
-
-        result = self.__residue_embedding.forward_and_sample(
-            protein_tensor,
-            SamplingConfig(return_per_residue_embeddings=True)
-        )
-
-        return result.per_residue_embedding
-
-    def aggregator_embedding(self, residue_embedding):
-        self.__check_aggregator_embedding()
-        return self.__aggregator_embedding(residue_embedding)
-
-    def __check_residue_embedding(self):
-        if self.__residue_embedding is None:
-            self.load_residue_embedding()
-
-    def __check_aggregator_embedding(self):
-        if self.__aggregator_embedding is None:
-            self.load_aggregator_embedding()
-
-
-def _get_structure_from_src(structure_src, format="pdb", chain_id=None, assembly_id=None):
-    if format == "pdb":
-        pdb_file = PDBFile.read(structure_src)
-        structure = _get_pdb_structure(pdb_file, assembly_id)
-    elif format == "mmcif":
-        cif_file = CIFFile.read(structure_src)
-        structure = _get_structure(cif_file, assembly_id)
-    elif format == "binarycif":
-        cif_file = BinaryCIFFile.read(structure_src)
-        structure = _get_structure(cif_file, assembly_id)
-    else:
-        raise RuntimeError(f"Unknown file format {format}")
-
-    if chain_id is not None:
-        structure = structure[structure.chain_id == chain_id]
-
-    return structure
-
-
-def _get_pdb_structure(pdb_file, assembly_id = None):
-    return get_pdb_structure(
-        pdb_file,
-        model=1
-    ) if assembly_id is None else get_pdb_assembly(
-        pdb_file,
-        assembly_id=assembly_id,
-        model=1
-    )
-
-
-def _get_structure(cif_file, assembly_id = None):
-    return get_structure(
-        cif_file,
-        model=1,
-        use_author_fields=False
-    ) if assembly_id is None else get_assembly(
-        cif_file,
-        assembly_id=assembly_id,
-        model=1,
-        use_author_fields=False
-    )
-
-
-def _download_model(
-        repo_id,
-        filename,
-        revision
-):
-    return hf_hub_download(
-        repo_id=repo_id,
-        filename=filename,
-        revision=revision
-    )
-
-
-def _load_model(model_path, device=None):
-    if not device:
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    weights = torch.load(model_path, weights_only=True, map_location=device)
-    aggregator_model = ResidueEmbeddingAggregator()
-    aggregator_model.load_state_dict(weights)
-    aggregator_model.to(device)
-    aggregator_model.eval()
-    return aggregator_model
-
-
-def _load_res_model(device=None):
-    return ESM3.from_pretrained(
-        ESM3_OPEN_SMALL,
-        device
-    )