rcsb-embedding-model 0.0.16__tar.gz → 0.0.18__tar.gz

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.16
+Version: 0.0.18
 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
@@ -12,13 +12,12 @@ 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.16
+**Version** 0.0.18
 
 
 ## Overview
@@ -48,11 +47,10 @@ If you are interested in training the model with a new dataset, visit the [rcsb-
 
 **Requirements:**
 
-- Python ≥ 3.10  
-- ESM ≥ 3.2.0  
-- PyTorch ≥ 2.2.0  
-- Lightning ≥ 2.5.0  
-- Typer ≥ 0.15.0  
+- Python ≥ 3.10
+- ESM == 3.1.1
+- Lightning ≥ 2.5.0 
+- Typer ≥ 0.15.0
 
 ---
 

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/README.md

@@ -1,6 +1,6 @@
 # RCSB Embedding Model
 
-**Version** 0.0.16
+**Version** 0.0.18
 
 
 ## Overview
@@ -30,11 +30,10 @@ If you are interested in training the model with a new dataset, visit the [rcsb-
 
 **Requirements:**
 
-- Python ≥ 3.10  
-- ESM ≥ 3.2.0  
-- PyTorch ≥ 2.2.0  
-- Lightning ≥ 2.5.0  
-- Typer ≥ 0.15.0  
+- Python ≥ 3.10
+- ESM == 3.1.1
+- Lightning ≥ 2.5.0 
+- Typer ≥ 0.15.0
 
 ---
 

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.16"
+version = "0.0.18"
 authors = [
     { name="Joan Segura", email="joan.segura@rcsb.org" },
 ]
@@ -15,7 +15,6 @@ license = "BSD-3-Clause"
 license-files = ["LICEN[CS]E*"]
 dependencies=[
     "esm >= 3.2.0",
-    "torch >= 2.2.0",
     "lightning >= 2.5.0",
     "typer >= 0.15.0"
 ]

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/cli/inference.py

@@ -5,7 +5,8 @@ import typer
 
 from rcsb_embedding_model.cli.args_utils import arg_devices
 from rcsb_embedding_model.types.api_types import StructureFormat, Accelerator, SrcLocation, SrcProteinFrom, \
-    StructureLocation, SrcAssemblyFrom, SrcTensorFrom
+    StructureLocation, SrcAssemblyFrom, SrcTensorFrom, OutFormat
+from rcsb_embedding_model.utils.data import adapt_csv_to_embedding_chain_stream
 
 app = typer.Typer(
     add_completion=False
@@ -22,7 +23,7 @@ def residue_embedding(
             file_okay=True,
             dir_okay=False,
             resolve_path=True,
-            help='CSV file 4 (or 3) columns: Structure Name | Structure File Path or URL (switch structure-location) | Chain Id (asym_i for cif files. This field is required if src-from=chain) | Output Embedding Name.'
+            help='CSV file 4 columns: Structure Name | Structure File Path or URL (switch structure-location) | Chain Id (asym_i for cif files) | Output Embedding Name.'
         )],
         output_path: Annotated[typer.FileText, typer.Option(
             exists=True,
@@ -31,9 +32,12 @@ def residue_embedding(
             resolve_path=True,
             help='Output path to store predictions. Embeddings are stored as torch tensor files.'
         )],
-        src_from: Annotated[SrcProteinFrom, typer.Option(
-            help='Use specific chains or all chains in a structure.'
-        )] = SrcProteinFrom.chain,
+        output_format: Annotated[OutFormat, typer.Option(
+            help='Format of the output. Options: separated (predictions are stored in single files) or grouped (predictions are stored in a single JSON file).'
+        )] = OutFormat.separated,
+        output_name: Annotated[str, typer.Option(
+            help='File name for storing embeddings as a single JSON file. Used when output-format=grouped.'
+        )] = 'inference',
         structure_location: Annotated[StructureLocation, typer.Option(
             help='Structure file location.'
         )] = StructureLocation.local,
@@ -62,8 +66,8 @@ def residue_embedding(
     from rcsb_embedding_model.inference.esm_inference import predict
     predict(
         src_stream=src_file,
-        src_location=SrcLocation.local,
-        src_from=src_from,
+        src_location=SrcLocation.file,
+        src_from=SrcProteinFrom.chain,
         structure_location=structure_location,
         structure_format=structure_format,
         min_res_n=min_res_n,
@@ -72,6 +76,8 @@ def residue_embedding(
         num_nodes=num_nodes,
         accelerator=accelerator,
         devices=arg_devices(devices),
+        out_format=output_format,
+        out_name=output_name,
         out_path=output_path
     )
 
@@ -86,7 +92,7 @@ def structure_embedding(
             file_okay=True,
             dir_okay=False,
             resolve_path=True,
-            help='CSV file 4 (or 3) columns: Structure Name | Structure File Path or URL (switch structure-location) | Chain Id (asym_i for cif files. This field is required if src-from=chain) | Output Embedding Name.'
+            help='CSV file 4 columns: Structure Name | Structure File Path or URL (switch structure-location) | Chain Id (asym_i for cif files) | Output Embedding Name.'
         )],
         output_path: Annotated[typer.FileText, typer.Option(
             exists=True,
@@ -95,12 +101,9 @@ def structure_embedding(
             resolve_path=True,
             help='Output path to store predictions. Embeddings are stored as a single DataFrame file (see out-df-name).'
         )],
-        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_from: Annotated[SrcProteinFrom, typer.Option(
-            help='Use specific chains or all chains in a structure.'
-        )] = SrcProteinFrom.chain,
+        output_name: Annotated[str, typer.Option(
+            help='File name for storing embeddings as a single JSON file.'
+        )] = 'inference',
         structure_location: Annotated[StructureLocation, typer.Option(
             help='Structure file location.'
         )] = StructureLocation.local,
@@ -129,8 +132,8 @@ def structure_embedding(
     from rcsb_embedding_model.inference.structure_inference import predict
     predict(
         src_stream=src_file,
-        src_location=SrcLocation.local,
-        src_from=src_from,
+        src_location=SrcLocation.file,
+        src_from=SrcProteinFrom.chain,
         structure_location=structure_location,
         structure_format=structure_format,
         min_res_n=min_res_n,
@@ -140,7 +143,7 @@ def structure_embedding(
         accelerator=accelerator,
         devices=arg_devices(devices),
         out_path=output_path,
-        out_df_name=out_df_name
+        out_name=output_name
     )
 
 
@@ -154,7 +157,7 @@ def chain_embedding(
             file_okay=True,
             dir_okay=False,
             resolve_path=True,
-            help='Option 1 (src-from=file) - CSV file 2 columns: Residue Embedding Torch Tensor File | Output Embedding Name. Option 2 (src-from=structure) - CSV file 3 columns: Structure Name | Structure File Path or URL (switch structure-location) | Output Embedding Name.'
+            help='CSV file 4 columns: Structure Name | Structure File Path or URL (switch structure-location) | Chain Id (asym_i for cif files) | Output Embedding Name.'
         )],
         output_path: Annotated[typer.FileText, typer.Option(
             exists=True,
@@ -168,11 +171,14 @@ def chain_embedding(
             file_okay=False,
             dir_okay=True,
             resolve_path=True,
-            help='Path where residue level embeddings are located. This argument is required if src-from=structure.'
-        )] = None,
-        src_from: Annotated[SrcTensorFrom, typer.Option(
-            help='Use file names or all chains in a structure.'
-        )] = SrcTensorFrom.file,
+            help='Path where residue level embeddings are located.'
+        )],
+        output_format: Annotated[OutFormat, typer.Option(
+            help='Format of the output. Options: separated (predictions are stored in single files) or grouped (predictions are stored in a single JSON file).'
+        )] = OutFormat.separated,
+        output_name: Annotated[str, typer.Option(
+            help='File name for storing embeddings as a single JSON file. Used when output-format=grouped.'
+        )] = 'inference',
         structure_location: Annotated[StructureLocation, typer.Option(
             help='Structure file location.'
         )] = StructureLocation.local,
@@ -200,10 +206,10 @@ def chain_embedding(
 ):
     from rcsb_embedding_model.inference.chain_inference import predict
     predict(
-        src_stream=src_file,
+        src_stream=adapt_csv_to_embedding_chain_stream(src_file, res_embedding_location),
         res_embedding_location=res_embedding_location,
-        src_location=SrcLocation.local,
-        src_from=src_from,
+        src_location=SrcLocation.stream,
+        src_from=SrcTensorFrom.file,
         structure_location=structure_location,
         structure_format=structure_format,
         min_res_n=min_res_n,
@@ -212,7 +218,9 @@ def chain_embedding(
         num_nodes=num_nodes,
         accelerator=accelerator,
         devices=arg_devices(devices),
-        out_path=output_path
+        out_path=output_path,
+        out_format=output_format,
+        out_name=output_name
     )
 
 @app.command(
@@ -241,9 +249,12 @@ def assembly_embedding(
             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,
+        output_format: Annotated[OutFormat, typer.Option(
+            help='Format of the output. Options: separated (predictions are stored in single files) or grouped (predictions are stored in a single JSON file).'
+        )] = OutFormat.separated,
+        output_name: Annotated[str, typer.Option(
+            help='File name for storing embeddings as a single JSON file. Used when output-format=grouped.'
+        )] = 'inference',
         structure_location: Annotated[StructureLocation, typer.Option(
             help='Structure file location.'
         )] = StructureLocation.local,
@@ -276,8 +287,8 @@ def assembly_embedding(
     predict(
         src_stream=src_file,
         res_embedding_location=res_embedding_location,
-        src_location=SrcLocation.local,
-        src_from=src_from,
+        src_location=SrcLocation.file,
+        src_from=SrcAssemblyFrom.assembly,
         structure_location=structure_location,
         structure_format=structure_format,
         min_res_n=min_res_n,
@@ -287,7 +298,117 @@ def assembly_embedding(
         num_nodes=num_nodes,
         accelerator=accelerator,
         devices=arg_devices(devices),
-        out_path=output_path
+        out_path=output_path,
+        out_format=output_format,
+        out_name=output_name
+    )
+
+@app.command(
+    name="complete-embedding",
+    help="Calculate chain and assembly embeddings from structural files. Predictions are stored as csv files."
+)
+def complete_embedding(
+        src_chain_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 or URL (switch structure-location) | Chain Id (asym_i for cif files) | Output Embedding Name.'
+        )],
+        src_assembly_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 or URL (switch structure-location) | Assembly Id | 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. Embeddings are stored as a single DataFrame file (see output_name).'
+        )],
+        res_embedding_location: Annotated[typer.FileText, typer.Option(
+            exists=True,
+            file_okay=False,
+            dir_okay=True,
+            resolve_path=True,
+            help='Output path to store ESM predictions.'
+        )],
+        output_format: Annotated[OutFormat, typer.Option(
+            help='Format of the output. Options: separated (predictions are stored in single files) or grouped (predictions are stored in a single JSON file).'
+        )] = OutFormat.separated,
+        output_name: Annotated[str, typer.Option(
+            help='File name for storing embeddings as a single JSON file. Used when output-format=grouped.'
+        )] = 'inference',
+        structure_location: Annotated[StructureLocation, typer.Option(
+            help='Structure file location.'
+        )] = StructureLocation.local,
+        structure_format: Annotated[StructureFormat, typer.Option(
+            help='Structure file format.'
+        )] = 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,
+        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'])
+):
+    residue_embedding(
+        src_file=src_chain_file,
+        output_path=res_embedding_location,
+        output_format=OutFormat.separated,
+        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=devices,
+    )
+    chain_embedding(
+        src_file=src_chain_file,
+        output_path=output_path,
+        output_format=output_format,
+        output_name=f"{output_name}-chain",
+        res_embedding_location=res_embedding_location,
+        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=devices
+    )
+    assembly_embedding(
+        src_file=src_assembly_file,
+        output_path=output_path,
+        output_format=output_format,
+        output_name=f"{output_name}-assembly",
+        res_embedding_location=res_embedding_location,
+        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=devices
     )
 
 

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/dataset/esm_prot_from_chain.py

@@ -27,7 +27,7 @@ class EsmProtFromChain(Dataset):
     def __init__(
         self,
         src_stream,
-        src_location=SrcLocation.local,
+        src_location=SrcLocation.file,
         structure_location=StructureLocation.local,
         structure_format=StructureFormat.mmcif,
         structure_provider=StructureProvider()
@@ -70,6 +70,7 @@ class EsmProtFromChain(Dataset):
         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
+        return None
 
 
 if __name__ == '__main__':

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/dataset/esm_prot_from_structure.py

@@ -19,7 +19,7 @@ class EsmProtFromStructure(EsmProtFromChain):
     def __init__(
             self,
             src_stream,
-            src_location=SrcLocation.local,
+            src_location=SrcLocation.file,
             structure_location=StructureLocation.local,
             structure_format=StructureFormat.mmcif,
             min_res_n=0,

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py

@@ -21,7 +21,7 @@ class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile
             self,
             src_stream,
             res_embedding_location,
-            src_location=SrcLocation.local,
+            src_location=SrcLocation.file,
             structure_location=StructureLocation.local,
             structure_format=StructureFormat.mmcif,
             min_res_n=0,
@@ -63,6 +63,6 @@ class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile
             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(structure=structure, structure_format=self.structure_format):
-                assemblies.append((src_name, src_structure, str(assembly_id), f"{item_name}.{assembly_id}"))
+                assemblies.append((src_name, src_structure, str(assembly_id), f"{item_name}-{assembly_id}"))
 
         return tuple(assemblies)

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/dataset/residue_assembly_embedding_from_tensor_file.py

@@ -22,7 +22,7 @@ class ResidueAssemblyEmbeddingFromTensorFile(Dataset):
             self,
             src_stream,
             res_embedding_location,
-            src_location=SrcLocation.local,
+            src_location=SrcLocation.file,
             structure_location=StructureLocation.local,
             structure_format=StructureFormat.mmcif,
             min_res_n=0,
@@ -79,7 +79,7 @@ if __name__ == "__main__":
     dataset = ResidueAssemblyEmbeddingFromTensorFile(
         src_stream="/Users/joan/tmp/assembly-test.csv",
         res_embedding_location="/Users/joan/tmp",
-        src_location=SrcLocation.local,
+        src_location=SrcLocation.file,
         structure_location=StructureLocation.local,
         structure_format=StructureFormat.mmcif
     )

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/dataset/residue_embedding_from_structure.py

@@ -21,7 +21,7 @@ class ResidueEmbeddingFromStructure(ResidueEmbeddingFromTensorFile):
             self,
             src_stream,
             res_embedding_location,
-            src_location=SrcLocation.local,
+            src_location=SrcLocation.file,
             structure_location=StructureLocation.local,
             structure_format=StructureFormat.mmcif,
             min_res_n=0,

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/dataset/residue_embedding_from_tensor_file.py

@@ -15,7 +15,7 @@ class ResidueEmbeddingFromTensorFile(Dataset):
     def __init__(
             self,
             src_stream,
-            src_location=SrcLocation.local
+            src_location=SrcLocation.file
     ):
         super().__init__()
         self.src_location = src_location

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/inference/assembly_inferece.py

@@ -2,14 +2,15 @@ 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.types.api_types import FileOrStreamTuple, SrcLocation, Accelerator, Devices, OptionalPath, \
+    EmbeddingPath, StructureLocation, StructureFormat, SrcAssemblyFrom, OutFormat
 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_location: SrcLocation = SrcLocation.file,
         src_from: SrcAssemblyFrom = SrcAssemblyFrom.assembly,
         structure_location: StructureLocation = StructureLocation.local,
         structure_format: StructureFormat = StructureFormat.mmcif,
@@ -20,6 +21,8 @@ def predict(
         num_nodes: int = 1,
         accelerator: Accelerator = Accelerator.auto,
         devices: Devices = 'auto',
+        out_format: OutFormat = OutFormat.separated,
+        out_name: str = 'inference',
         out_path: OptionalPath = None
 ):
     inference_set = ResidueAssemblyEmbeddingFromTensorFile(
@@ -48,6 +51,8 @@ def predict(
         num_nodes=num_nodes,
         accelerator=accelerator,
         devices=devices,
+        out_format=out_format,
+        out_name=out_name,
         out_path=out_path,
         inference_set=inference_set
     )

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/inference/chain_inference.py

@@ -5,15 +5,15 @@ from rcsb_embedding_model.dataset.residue_embedding_from_structure import Residu
 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, FileOrStreamTuple, SrcLocation, \
-    SrcTensorFrom, StructureLocation, StructureFormat
+    SrcTensorFrom, StructureLocation, StructureFormat, OutFormat
 from rcsb_embedding_model.utils.data import collate_seq_embeddings
-from rcsb_embedding_model.writer.batch_writer import CsvBatchWriter
+from rcsb_embedding_model.writer.batch_writer import CsvBatchWriter, JsonStorage
 
 
 def predict(
         src_stream: FileOrStreamTuple,
         res_embedding_location: OptionalPath = None,
-        src_location: SrcLocation = SrcLocation.local,
+        src_location: SrcLocation = SrcLocation.file,
         src_from: SrcTensorFrom = SrcTensorFrom.file,
         structure_location: StructureLocation = StructureLocation.local,
         structure_format: StructureFormat = StructureFormat.mmcif,
@@ -23,6 +23,8 @@ def predict(
         num_nodes: int = 1,
         accelerator: Accelerator = Accelerator.auto,
         devices: Devices = 'auto',
+        out_format: OutFormat = OutFormat.separated,
+        out_name: str = 'inference',
         out_path: OptionalPath = None,
         inference_set=None
 ):
@@ -51,13 +53,13 @@ def predict(
     )
 
     module = ChainModule()
-
-    inference_writer = CsvBatchWriter(out_path) if out_path is not None else None
+    inference_writer = (JsonStorage(out_path, out_name) if out_format == OutFormat.grouped else 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
+        devices=devices,
+        logger=False
     )
 
     prediction = trainer.predict(

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/inference/esm_inference.py

@@ -4,13 +4,14 @@ from lightning import Trainer
 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 StructureFormat, Accelerator, Devices, OptionalPath, StructureLocation, SrcProteinFrom, FileOrStreamTuple, SrcLocation
-from rcsb_embedding_model.writer.batch_writer import TensorBatchWriter
+from rcsb_embedding_model.types.api_types import StructureFormat, Accelerator, Devices, OptionalPath, StructureLocation, \
+    SrcProteinFrom, FileOrStreamTuple, SrcLocation, OutFormat
+from rcsb_embedding_model.writer.batch_writer import TensorBatchWriter, JsonStorage
 
 
 def predict(
         src_stream: FileOrStreamTuple,
-        src_location: SrcLocation = SrcLocation.local,
+        src_location: SrcLocation = SrcLocation.file,
         src_from: SrcProteinFrom = SrcProteinFrom.chain,
         structure_location: StructureLocation = StructureLocation.local,
         structure_format: StructureFormat = StructureFormat.mmcif,
@@ -20,6 +21,8 @@ def predict(
         num_nodes: int = 1,
         accelerator: Accelerator = Accelerator.auto,
         devices: Devices = 'auto',
+        out_format: OutFormat = OutFormat.separated,
+        out_name: str = 'inference',
         out_path: OptionalPath = None
 ):
 
@@ -44,12 +47,13 @@ def predict(
     )
 
     module = EsmModule()
-    inference_writer = TensorBatchWriter(out_path) if out_path is not None else None
+    inference_writer = (JsonStorage(out_path, out_name) if out_format == OutFormat.grouped else 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
+        devices=devices,
+        logger=False
     )
 
     prediction = trainer.predict(

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/inference/structure_inference.py

@@ -4,13 +4,14 @@ from lightning import Trainer
 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
+from rcsb_embedding_model.types.api_types import StructureFormat, Accelerator, Devices, OptionalPath, StructureLocation, \
+    SrcProteinFrom, FileOrStreamTuple, SrcLocation
+from rcsb_embedding_model.writer.batch_writer import JsonStorage
 
 
 def predict(
         src_stream: FileOrStreamTuple,
-        src_location: SrcLocation = SrcLocation.local,
+        src_location: SrcLocation = SrcLocation.file,
         src_from: SrcProteinFrom = SrcProteinFrom.chain,
         structure_location: StructureLocation = StructureLocation.local,
         structure_format: StructureFormat = StructureFormat.mmcif,
@@ -20,8 +21,8 @@ def predict(
         num_nodes: int = 1,
         accelerator: Accelerator = Accelerator.auto,
         devices: Devices = 'auto',
-        out_path: OptionalPath = None,
-        out_df_name: str = None
+        out_name: str = 'inference',
+        out_path: OptionalPath = None
 ):
 
     inference_set = EsmProtFromChain(
@@ -45,12 +46,13 @@ def predict(
     )
 
     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
+    inference_writer = JsonStorage(out_path, out_name) if out_path is not None and out_name 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
+        devices=devices,
+        logger=False
     )
 
     prediction = trainer.predict(

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/modules/esm_module.py

@@ -1,4 +1,5 @@
 from esm.sdk.api import SamplingConfig
+from esm.sdk import batch_executor
 from lightning import LightningModule
 
 from rcsb_embedding_model.utils.model import get_residue_model
@@ -14,11 +15,13 @@ class EsmModule(LightningModule):
 
     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(
+        def __batch_embedding(esm_prot):
+            return 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)
+        with batch_executor() as executor:
+            prot_embeddings = executor.execute_batch(
+                user_func=__batch_embedding,
+                esm_prot=[esm_prot for esm_prot, name in prot_batch]
+            )
+        return tuple(prot_embeddings), tuple([name for esm_prot, name in prot_batch])

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/types/api_types.py

@@ -32,7 +32,7 @@ class Accelerator(str, Enum):
 
 
 class SrcLocation(str, Enum):
-    local = "local"
+    file = "file"
     stream = "stream"
 
 
@@ -54,3 +54,7 @@ class SrcAssemblyFrom(str, Enum):
 class SrcTensorFrom(str, Enum):
     file = "file"
     structure = "structure"
+
+class OutFormat(str, Enum):
+    separated = "separated"
+    grouped = "grouped"

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/utils/data.py

@@ -1,3 +1,4 @@
+import os
 from io import StringIO
 
 import requests
@@ -76,3 +77,9 @@ def concatenate_tensors(file_list, max_residues, dim=0):
         return tensor_cat
     else:
         raise ValueError("No valid tensors were loaded to concatenate.")
+
+def adapt_csv_to_embedding_chain_stream(src_file, res_embedding_location):
+    def __parse_row(row):
+        r = row.split(",")
+        return os.path.join(res_embedding_location, f"{r[0]}.{r[2]}.pt"), f"{r[0]}.{r[2]}"
+    return tuple([__parse_row(r) for r in open(src_file)])

{rcsb_embedding_model-0.0.16 → rcsb_embedding_model-0.0.18}/src/rcsb_embedding_model/writer/batch_writer.py

@@ -111,3 +111,21 @@ class DataFrameStorage(CoreBatchWriter, ABC):
             f"{self.out_path}/{self.df_id}.pkl.gz",
             compression='gzip'
         )
+
+
+class JsonStorage(DataFrameStorage, ABC):
+    def __init__(
+            self,
+            output_path,
+            df_id,
+            postfix="pkl",
+            write_interval="batch"
+    ):
+        super().__init__(output_path, df_id, postfix, write_interval)
+
+    def on_predict_end(self, trainer, pl_module):
+        self.embedding.to_json(
+            f"{self.out_path}/{self.df_id}.json.gz",
+            orient='records',
+            compression='gzip'
+        )