rcsb-embedding-model 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

rcsb_embedding_model/cli/args_utils.py

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

rcsb_embedding_model/cli/inference.py

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

rcsb_embedding_model/dataset/esm_prot_from_csv.py

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

rcsb_embedding_model/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/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/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/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/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/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/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/rcsb_structure_embedding.py

@@ -1,55 +1,63 @@
 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
+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
-    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):
+    def load_models(
+            self,
+            device: torch.device = None
+    ):
         self.load_residue_embedding(device)
         self.load_aggregator_embedding(device)
 
-    def load_residue_embedding(self, device=None):
+    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=None):
+    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(
-            _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)
+        self.__aggregator_embedding = _load_model(device)
+
+    def structure_embedding(
+            self,
+            src_structure: StreamSrc,
+            src_format: SrcFormat = SrcFormat.mmcif,
+            chain_id: str = None,
+            assembly_id: str = None
+    ):
+        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, format="pdb", chain_id=None, assembly_id=None):
+    def residue_embedding(
+            self,
+            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, 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)]
@@ -66,7 +74,10 @@ class RcsbStructureEmbedding:
             dim=0
         )
 
-    def sequence_embedding(self, sequence):
+    def sequence_embedding(
+            self,
+            sequence: str
+    ):
         self.__check_residue_embedding()
 
         if sequence.startswith(">"):
@@ -85,7 +96,10 @@ class RcsbStructureEmbedding:
 
         return result.per_residue_embedding
 
-    def aggregator_embedding(self, residue_embedding):
+    def aggregator_embedding(
+            self,
+            residue_embedding: torch.Tensor
+    ):
         self.__check_aggregator_embedding()
         return self.__aggregator_embedding(residue_embedding)
 
@@ -98,74 +112,16 @@ class RcsbStructureEmbedding:
             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):
+def _load_model(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 = get_aggregator_model(device=device)
     aggregator_model.to(device)
     aggregator_model.eval()
     return aggregator_model
 
 
 def _load_res_model(device=None):
-    return ESM3.from_pretrained(
-        ESM3_OPEN_SMALL,
-        device
-    )
+    if not device:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return get_residue_model(device)

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/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/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/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(
+        src_structure,
+        src_format="mmcif",
+        chain_id=None,
+        assembly_id=None
+):
+    if src_format == "pdb":
+        pdb_file = PDBFile.read(src_structure)
+        structure = __get_pdb_structure(pdb_file, assembly_id)
+    elif src_format == "mmcif":
+        cif_file = CIFFile.read(src_structure)
+        structure = __get_structure(cif_file, assembly_id)
+    elif src_format == "binarycif":
+        cif_file = BinaryCIFFile.read(src_structure)
+        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/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.7.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,24 @@
+rcsb_embedding_model/__init__.py,sha256=r3gLdeBIXkQEQA_K6QcRPO-TtYuAQSutk6pXRUE_nas,120
+rcsb_embedding_model/rcsb_structure_embedding.py,sha256=qGUEdRPjYbsFWThsQa_ZVaSJ7nURnfRBLBqJlLbcY0I,4433
+rcsb_embedding_model/cli/args_utils.py,sha256=7nP2q8pL5dWK_U7opxtWmoFcYVwasky6elHk-dASFaI,165
+rcsb_embedding_model/cli/inference.py,sha256=sx8cGiq_japc0mKFarK1aVkGfK-FhTeZdn_Ng0ijezE,6590
+rcsb_embedding_model/dataset/esm_prot_from_csv.py,sha256=1XMiYyJXfodXZGSrU07uyoYbdKR9-KvNfb1xNqab_W8,2722
+rcsb_embedding_model/dataset/residue_embedding_from_csv.py,sha256=0-5L64tyER-RpT166pC71qxOpUdVZbcuBQONPcAIuno,862
+rcsb_embedding_model/inference/chain_inference.py,sha256=SgXDa-TkDcvlkQxqEwDt81RdE7NmgiaJD8uaROgMbl8,1506
+rcsb_embedding_model/inference/esm_inference.py,sha256=pX-_RhzAIvL0Zdg9wjScLBP6Y1sq4RLNio4-vdR5MLU,1498
+rcsb_embedding_model/inference/structure_inference.py,sha256=qPzAGWyzFWqeKV9yoPSw4LrEB9XgKTJnRQysSBhfg14,1564
+rcsb_embedding_model/model/layers.py,sha256=lhKaWC4gTS_T5lHOP0mgnnP8nKTPEOm4MrjhESA4hE8,743
+rcsb_embedding_model/model/residue_embedding_aggregator.py,sha256=k3UW63Ax8DtjCMdD3O5xNxtyAu28l2n3-Ab6nS0atm0,1967
+rcsb_embedding_model/modules/chain_module.py,sha256=sDSPXJmWuU2C3lt1NorlbUVWZvRSLzumPdFQk01h3VI,403
+rcsb_embedding_model/modules/esm_module.py,sha256=CTHGOATXiarqZsBsZ8oxGJBj20A73186Slpr0EzMJsE,770
+rcsb_embedding_model/modules/structure_module.py,sha256=dEtDNdWo1j2sSDa0JiOHQfEfQzIWqSLEKpvOX0GrXZ4,1048
+rcsb_embedding_model/types/api_types.py,sha256=x7274MyjkRXn8B-W-PY5PK9g0CP1pT_clZbrAuFuHPA,626
+rcsb_embedding_model/utils/data.py,sha256=LGw3wvq_LCcqSovHZacOqxEczn12SZk2i51WK9xkk0k,1877
+rcsb_embedding_model/utils/model.py,sha256=rpZa-gfm3cEtbBd7UXMHrZv3x6f0AC8TJT3gtrSxr5I,852
+rcsb_embedding_model/utils/structure_parser.py,sha256=0lcjCuQMCh0lb3OMj76rqf7kACzJgOwdk3EZ7-ZOQfI,1492
+rcsb_embedding_model/writer/batch_writer.py,sha256=ekgzFZyoKpcnZ3IDP9hfOWBpuHxUQ31P35ViDAi-Edw,2843
+rcsb_embedding_model-0.0.7.dist-info/METADATA,sha256=mfl1YYB48Um5FdZZkHOwzzMPRvsw_HlFHeqXsCGWs0Q,4959
+rcsb_embedding_model-0.0.7.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+rcsb_embedding_model-0.0.7.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
+rcsb_embedding_model-0.0.7.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
+rcsb_embedding_model-0.0.7.dist-info/RECORD,,

rcsb_embedding_model-0.0.7.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+inference = rcsb_embedding_model.cli.inference:app

rcsb_embedding_model-0.0.5.dist-info/METADATA

@@ -1,115 +0,0 @@
-Metadata-Version: 2.4
-Name: rcsb-embedding-model
-Version: 0.0.5
-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: torch>=2.2.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.5.dist-info/RECORD

@@ -1,8 +0,0 @@
-rcsb_embedding_model/__init__.py,sha256=r3gLdeBIXkQEQA_K6QcRPO-TtYuAQSutk6pXRUE_nas,120
-rcsb_embedding_model/rcsb_structure_embedding.py,sha256=CIgM3euWK-QO19V_D8flAB9Suv5aLJnU1DDEYea-n6w,6016
-rcsb_embedding_model/model/layers.py,sha256=lhKaWC4gTS_T5lHOP0mgnnP8nKTPEOm4MrjhESA4hE8,743
-rcsb_embedding_model/model/residue_embedding_aggregator.py,sha256=k3UW63Ax8DtjCMdD3O5xNxtyAu28l2n3-Ab6nS0atm0,1967
-rcsb_embedding_model-0.0.5.dist-info/METADATA,sha256=_K496_MAxRYfDiLun4L4LlvCQEUgOSmX7aC7vIrEahw,5381
-rcsb_embedding_model-0.0.5.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-rcsb_embedding_model-0.0.5.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
-rcsb_embedding_model-0.0.5.dist-info/RECORD,,