rcsb-embedding-model 0.0.9__tar.gz → 0.0.10__tar.gz

{rcsb_embedding_model-0.0.9 → rcsb_embedding_model-0.0.10}/PKG-INFO

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

{rcsb_embedding_model-0.0.9 → rcsb_embedding_model-0.0.10}/README.md

@@ -1,6 +1,6 @@
 # RCSB Embedding Model
 
-**Version** 0.0.9
+**Version** 0.0.10
 
 
 ## Overview

{rcsb_embedding_model-0.0.9 → rcsb_embedding_model-0.0.10}/pyproject.toml

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

{rcsb_embedding_model-0.0.9 → rcsb_embedding_model-0.0.10}/src/rcsb_embedding_model/utils/structure_parser.py

@@ -1,4 +1,4 @@
-from biotite.structure import filter_amino_acids, chain_iter, get_chains, get_residues, AtomArray
+from biotite.structure import filter_amino_acids, filter_polymer, 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
 
@@ -30,7 +30,8 @@ def get_structure_from_src(
 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)]
+        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:
             chain_ids.append(str(get_chains(atom_res)[0]))
     return tuple(chain_ids)

rcsb_embedding_model-0.0.10/tests/test_remote_inference.py

@@ -0,0 +1,28 @@
+import unittest
+
+from rcsb_embedding_model.types.api_types import SrcLocation, SrcProteinFrom, StructureLocation, StructureFormat, \
+    Accelerator
+
+
+class TestRemoteInference(unittest.TestCase):
+
+    def test_esm_inference_from_structure(self):
+        from rcsb_embedding_model.inference.esm_inference import predict
+
+        esm_embeddings = predict(
+            src_stream=[
+                ("3jce", "https://files.rcsb.org/download/3jce.cif", "3jce"),
+                ("9qxy", "https://files.rcsb.org/download/9qxy.cif", "9qxy")
+            ],
+            src_location=SrcLocation.stream,
+            src_from=SrcProteinFrom.structure,
+            structure_location=StructureLocation.remote,
+            structure_format=StructureFormat.mmcif,
+            accelerator=Accelerator.cpu
+        )
+
+        self.assertEqual(len(esm_embeddings), 52)
+        shapes = ((208, 1536), (207, 1536), (152, 1536), (104, 1536), (153, 1536), (131, 1536), (129, 1536), (100, 1536), (119, 1536), (125, 1536), (116, 1536), (102, 1536), (90, 1536), (84, 1536), (82, 1536), (57, 1536), (81, 1536), (87, 1536), (53, 1536), (220, 1536), (58, 1536), (52, 1536), (48, 1536), (66, 1536), (40, 1536), (236, 1536), (272, 1536), (211, 1536), (203, 1536), (179, 1536), (178, 1536), (151, 1536), (143, 1536), (144, 1536), (124, 1536), (145, 1536), (138, 1536), (122, 1536), (118, 1536), (116, 1536), (119, 1536), (105, 1536), (112, 1536), (95, 1536), (104, 1536), (96, 1536), (81, 1536), (79, 1536), (65, 1536), (60, 1536), (588, 1536), (208, 1536))
+        for idx, shape in enumerate(shapes):
+            self.assertEqual(tuple(esm_embeddings[idx][0][0].shape), shape)
+