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

rcsb_embedding_model/cli/inference.py

@@ -1,9 +1,11 @@
+import sys
 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
+from rcsb_embedding_model.types.api_types import StructureFormat, Accelerator, SrcLocation, SrcProteinFrom, \
+    StructureLocation, SrcAssemblyFrom
 
 app = typer.Typer(
     add_completion=False
@@ -12,7 +14,7 @@ app = typer.Typer(
 
 @app.command(
     name="residue-embedding",
-    help="Calculate residue level embeddings of protein structures using ESM3."
+    help="Calculate residue level embeddings of protein structures using ESM3. Predictions are stored as torch tensor files."
 )
 def residue_embedding(
         src_file: Annotated[typer.FileText, typer.Option(
@@ -20,21 +22,27 @@ def residue_embedding(
             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.'
+            help='CSV file 4 (or 3) columns: Structure Name | Structure File Path | Chain Id (asym_i for cif files. This field is required if src-from=chain) | Output Embedding 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.'
+            help='Output path to store predictions. Embeddings are stored as torch tensor files.'
         )],
-        src_location: Annotated[SrcLocation, typer.Option(
-            help='Source input location.'
-        )] = SrcLocation.local,
-        src_format: Annotated[SrcFormat, typer.Option(
+        src_from: Annotated[SrcProteinFrom, typer.Option(
+            help='Use specific chains or all chains in a structure.'
+        )] = SrcProteinFrom.chain,
+        structure_location: Annotated[StructureLocation, typer.Option(
+            help='Structure file location.'
+        )] = StructureLocation.local,
+        structure_format: Annotated[StructureFormat, typer.Option(
             help='Structure file format.'
-        )] = SrcFormat.mmcif,
+        )] = StructureFormat.mmcif,
+        min_res_n: Annotated[int, typer.Option(
+            help='When using all chains in a structure, consider only chains with more than <min_res_n> residues.'
+        )] = 0,
         batch_size: Annotated[int, typer.Option(
             help='Number of samples processed together in one iteration.'
         )] = 1,
@@ -53,9 +61,12 @@ def residue_embedding(
 ):
     from rcsb_embedding_model.inference.esm_inference import predict
     predict(
-        csv_file=src_file,
-        src_location=src_location,
-        src_format=src_format,
+        src_stream=src_file,
+        src_location=SrcLocation.local,
+        src_from=src_from,
+        structure_location=structure_location,
+        structure_format=structure_format,
+        min_res_n=min_res_n,
         batch_size=batch_size,
         num_workers=num_workers,
         num_nodes=num_nodes,
@@ -67,7 +78,7 @@ def residue_embedding(
 
 @app.command(
     name="structure-embedding",
-    help="Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas data-frame file."
+    help="Calculate single-chain protein embeddings from structural files. Predictions are stored in a single pandas DataFrame file."
 )
 def structure_embedding(
         src_file: Annotated[typer.FileText, typer.Option(
@@ -75,24 +86,30 @@ def structure_embedding(
             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.'
+            help='CSV file 4 (or 3) columns: Structure Name | Structure File Path | Chain Id (asym_i for cif files. This field is required if src-from=chain) | Output Embedding 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.'
+            help='Output path to store predictions. Embeddings are stored as a single DataFrame file (see out-df-name).'
         )],
-        out_df_id: Annotated[str, typer.Option(
-            help='File name to store predicted embeddings.'
+        out_df_name: Annotated[str, typer.Option(
+            help='File name (without extension) for storing embeddings as a pandas DataFrame pickle (.pkl). The DataFrame contains 2 columns: Id | Embedding'
         )],
-        src_location: Annotated[SrcLocation, typer.Option(
+        src_from: Annotated[SrcProteinFrom, typer.Option(
+            help='Use specific chains or all chains in a structure.'
+        )] = SrcProteinFrom.chain,
+        structure_location: Annotated[StructureLocation, typer.Option(
             help='Source input location.'
-        )] = SrcLocation.local,
-        src_format: Annotated[SrcFormat, typer.Option(
+        )] = StructureLocation.local,
+        structure_format: Annotated[StructureFormat, typer.Option(
             help='Structure file format.'
-        )] = SrcFormat.mmcif,
+        )] = StructureFormat.mmcif,
+        min_res_n: Annotated[int, typer.Option(
+            help='When using all chains in a structure, consider only chains with more than <min_res_n> residues.'
+        )] = 0,
         batch_size: Annotated[int, typer.Option(
             help='Number of samples processed together in one iteration.'
         )] = 1,
@@ -111,22 +128,25 @@ def structure_embedding(
 ):
     from rcsb_embedding_model.inference.structure_inference import predict
     predict(
-        csv_file=src_file,
-        src_location=src_location,
-        src_format=src_format,
+        src_stream=src_file,
+        src_location=SrcLocation.local,
+        src_from=src_from,
+        structure_location=structure_location,
+        structure_format=structure_format,
+        min_res_n=min_res_n,
         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
+        out_df_name=out_df_name
     )
 
 
 @app.command(
     name="chain-embedding",
-    help="Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files."
+    help="Calculate single-chain protein embeddings from residue level embeddings stored as torch tensor files. Predictions are stored as csv files."
 )
 def chain_embedding(
         src_file: Annotated[typer.FileText, typer.Option(
@@ -134,14 +154,14 @@ def chain_embedding(
             file_okay=True,
             dir_okay=False,
             resolve_path=True,
-            help='CSV file 2 columns: Residue Embedding Tensor File | Output file name.'
+            help='CSV file 2 columns: Residue embedding torch tensor file | Output embedding 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.'
+            help='Output path to store predictions. Embeddings are stored as csv files.'
         )],
         batch_size: Annotated[int, typer.Option(
             help='Number of samples processed together in one iteration.'
@@ -161,7 +181,83 @@ def chain_embedding(
 ):
     from rcsb_embedding_model.inference.chain_inference import predict
     predict(
-        csv_file=src_file,
+        src_stream=src_file,
+        src_location=SrcLocation.local,
+        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="assembly-embedding",
+    help="Calculate assembly embeddings from residue level embeddings stored as torch tensor files. Predictions are stored as csv files."
+)
+def assembly_embedding(
+        src_file: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            resolve_path=True,
+            help='CSV file 4 columns: Structure Name | Structure File Path | Assembly Id | Output embedding name.'
+        )],
+        res_embedding_location: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Path where residue level embeddings for single chains are located.'
+        )],
+        output_path: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Output path to store predictions. Embeddings are stored as csv files.'
+        )],
+        src_from: Annotated[SrcAssemblyFrom, typer.Option(
+            help='Use specific assembly or all assemblies in a structure.'
+        )] = SrcAssemblyFrom.assembly,
+        structure_location: Annotated[StructureLocation, typer.Option(
+            help='Source input location.'
+        )] = StructureLocation.local,
+        structure_format: Annotated[StructureFormat, typer.Option(
+            help='Structure file format.'
+        )] = StructureFormat.mmcif,
+        min_res_n: Annotated[int, typer.Option(
+            help='Consider only assembly chains with more than <min_res_n> residues.'
+        )] = 0,
+        max_res_n: Annotated[int, typer.Option(
+            help='Stop adding assembly chains when number of residues is greater than <max_res_n> residues.'
+        )] = sys.maxsize,
+        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.assembly_inferece import predict
+    predict(
+        src_stream=src_file,
+        res_embedding_location=res_embedding_location,
+        src_location=SrcLocation.local,
+        src_from=src_from,
+        structure_location=structure_location,
+        structure_format=structure_format,
+        min_res_n=min_res_n,
+        max_res_n=max_res_n,
         batch_size=batch_size,
         num_workers=num_workers,
         num_nodes=num_nodes,

rcsb_embedding_model/dataset/esm_prot_from_chain.py

@@ -0,0 +1,102 @@
+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 StructureFormat, StructureLocation, SrcLocation
+from rcsb_embedding_model.utils.data import stringio_from_url
+from rcsb_embedding_model.utils.structure_parser import rename_atom_ch
+from rcsb_embedding_model.utils.structure_provider import StructureProvider
+
+
+class EsmProtFromChain(Dataset):
+
+    STREAM_NAME_ATTR = 'stream_name'
+    STREAM_ATTR = 'stream'
+    CH_ATTR = 'chain_id'
+    ITEM_NAME_ATTR = 'item_name'
+
+    COLUMNS = [STREAM_NAME_ATTR, STREAM_ATTR, CH_ATTR, ITEM_NAME_ATTR]
+
+    def __init__(
+        self,
+        src_stream,
+        src_location=SrcLocation.local,
+        structure_location=StructureLocation.local,
+        structure_format=StructureFormat.mmcif,
+        structure_provider=StructureProvider()
+    ):
+        super().__init__()
+        self.__structure_provider = structure_provider
+        self.src_location = src_location
+        self.structure_location = structure_location
+        self.structure_format = structure_format
+        self.data = pd.DataFrame()
+        self.__load_stream(src_stream)
+
+    def __load_stream(self, src_stream):
+        self.data = pd.DataFrame(
+            src_stream,
+            dtype=str,
+            columns=EsmProtFromChain.COLUMNS
+        ) if self.src_location == SrcLocation.stream else pd.read_csv(
+            src_stream,
+            header=None,
+            index_col=None,
+            dtype=str,
+            names=EsmProtFromChain.COLUMNS
+        )
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        src_name = self.data.loc[idx, EsmProtFromChain.STREAM_NAME_ATTR]
+        src_structure = self.data.loc[idx, EsmProtFromChain.STREAM_ATTR]
+        chain_id = self.data.loc[idx, EsmProtFromChain.CH_ATTR]
+        item_name = self.data.loc[idx, EsmProtFromChain.ITEM_NAME_ATTR]
+        structure = self.__structure_provider.get_structure(
+            src_name=src_name,
+            src_structure=stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure,
+            structure_format=self.structure_format,
+            chain_id=chain_id
+        )
+        for atom_ch in chain_iter(structure):
+            protein_chain = ProteinChain.from_atomarray(rename_atom_ch(atom_ch))
+            return ESMProtein.from_protein_chain(protein_chain), item_name
+
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--file_list', type=argparse.FileType('r'), required=True)
+    args = parser.parse_args()
+
+    dataset = EsmProtFromChain(
+        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, 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(prot_name, embeddings.shape)

rcsb_embedding_model/dataset/esm_prot_from_structure.py

@@ -0,0 +1,63 @@
+
+import pandas as pd
+
+from rcsb_embedding_model.dataset.esm_prot_from_chain import EsmProtFromChain
+from rcsb_embedding_model.types.api_types import StructureLocation, StructureFormat, SrcLocation
+from rcsb_embedding_model.utils.data import stringio_from_url
+from rcsb_embedding_model.utils.structure_parser import get_protein_chains
+from rcsb_embedding_model.utils.structure_provider import StructureProvider
+
+
+class EsmProtFromStructure(EsmProtFromChain):
+
+    STREAM_NAME_ATTR = 'stream_name'
+    STREAM_ATTR = 'stream'
+    ITEM_NAME_ATTR = 'item_name'
+
+    COLUMNS = [STREAM_NAME_ATTR, STREAM_ATTR, ITEM_NAME_ATTR]
+
+    def __init__(
+            self,
+            src_stream,
+            src_location=SrcLocation.local,
+            structure_location=StructureLocation.local,
+            structure_format=StructureFormat.mmcif,
+            min_res_n=0,
+            structure_provider=StructureProvider()
+    ):
+        self.min_res_n = min_res_n
+        self.src_location = src_location
+        self.structure_location = structure_location
+        self.structure_format = structure_format
+        self.__structure_provider = structure_provider
+        super().__init__(
+            src_stream=self.__get_chains(src_stream),
+            src_location=SrcLocation.stream,
+            structure_location=StructureLocation.local,
+            structure_format=structure_format
+        )
+
+    def __get_chains(self, src_stream):
+        chains = []
+        for idx, row in (pd.DataFrame(
+            src_stream,
+            dtype=str,
+            columns=self.COLUMNS
+        ) if self.src_location == SrcLocation.stream else pd.read_csv(
+            src_stream,
+            header=None,
+            index_col=None,
+            dtype=str,
+            names=EsmProtFromStructure.COLUMNS
+        )).iterrows():
+            src_name = row[EsmProtFromStructure.STREAM_NAME_ATTR]
+            src_structure = row[EsmProtFromStructure.STREAM_ATTR]
+            item_name = row[EsmProtFromStructure.ITEM_NAME_ATTR]
+            structure = self.__structure_provider.get_structure(
+                src_name=src_name,
+                src_structure=stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure,
+                structure_format=self.structure_format
+            )
+            for ch in get_protein_chains(structure, self.min_res_n):
+                chains.append((src_name, src_structure, ch, f"{item_name}.{ch}"))
+        return tuple(chains)

rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py

@@ -0,0 +1,68 @@
+import sys
+
+import pandas as pd
+
+from rcsb_embedding_model.dataset.residue_assembly_embedding_from_tensor_file import ResidueAssemblyEmbeddingFromTensorFile
+from rcsb_embedding_model.types.api_types import SrcLocation, StructureLocation, StructureFormat
+from rcsb_embedding_model.utils.data import stringio_from_url
+from rcsb_embedding_model.utils.structure_parser import get_assemblies
+from rcsb_embedding_model.utils.structure_provider import StructureProvider
+
+
+class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile):
+
+    STREAM_NAME_ATTR = 'stream_name'
+    STREAM_ATTR = 'stream'
+    ITEM_NAME_ATTR = 'item_name'
+
+    COLUMNS = [STREAM_NAME_ATTR, STREAM_ATTR, ITEM_NAME_ATTR]
+
+    def __init__(
+            self,
+            src_stream,
+            res_embedding_location,
+            src_location=SrcLocation.local,
+            structure_location=StructureLocation.local,
+            structure_format=StructureFormat.mmcif,
+            min_res_n=0,
+            max_res_n=sys.maxsize,
+            structure_provider=StructureProvider()
+    ):
+        self.src_location = src_location
+        self.structure_location = structure_location
+        self.structure_format = structure_format
+        self.min_res_n = min_res_n
+        self.max_res_n = max_res_n
+        self.__structure_provider = structure_provider
+        super().__init__(
+            src_stream=self.__get_assemblies(src_stream),
+            res_embedding_location=res_embedding_location,
+            src_location=src_location,
+            structure_location=structure_location,
+            structure_format=structure_format,
+            min_res_n=min_res_n,
+            max_res_n=max_res_n,
+            structure_provider=structure_provider
+        )
+
+    def __get_assemblies(self, src_stream):
+        assemblies = []
+        for idx, row in (pd.DataFrame(
+                src_stream,
+                dtype=str,
+                columns=self.COLUMNS
+        ) if self.src_location == SrcLocation.stream else pd.read_csv(
+            src_stream,
+            header=None,
+            index_col=None,
+            dtype=str,
+            names=ResidueAssemblyDatasetFromStructure.COLUMNS
+        )).iterrows():
+            src_name = row[ResidueAssemblyDatasetFromStructure.STREAM_NAME_ATTR]
+            src_structure = row[ResidueAssemblyDatasetFromStructure.STREAM_ATTR]
+            src_structure = stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure
+            item_name = row[ResidueAssemblyDatasetFromStructure.ITEM_NAME_ATTR]
+            for assembly_id in get_assemblies(src_structure=src_structure, structure_format=self.structure_format):
+                assemblies.append((src_name, src_structure, str(assembly_id), f"{item_name}.{assembly_id}"))
+
+        return tuple(assemblies)

rcsb_embedding_model/dataset/residue_assembly_embedding_from_tensor_file.py

@@ -0,0 +1,94 @@
+import sys
+
+import pandas as pd
+from torch.utils.data import Dataset, DataLoader
+
+from rcsb_embedding_model.types.api_types import StructureLocation, StructureFormat, SrcLocation
+from rcsb_embedding_model.utils.data import stringio_from_url, concatenate_tensors
+from rcsb_embedding_model.utils.structure_parser import get_protein_chains
+from rcsb_embedding_model.utils.structure_provider import StructureProvider
+
+
+class ResidueAssemblyEmbeddingFromTensorFile(Dataset):
+
+    STREAM_NAME_ATTR = 'stream_name'
+    STREAM_ATTR = 'stream'
+    ASSEMBLY_ATTR = 'assembly_id'
+    ITEM_NAME_ATTR = 'item_name'
+
+    COLUMNS = [STREAM_NAME_ATTR, STREAM_ATTR, ASSEMBLY_ATTR, ITEM_NAME_ATTR]
+
+    def __init__(
+            self,
+            src_stream,
+            res_embedding_location,
+            src_location=SrcLocation.local,
+            structure_location=StructureLocation.local,
+            structure_format=StructureFormat.mmcif,
+            min_res_n=0,
+            max_res_n=sys.maxsize,
+            structure_provider=StructureProvider()
+    ):
+        super().__init__()
+        self.res_embedding_location = res_embedding_location
+        self.src_location = src_location
+        self.structure_location = structure_location
+        self.structure_format = structure_format
+        self.min_res_n = min_res_n
+        self.max_res_n = max_res_n
+        self.data = pd.DataFrame()
+        self.__load_stream(src_stream)
+        self.__structure_provider = structure_provider
+
+    def __load_stream(self, src_stream):
+        self.data = pd.DataFrame(
+            src_stream,
+            dtype=str,
+            columns=ResidueAssemblyEmbeddingFromTensorFile.COLUMNS
+        ) if self.src_location == SrcLocation.stream else pd.read_csv(
+            src_stream,
+            header=None,
+            index_col=None,
+            dtype=str,
+            names=ResidueAssemblyEmbeddingFromTensorFile.COLUMNS
+        )
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        src_name = self.data.loc[idx, ResidueAssemblyEmbeddingFromTensorFile.STREAM_NAME_ATTR]
+        src_structure = self.data.loc[idx, ResidueAssemblyEmbeddingFromTensorFile.STREAM_ATTR]
+        assembly_id = self.data.loc[idx, ResidueAssemblyEmbeddingFromTensorFile.ASSEMBLY_ATTR]
+        item_name = self.data.loc[idx, ResidueAssemblyEmbeddingFromTensorFile.ITEM_NAME_ATTR]
+
+        structure = self.__structure_provider.get_structure(
+            src_name=src_name,
+            src_structure=stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure,
+            structure_format=self.structure_format,
+            assembly_id=assembly_id
+        )
+        residue_embedding_files = [
+            f"{self.res_embedding_location}/{src_name}.{ch}.pt" for ch in get_protein_chains(structure, self.min_res_n)
+        ]
+        return concatenate_tensors(residue_embedding_files, self.max_res_n), item_name
+
+
+if __name__ == "__main__":
+
+    dataset = ResidueAssemblyEmbeddingFromTensorFile(
+        src_stream="/Users/joan/tmp/assembly-test.csv",
+        res_embedding_location="/Users/joan/tmp",
+        src_location=SrcLocation.local,
+        structure_location=StructureLocation.local,
+        structure_format=StructureFormat.mmcif
+    )
+
+    dataloader = DataLoader(
+        dataset,
+        batch_size=1,
+        collate_fn=lambda _: _
+    )
+
+    for _batch in dataloader:
+        print(_batch)

rcsb_embedding_model/dataset/residue_embedding_from_tensor_file.py

@@ -0,0 +1,43 @@
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+from rcsb_embedding_model.types.api_types import StructureLocation, SrcLocation
+
+
+class ResidueEmbeddingFromTensorFile(Dataset):
+
+    FILE_ATTR = 'file'
+    ITEM_NAME_ATTR = 'item_name'
+
+    COLUMNS = [FILE_ATTR, ITEM_NAME_ATTR]
+
+    def __init__(
+            self,
+            src_stream,
+            src_location=SrcLocation.local
+    ):
+        super().__init__()
+        self.src_location = src_location
+        self.data = pd.DataFrame()
+        self.__load_stream(src_stream)
+
+    def __load_stream(self, src_stream):
+        self.data = pd.DataFrame(
+            src_stream,
+            dtype=str,
+            columns=self.COLUMNS
+        ) if self.src_location == SrcLocation.stream else pd.read_csv(
+            src_stream,
+            header=None,
+            index_col=None,
+            names=ResidueEmbeddingFromTensorFile.COLUMNS
+        )
+
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, idx):
+        embedding_src = self.data.loc[idx, ResidueEmbeddingFromTensorFile.FILE_ATTR]
+        item_name = self.data.loc[idx, ResidueEmbeddingFromTensorFile.ITEM_NAME_ATTR]
+        return torch.load(embedding_src, map_location=torch.device('cpu')), item_name

rcsb_embedding_model/inference/assembly_inferece.py

@@ -0,0 +1,53 @@
+import sys
+
+from rcsb_embedding_model.dataset.resdiue_assembly_embedding_from_structure import ResidueAssemblyDatasetFromStructure
+from rcsb_embedding_model.dataset.residue_assembly_embedding_from_tensor_file import ResidueAssemblyEmbeddingFromTensorFile
+from rcsb_embedding_model.types.api_types import FileOrStreamTuple, SrcLocation, Accelerator, Devices, OptionalPath, EmbeddingPath, StructureLocation, StructureFormat, SrcAssemblyFrom
+from rcsb_embedding_model.inference.chain_inference import predict as chain_predict
+
+
+def predict(
+        src_stream: FileOrStreamTuple,
+        res_embedding_location: EmbeddingPath,
+        src_location: SrcLocation = SrcLocation.local,
+        src_from: SrcAssemblyFrom = SrcAssemblyFrom.assembly,
+        structure_location: StructureLocation = StructureLocation.local,
+        structure_format: StructureFormat = StructureFormat.mmcif,
+        min_res_n: int = 0,
+        max_res_n: int = sys.maxsize,
+        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 = ResidueAssemblyEmbeddingFromTensorFile(
+        src_stream=src_stream,
+        res_embedding_location=res_embedding_location,
+        src_location=src_location,
+        structure_location=structure_location,
+        structure_format=structure_format,
+        min_res_n=min_res_n,
+        max_res_n=max_res_n
+    ) if src_from == SrcAssemblyFrom.assembly else ResidueAssemblyDatasetFromStructure(
+        src_stream=src_stream,
+        res_embedding_location=res_embedding_location,
+        src_location=src_location,
+        structure_location=structure_location,
+        structure_format=structure_format,
+        min_res_n=min_res_n,
+        max_res_n=max_res_n
+    )
+
+    return chain_predict(
+        src_stream=src_stream,
+        src_location=src_location,
+        batch_size=batch_size,
+        num_workers=num_workers,
+        num_nodes=num_nodes,
+        accelerator=accelerator,
+        devices=devices,
+        out_path=out_path,
+        inference_set=inference_set
+    )

rcsb_embedding_model/inference/chain_inference.py

@@ -1,26 +1,30 @@
 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.dataset.residue_embedding_from_tensor_file import ResidueEmbeddingFromTensorFile
 from rcsb_embedding_model.modules.chain_module import ChainModule
-from rcsb_embedding_model.types.api_types import Accelerator, Devices, OptionalPath
+from rcsb_embedding_model.types.api_types import Accelerator, Devices, OptionalPath, FileOrStreamTuple, SrcLocation
 from rcsb_embedding_model.utils.data import collate_seq_embeddings
 from rcsb_embedding_model.writer.batch_writer import CsvBatchWriter
 
 
 def predict(
-        csv_file: FileText,
+        src_stream: FileOrStreamTuple,
+        src_location: SrcLocation = SrcLocation.local,
         batch_size: int = 1,
         num_workers: int = 0,
         num_nodes: int = 1,
         accelerator: Accelerator = Accelerator.auto,
         devices: Devices = 'auto',
-        out_path: OptionalPath = None
+        out_path: OptionalPath = None,
+        inference_set=None
 ):
-    inference_set = ResidueEmbeddingFromCSV(
-        csv_file=csv_file
-    )
+
+    if inference_set is None:
+        inference_set = ResidueEmbeddingFromTensorFile(
+            src_stream=src_stream,
+            src_location=src_location
+        )
 
     inference_dataloader = DataLoader(
         dataset=inference_set,

rcsb_embedding_model/inference/esm_inference.py

@@ -1,17 +1,20 @@
 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.dataset.esm_prot_from_structure import EsmProtFromStructure
+from rcsb_embedding_model.dataset.esm_prot_from_chain import EsmProtFromChain
 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.types.api_types import StructureFormat, Accelerator, Devices, OptionalPath, StructureLocation, SrcProteinFrom, FileOrStreamTuple, SrcLocation
 from rcsb_embedding_model.writer.batch_writer import TensorBatchWriter
 
 
 def predict(
-        csv_file: FileText,
+        src_stream: FileOrStreamTuple,
         src_location: SrcLocation = SrcLocation.local,
-        src_format: SrcFormat = SrcFormat.mmcif,
+        src_from: SrcProteinFrom = SrcProteinFrom.chain,
+        structure_location: StructureLocation = StructureLocation.local,
+        structure_format: StructureFormat = StructureFormat.mmcif,
+        min_res_n: int = 0,
         batch_size: int = 1,
         num_workers: int = 0,
         num_nodes: int = 1,
@@ -20,10 +23,17 @@ def predict(
         out_path: OptionalPath = None
 ):
 
-    inference_set = EsmProtFromCsv(
-        csv_file=csv_file,
+    inference_set = EsmProtFromChain(
+        src_stream=src_stream,
         src_location=src_location,
-        src_format=src_format
+        structure_location=structure_location,
+        structure_format=structure_format
+    ) if src_from == SrcProteinFrom.chain else EsmProtFromStructure(
+        src_stream=src_stream,
+        src_location=src_location,
+        structure_location=structure_location,
+        structure_format=structure_format,
+        min_res_n=min_res_n
     )
 
     inference_dataloader = DataLoader(

rcsb_embedding_model/inference/structure_inference.py

@@ -1,30 +1,40 @@
 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.dataset.esm_prot_from_structure import EsmProtFromStructure
+from rcsb_embedding_model.dataset.esm_prot_from_chain import EsmProtFromChain
+from rcsb_embedding_model.modules.structure_module import StructureModule
+from rcsb_embedding_model.types.api_types import StructureFormat, Accelerator, Devices, OptionalPath, StructureLocation, SrcProteinFrom, FileOrStreamTuple, SrcLocation
 from rcsb_embedding_model.writer.batch_writer import DataFrameStorage
 
 
 def predict(
-        csv_file: FileText,
+        src_stream: FileOrStreamTuple,
         src_location: SrcLocation = SrcLocation.local,
-        src_format: SrcFormat = SrcFormat.mmcif,
+        src_from: SrcProteinFrom = SrcProteinFrom.chain,
+        structure_location: StructureLocation = StructureLocation.local,
+        structure_format: StructureFormat = StructureFormat.mmcif,
+        min_res_n: int = 0,
         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
+        out_df_name: str = None
 ):
 
-    inference_set = EsmProtFromCsv(
-        csv_file=csv_file,
+    inference_set = EsmProtFromChain(
+        src_stream=src_stream,
         src_location=src_location,
-        src_format=src_format
+        structure_location=structure_location,
+        structure_format=structure_format
+    ) if src_from == SrcProteinFrom.chain else EsmProtFromStructure(
+        src_stream=src_stream,
+        src_location=src_location,
+        structure_location=structure_location,
+        structure_format=structure_format,
+        min_res_n=min_res_n
     )
 
     inference_dataloader = DataLoader(
@@ -34,8 +44,8 @@ def predict(
         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
+    module = StructureModule()
+    inference_writer = DataFrameStorage(out_path, out_df_name) if out_path is not None and out_df_name is not None else None
     trainer = Trainer(
         callbacks=[inference_writer] if inference_writer is not None else None,
         num_nodes=num_nodes,

rcsb_embedding_model/rcsb_structure_embedding.py

@@ -3,7 +3,7 @@ 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 rcsb_embedding_model.types.api_types import StreamSrc, SrcFormat
+from rcsb_embedding_model.types.api_types import StreamSrc, StructureFormat
 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
 
@@ -42,22 +42,22 @@ class RcsbStructureEmbedding:
     def structure_embedding(
             self,
             src_structure: StreamSrc,
-            src_format: SrcFormat = SrcFormat.mmcif,
+            structure_format: StructureFormat = StructureFormat.mmcif,
             chain_id: str = None,
             assembly_id: str = None
     ):
-        res_embedding = self.residue_embedding(src_structure, src_format, chain_id, assembly_id)
+        res_embedding = self.residue_embedding(src_structure, structure_format, chain_id, assembly_id)
         return self.aggregator_embedding(res_embedding)
 
     def residue_embedding(
             self,
             src_structure: StreamSrc,
-            src_format: SrcFormat = SrcFormat.mmcif,
+            structure_format: StructureFormat = StructureFormat.mmcif,
             chain_id: str = None,
             assembly_id: str = None
     ):
         self.__check_residue_embedding()
-        structure = get_structure_from_src(src_structure, src_format, chain_id, assembly_id)
+        structure = get_structure_from_src(src_structure, structure_format, chain_id, assembly_id)
         embedding_ch = []
         for atom_ch in chain_iter(structure):
             atom_res = atom_ch[filter_amino_acids(atom_ch)]

rcsb_embedding_model/types/api_types.py

@@ -1,16 +1,23 @@
 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])
+from typer import FileText
+
+StreamSrc = NewType('StreamSrc', Union[FileText, IO])
+StreamTuple = NewType('StreamTuple', Union[
+    Tuple[str, StreamSrc, str, str],
+    Tuple[str, StreamSrc, str],
+    Tuple[str, str]
+])
+FileOrStreamTuple = NewType('FileOrStreamTuple', Union[FileText, StreamTuple])
 
 Devices = NewType('Devices', Union[int, List[int], "auto"])
 
-OptionalPath = NewType('OptionalPath', Optional[PathLike])
+EmbeddingPath = Union[str, FileText]
+OptionalPath = NewType('OptionalPath', Optional[FileText])
 
 
-class SrcFormat(str, Enum):
+class StructureFormat(str, Enum):
     pdb = "pdb"
     mmcif = "mmcif"
     bciff = "binarycif"
@@ -25,5 +32,20 @@ class Accelerator(str, Enum):
 
 
 class SrcLocation(str, Enum):
+    local = "local"
+    stream = "stream"
+
+
+class StructureLocation(str, Enum):
     local = "local"
     remote = "remote"
+
+
+class SrcProteinFrom(str, Enum):
+    chain = "chain"
+    structure = "structure"
+
+
+class SrcAssemblyFrom(str, Enum):
+    assembly = "assembly"
+    structure = "structure"

rcsb_embedding_model/utils/data.py

@@ -44,4 +44,34 @@ def stringio_from_url(url):
         print(f"Error fetching URL: {e}")
         return None
 
+def concatenate_tensors(file_list, max_residues, dim=0):
+    """
+    Concatenates a list of tensors stored in individual files along a specified dimension.
+
+    Args:
+        file_list (list of str): List of file paths to tensor files.
+        max_residues (int): Maximum number of residues allowed in the assembly
+        dim (int): The dimension along which to concatenate the tensors. Default is 0.
 
+    Returns:
+        torch.Tensor: The concatenated tensor.
+    """
+    tensors = []
+    total_residues = 0
+    for file in file_list:
+        try:
+            tensor = torch.load(
+                file,
+                map_location=torch.device('cpu')
+            )
+            total_residues += tensor.shape[0]
+            tensors.append(tensor)
+        except Exception as e:
+            continue
+        if total_residues > max_residues:
+            break
+    if tensors and len(tensors) > 0:
+        tensor_cat = torch.cat(tensors, dim=dim)
+        return tensor_cat
+    else:
+        raise ValueError("No valid tensors were loaded to concatenate.")

rcsb_embedding_model/utils/structure_parser.py

@@ -1,32 +1,62 @@
-
-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 biotite.structure import filter_amino_acids, chain_iter, get_chains, get_residues, AtomArray
+from biotite.structure.io.pdb import PDBFile, get_structure as get_pdb_structure, get_assembly as get_pdb_assembly, list_assemblies as list_pdb_assemblies
+from biotite.structure.io.pdbx import CIFFile, get_structure, get_assembly, BinaryCIFFile, list_assemblies
 
 
 def get_structure_from_src(
         src_structure,
-        src_format="mmcif",
+        structure_format="mmcif",
         chain_id=None,
         assembly_id=None
 ):
-    if src_format == "pdb":
+    if structure_format == "pdb":
         pdb_file = PDBFile.read(src_structure)
         structure = __get_pdb_structure(pdb_file, assembly_id)
-    elif src_format == "mmcif":
+    elif structure_format == "mmcif":
         cif_file = CIFFile.read(src_structure)
         structure = __get_structure(cif_file, assembly_id)
-    elif src_format == "binarycif":
+    elif structure_format == "binarycif":
         cif_file = BinaryCIFFile.read(src_structure)
         structure = __get_structure(cif_file, assembly_id)
     else:
-        raise RuntimeError(f"Unknown file format {src_format}")
+        raise RuntimeError(f"Unknown file format {structure_format}")
 
     if chain_id is not None:
-        structure = structure[structure.chain_id == chain_id]
+        return structure[structure.chain_id == chain_id]
 
     return structure
 
 
+def get_protein_chains(structure, min_res_n=0):
+    chain_ids = []
+    for atom_ch in chain_iter(structure):
+        atom_res = atom_ch[filter_amino_acids(atom_ch)]
+        if len(atom_res) > 0 and len(get_residues(atom_res)) > min_res_n:
+            chain_ids.append(str(get_chains(atom_res)[0]))
+    return tuple(chain_ids)
+
+
+def get_assemblies(src_structure, structure_format="mmcif"):
+    if structure_format == "pdb":
+        return tuple(list_pdb_assemblies(PDBFile.read(src_structure)))
+    elif structure_format == "mmcif":
+        return tuple(list_assemblies(CIFFile.read(src_structure)).keys())
+    elif structure_format == "binarycif":
+        return tuple(list_assemblies(BinaryCIFFile.read(src_structure)))
+    else:
+        raise RuntimeError(f"Unknown file format {structure_format}")
+
+
+def rename_atom_ch(atom_ch, ch="A"):
+    renamed_atom_ch = AtomArray(len(atom_ch))
+    n = 0
+    for atom in atom_ch:
+        atom.chain_id = ch
+        renamed_atom_ch[n] = atom
+        n += 1
+    return renamed_atom_ch
+
+
 def __get_pdb_structure(pdb_file, assembly_id=None):
     return get_pdb_structure(
         pdb_file,

rcsb_embedding_model/utils/structure_provider.py

@@ -0,0 +1,27 @@
+from rcsb_embedding_model.utils.structure_parser import get_structure_from_src
+
+
+class StructureProvider:
+
+    def __init__(self):
+        self.__src_name = None
+        self.__structure = None
+
+    def get_structure(
+        self,
+        src_name,
+        src_structure,
+        structure_format="mmcif",
+        chain_id=None,
+        assembly_id=None
+    ):
+        if src_name != self.__src_name:
+            self.__src_name = src_name
+            self.__structure = get_structure_from_src(
+                src_structure=src_structure,
+                structure_format=structure_format,
+                assembly_id=assembly_id
+            )
+        if chain_id is not None:
+            return self.__structure[self.__structure.chain_id == chain_id]
+        return self.__structure

{rcsb_embedding_model-0.0.7.dist-info → rcsb_embedding_model-0.0.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.7
+Version: 0.0.8
 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
@@ -18,7 +18,7 @@ Description-Content-Type: text/markdown
 
 # RCSB Embedding Model
 
-**Version** 0.0.7
+**Version** 0.0.8
 
 
 ## Overview
@@ -60,14 +60,17 @@ If you are interested in training the model with a new dataset, visit the [rcsb-
 
 ### 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
+    # 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 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
+    # 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.
-    inference chain-embedding --src-file results/residue_embeddings.csv --output-path results/chain_embeddings --batch-size 4
+    # 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
 
@@ -85,7 +88,7 @@ If you are interested in training the model with a new dataset, visit the [rcsb-
     # Aggregate to structure-level embedding
     struct_emb = model.aggregator_embedding(res_emb)
 
-See the examples directory for complete scripts.
+See the examples and tests directories for more use cases.
 
 ---
 
@@ -117,7 +120,7 @@ The aggregation component consists of six transformer encoder layers, each with
 
 ## 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)
+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)
 
 ---
 

rcsb_embedding_model-0.0.8.dist-info/RECORD

@@ -0,0 +1,29 @@
+rcsb_embedding_model/__init__.py,sha256=r3gLdeBIXkQEQA_K6QcRPO-TtYuAQSutk6pXRUE_nas,120
+rcsb_embedding_model/rcsb_structure_embedding.py,sha256=dKp9hXQO0JAnO4SEfjJ_mG_jHu3UxAPguv6jkOjp-BI,4487
+rcsb_embedding_model/cli/args_utils.py,sha256=7nP2q8pL5dWK_U7opxtWmoFcYVwasky6elHk-dASFaI,165
+rcsb_embedding_model/cli/inference.py,sha256=KPZLqznbxZE_CBCGigUGg7yOfGsi8ID4aWMTExniRj4,11464
+rcsb_embedding_model/dataset/esm_prot_from_chain.py,sha256=dBD2N0Y-GoN6p3z2yLnOvv6JGn-skAxwgbOYhXKDngc,3487
+rcsb_embedding_model/dataset/esm_prot_from_structure.py,sha256=kOqgHfHjiym5InaAgpgMmHBgCAPEqW88PCoHHQy0ROI,2490
+rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py,sha256=d8C7HRJBZWuOKhPQpihv1koT4aIvyt5QN2yndC2ABuE,2842
+rcsb_embedding_model/dataset/residue_assembly_embedding_from_tensor_file.py,sha256=KXiohnPjjfZEFbPZQ46HGE8eEYWrVX8bfbTz4zPlo7o,3451
+rcsb_embedding_model/dataset/residue_embedding_from_tensor_file.py,sha256=cOxT--Spkel10JJCeGlqgLXN5vNCZzPdfSxDgUSrdPI,1268
+rcsb_embedding_model/inference/assembly_inferece.py,sha256=MPssN5bsOqOU-LGwa6AKX99cv5LD43Mnbaqhuuww1Tw,2165
+rcsb_embedding_model/inference/chain_inference.py,sha256=R9gi0MZ_HaM3v9c433W_5w4suse4nJmy4SgUTHJVZLg,1713
+rcsb_embedding_model/inference/esm_inference.py,sha256=oVN4r9_6V8TS0pYoNn7GR92Xo0Zn7eBsnt_OfDSaH6g,2126
+rcsb_embedding_model/inference/structure_inference.py,sha256=QIUEo8eEc-kTSYKGdlX2rxT74huw4ZAw6U8Px9kYajE,2216
+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=3sPh33yb3Ya9r3O5vuiTfhb1WyFuhQWCQmewSbqEyG0,1076
+rcsb_embedding_model/utils/data.py,sha256=x6ca_bVdBXEAp9ugCi1rVEQ-G5nGTFKpzDKqZKpkFBE,2933
+rcsb_embedding_model/utils/model.py,sha256=rpZa-gfm3cEtbBd7UXMHrZv3x6f0AC8TJT3gtrSxr5I,852
+rcsb_embedding_model/utils/structure_parser.py,sha256=jat4SCtPHYMZ6JJR-T7lPQoMbT_E8CwYSGDNSZjG86U,2697
+rcsb_embedding_model/utils/structure_provider.py,sha256=eWtxjkPpmRfmil_DKR1J6miaXR3lQ28DF5O0qrqSgGA,786
+rcsb_embedding_model/writer/batch_writer.py,sha256=ekgzFZyoKpcnZ3IDP9hfOWBpuHxUQ31P35ViDAi-Edw,2843
+rcsb_embedding_model-0.0.8.dist-info/METADATA,sha256=XvNb99X9GWdMEdz-A_o-ngTxlfiWrr8KMvjKh_rk3x0,5366
+rcsb_embedding_model-0.0.8.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+rcsb_embedding_model-0.0.8.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
+rcsb_embedding_model-0.0.8.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
+rcsb_embedding_model-0.0.8.dist-info/RECORD,,

rcsb_embedding_model/dataset/esm_prot_from_csv.py

@@ -1,90 +0,0 @@
-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

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

rcsb_embedding_model-0.0.7.dist-info/RECORD

@@ -1,24 +0,0 @@
-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,,