rcsb-embedding-model 0.0.11__tar.gz → 0.0.13__tar.gz

{rcsb_embedding_model-0.0.11 → rcsb_embedding_model-0.0.13}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.11
+Version: 0.0.13
 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.11
+**Version** 0.0.13
 
 
 ## Overview

{rcsb_embedding_model-0.0.11 → rcsb_embedding_model-0.0.13}/README.md

@@ -1,6 +1,6 @@
 # RCSB Embedding Model
 
-**Version** 0.0.11
+**Version** 0.0.13
 
 
 ## Overview

{rcsb_embedding_model-0.0.11 → rcsb_embedding_model-0.0.13}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.11"
+version = "0.0.13"
 authors = [
     { name="Joan Segura", email="joan.segura@rcsb.org" },
 ]

{rcsb_embedding_model-0.0.11 → rcsb_embedding_model-0.0.13}/src/rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py

@@ -37,7 +37,7 @@ class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile
         super().__init__(
             src_stream=self.__get_assemblies(src_stream),
             res_embedding_location=res_embedding_location,
-            src_location=src_location,
+            src_location=SrcLocation.stream,
             structure_location=structure_location,
             structure_format=structure_format,
             min_res_n=min_res_n,
@@ -60,9 +60,9 @@ class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile
         )).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
+            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):
+            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}"))
 
         return tuple(assemblies)

{rcsb_embedding_model-0.0.11 → rcsb_embedding_model-0.0.13}/src/rcsb_embedding_model/utils/data.py

@@ -44,6 +44,7 @@ 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.

{rcsb_embedding_model-0.0.11 → rcsb_embedding_model-0.0.13}/src/rcsb_embedding_model/utils/structure_parser.py

@@ -37,13 +37,13 @@ def get_protein_chains(structure, min_res_n=0):
     return tuple(chain_ids)
 
 
-def get_assemblies(src_structure, structure_format="mmcif"):
+def get_assemblies(structure, structure_format="mmcif"):
     if structure_format == "pdb":
-        return tuple(list_pdb_assemblies(PDBFile.read(src_structure)))
+        return tuple(list_pdb_assemblies(PDBFile.read(structure)))
     elif structure_format == "mmcif":
-        return tuple(list_assemblies(CIFFile.read(src_structure)).keys())
+        return tuple(list_assemblies(CIFFile.read(structure)).keys())
     elif structure_format == "binarycif":
-        return tuple(list_assemblies(BinaryCIFFile.read(src_structure)))
+        return tuple(list_assemblies(BinaryCIFFile.read(structure)))
     else:
         raise RuntimeError(f"Unknown file format {structure_format}")
 

{rcsb_embedding_model-0.0.11 → rcsb_embedding_model-0.0.13}/tests/test_remote_inference.py

@@ -1,11 +1,14 @@
+import os.path
 import unittest
 
 from rcsb_embedding_model.types.api_types import SrcLocation, SrcProteinFrom, StructureLocation, StructureFormat, \
-    Accelerator
+    Accelerator, SrcAssemblyFrom
 
 
 class TestRemoteInference(unittest.TestCase):
 
+    __test_path = os.path.dirname(__file__)
+
     def test_esm_inference_from_structure(self):
         from rcsb_embedding_model.inference.esm_inference import predict
 
@@ -26,3 +29,20 @@ class TestRemoteInference(unittest.TestCase):
         for idx, shape in enumerate(shapes):
             self.assertEqual(tuple(esm_embeddings[idx][0][0].shape), shape)
 
+    def test_assembly_inference_from_structure(self):
+        from rcsb_embedding_model.inference.assembly_inferece import predict
+
+        assembly_embeddings = predict(
+            src_stream=[
+                ("1acb", "https://files.rcsb.org/download/1acb.cif", "1acb"),
+                ("2uzi", "https://files.rcsb.org/download/2uzi.cif", "2uzi")
+            ],
+            res_embedding_location=f"{self.__test_path}/resources/embeddings",
+            src_location=SrcLocation.stream,
+            src_from=SrcAssemblyFrom.structure,
+            structure_location=StructureLocation.remote,
+            structure_format=StructureFormat.mmcif,
+            accelerator=Accelerator.cpu
+        )
+
+        self.assertEqual(len(assembly_embeddings), 2)