rcsb-embedding-model 0.0.29__tar.gz → 0.0.31__tar.gz

{rcsb_embedding_model-0.0.29 → rcsb_embedding_model-0.0.31}/.gitignore

@@ -2,4 +2,5 @@
 /rcsb-embedding-model.iml
 /dist/
 /.pypi.rc
-__pycache__
+__pycache__
+/tests/resources/tmp

{rcsb_embedding_model-0.0.29 → rcsb_embedding_model-0.0.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.29
+Version: 0.0.31
 Summary: Protein Embedding Model for Structure Search
 Project-URL: Homepage, https://github.com/rcsb/rcsb-embedding-model
 Project-URL: Issues, https://github.com/rcsb/rcsb-embedding-model/issues
@@ -11,6 +11,7 @@ Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Requires-Dist: esm>=3.2.0
+Requires-Dist: importlib-metadata>=8.7.0
 Requires-Dist: lightning>=2.5.0
 Requires-Dist: typer>=0.15.0
 Description-Content-Type: text/markdown

{rcsb_embedding_model-0.0.29 → rcsb_embedding_model-0.0.31}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.29"
+version = "0.0.31"
 authors = [
     { name="Joan Segura", email="joan.segura@rcsb.org" },
 ]
@@ -14,6 +14,7 @@ classifiers = [
 license = "BSD-3-Clause"
 license-files = ["LICEN[CS]E*"]
 dependencies=[
+    "importlib-metadata >= 8.7.0",
     "esm >= 3.2.0",
     "lightning >= 2.5.0",
     "typer >= 0.15.0"

rcsb_embedding_model-0.0.31/src/rcsb_embedding_model/__init__.py

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

{rcsb_embedding_model-0.0.29 → rcsb_embedding_model-0.0.31}/src/rcsb_embedding_model/cli/inference.py

@@ -3,6 +3,7 @@ from typing import Annotated, List
 
 import typer
 
+from rcsb_embedding_model import __version__
 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, OutFormat
@@ -437,5 +438,22 @@ def complete_embedding(
     )
 
 
+def version_callback(value: bool):
+    if value:
+        typer.echo(f"{__version__}")
+        raise typer.Exit()
+
+@app.callback()
+def main(
+        version: bool = typer.Option(
+            None,
+            "--version",
+            callback=version_callback,
+            is_eager=True,
+            help="Show the version and exit",
+        )
+):
+    pass
+
 if __name__ == "__main__":
     app()

{rcsb_embedding_model-0.0.29 → rcsb_embedding_model-0.0.31}/src/rcsb_embedding_model/dataset/residue_assembly_embedding_from_tensor_file.py

@@ -1,3 +1,4 @@
+import argparse
 import sys
 
 import pandas as pd
@@ -76,12 +77,17 @@ class ResidueAssemblyEmbeddingFromTensorFile(Dataset):
 
 if __name__ == "__main__":
 
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--file_list', type=argparse.FileType('r'), required=True)
+    parser.add_argument('--res_embeddings_path', required=True)
+    args = parser.parse_args()
+
     dataset = ResidueAssemblyEmbeddingFromTensorFile(
-        src_stream="/Users/joan/tmp/assembly-test.csv",
-        res_embedding_location="/Users/joan/tmp",
+        src_stream=args.file_list,
+        res_embedding_location=args.res_embeddings_path,
         src_location=SrcLocation.file,
-        structure_location=StructureLocation.local,
-        structure_format=StructureFormat.mmcif
+        structure_location=StructureLocation.remote,
+        structure_format=StructureFormat.bciff
     )
 
     dataloader = DataLoader(

{rcsb_embedding_model-0.0.29 → rcsb_embedding_model-0.0.31}/src/rcsb_embedding_model/utils/structure_parser.py

@@ -32,7 +32,7 @@ def get_protein_chains(structure, min_res_n=0):
     for atom_ch in chain_iter(structure):
         atom_res = atom_ch[filter_polymer(atom_ch)]
         atom_res = atom_res[filter_amino_acids(atom_res)]
-        if len(atom_res) > 0 and len(get_residues(atom_res)) > min_res_n:
+        if len(atom_res) > 0 and len(get_residues(atom_res)[0]) > min_res_n:
             chain_ids.append(str(get_chains(atom_res)[0]))
     return tuple(chain_ids)
 

rcsb_embedding_model-0.0.31/tests/resources/src_stream/assembly-complete-test.csv

@@ -0,0 +1,7 @@
+1A21,https://models.rcsb.org/1A21.bcif.gz,1,1A21-1
+1A21,https://models.rcsb.org/1A21.bcif.gz,2,1A21-2
+1A3J,https://models.rcsb.org/1A3J.bcif.gz,1,1A3J-1
+1A3X,https://models.rcsb.org/1A3X.bcif.gz,1,1A3X-1
+1A3X,https://models.rcsb.org/1A3X.bcif.gz,2,1A3X-2
+1AIV,https://models.rcsb.org/1AIV.bcif.gz,1,1AIV-1
+1AU1,https://models.rcsb.org/1AU1.bcif.gz,1,1AU1-1

rcsb_embedding_model-0.0.31/tests/resources/src_stream/instance-complete-test.csv

@@ -0,0 +1,10 @@
+1A21,https://models.rcsb.org/1A21.bcif.gz,A,1A21.A
+1A21,https://models.rcsb.org/1A21.bcif.gz,B,1A21.B
+1A3J,https://models.rcsb.org/1A3J.bcif.gz,A,1A3J.A
+1A3J,https://models.rcsb.org/1A3J.bcif.gz,B,1A3J.B
+1A3J,https://models.rcsb.org/1A3J.bcif.gz,C,1A3J.C
+1A3X,https://models.rcsb.org/1A3X.bcif.gz,A,1A3X.A
+1A3X,https://models.rcsb.org/1A3X.bcif.gz,B,1A3X.B
+1AIV,https://models.rcsb.org/1AIV.bcif.gz,A,1AIV.A
+1AU1,https://models.rcsb.org/1AU1.bcif.gz,A,1AU1.A
+1AU1,https://models.rcsb.org/1AU1.bcif.gz,B,1AU1.B

rcsb_embedding_model-0.0.31/tests/test_cli_inference.py

@@ -0,0 +1,53 @@
+
+import os
+import shutil
+import unittest
+
+from rcsb_embedding_model.types.api_types import OutFormat, StructureLocation, StructureFormat, Accelerator
+
+
+
+
+
+class TestCliInference(unittest.TestCase):
+    __test_path = os.path.dirname(__file__)
+
+    def test_complete_inference(self):
+        _remove_files_in_directory(f"{self.__test_path}/resources/tmp")
+        from rcsb_embedding_model.cli.inference import complete_embedding
+        complete_embedding(
+            src_chain_file=f"{self.__test_path}/resources/src_stream/instance-complete-test.csv",
+            src_assembly_file=f"{self.__test_path}/resources/src_stream/assembly-complete-test.csv",
+            output_res_path=f"{self.__test_path}/resources/tmp",
+            output_chain_path=f"{self.__test_path}/resources/tmp",
+            output_assembly_path=f"{self.__test_path}/resources/tmp",
+            output_format=OutFormat.grouped,
+            output_chain_name="instance-inference",
+            output_assembly_name="assembly-inference",
+            structure_location=StructureLocation.remote,
+            structure_format=StructureFormat.bciff,
+            min_res_n=0,
+            batch_size_res=1,
+            num_workers_res=0,
+            batch_size_chain=1,
+            num_workers_chain=0,
+            batch_size_assembly=1,
+            num_workers_assembly=0,
+            num_nodes=1,
+            accelerator=Accelerator.cpu
+        )
+        self.assertTrue(os.path.exists(f"{self.__test_path}/resources/tmp/instance-inference.json.gz"))
+        self.assertTrue(os.path.exists(f"{self.__test_path}/resources/tmp/assembly-inference.json.gz"))
+
+
+def _remove_files_in_directory(directory_path):
+    os.makedirs(directory_path, exist_ok=True)
+    for filename in os.listdir(directory_path):
+        file_path = os.path.join(directory_path, filename)
+        try:
+            if os.path.isfile(file_path):
+                os.unlink(file_path)
+            elif os.path.isdir(file_path):
+                shutil.rmtree(file_path)
+        except Exception as e:
+            print(f"Failed to delete {file_path}. Reason: {e}")

rcsb_embedding_model-0.0.29/src/rcsb_embedding_model/__init__.py

@@ -1,4 +0,0 @@
-
-from rcsb_embedding_model.rcsb_structure_embedding import RcsbStructureEmbedding
-
-__all__ = ["RcsbStructureEmbedding"]