rcsb-embedding-model 0.0.26__tar.gz → 0.0.28__tar.gz

{rcsb_embedding_model-0.0.26 → rcsb_embedding_model-0.0.28}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.26
+Version: 0.0.28
 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

{rcsb_embedding_model-0.0.26 → rcsb_embedding_model-0.0.28}/pyproject.toml

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

{rcsb_embedding_model-0.0.26 → rcsb_embedding_model-0.0.28}/src/rcsb_embedding_model/dataset/esm_prot_from_chain.py

@@ -53,6 +53,7 @@ class EsmProtFromChain(Dataset):
             dtype=str,
             names=EsmProtFromChain.COLUMNS
         )
+        self.data = self.data.sort_values(by=self.data.columns[0])
 
     def __len__(self):
         return len(self.data)

{rcsb_embedding_model-0.0.26 → rcsb_embedding_model-0.0.28}/src/rcsb_embedding_model/dataset/residue_assembly_embedding_from_tensor_file.py

@@ -52,6 +52,7 @@ class ResidueAssemblyEmbeddingFromTensorFile(Dataset):
             dtype=str,
             names=ResidueAssemblyEmbeddingFromTensorFile.COLUMNS
         )
+        self.data = self.data.sort_values(by=self.data.columns[0])
 
     def __len__(self):
         return len(self.data)

{rcsb_embedding_model-0.0.26 → rcsb_embedding_model-0.0.28}/src/rcsb_embedding_model/dataset/residue_embedding_from_tensor_file.py

@@ -33,6 +33,7 @@ class ResidueEmbeddingFromTensorFile(Dataset):
             index_col=None,
             names=ResidueEmbeddingFromTensorFile.COLUMNS
         )
+        self.data = self.data.sort_values(by=self.data.columns[0])
 
     def __len__(self):
         return len(self.data)

{rcsb_embedding_model-0.0.26 → rcsb_embedding_model-0.0.28}/src/rcsb_embedding_model/utils/data.py

@@ -1,7 +1,8 @@
 import os
-from io import StringIO
-
 import requests
+import gzip
+from io import StringIO, BytesIO
+
 import torch
 
 
@@ -40,10 +41,24 @@ def stringio_from_url(url):
     try:
         response = requests.get(url)
         response.raise_for_status()
-        return StringIO(response.text)
+        data = response.content
+        if url.endswith('.bcif.gz'):
+            with gzip.GzipFile(fileobj=BytesIO(data), mode='rb') as gz:
+                decompressed_data = gz.read()
+                return BytesIO(decompressed_data)
+        if url.endswith('.gz'):
+            compressed = BytesIO(data)
+            with gzip.open(compressed, 'rt') as f:
+                return StringIO(f.read())
+        else:
+            return StringIO(response.text)
     except requests.exceptions.RequestException as e:
         print(f"Error fetching URL: {e}")
         return None
+    except (OSError, gzip.BadGzipFile) as e:
+        print(f"Error decompressing gzip file: {e}")
+        return None
+
 
 
 def concatenate_tensors(file_list, max_residues, dim=0):

{rcsb_embedding_model-0.0.26 → rcsb_embedding_model-0.0.28}/tests/test_remote_inference.py

@@ -29,6 +29,48 @@ class TestRemoteInference(unittest.TestCase):
         for idx, shape in enumerate(shapes):
             self.assertEqual(tuple(esm_embeddings[idx][0][0].shape), shape)
 
+    def test_esm_inference_from_bcif_gz(self):
+        from rcsb_embedding_model.inference.esm_inference import predict
+
+        esm_embeddings = predict(
+            src_stream=[
+                ("1acb", "https://models.rcsb.org/1acb.bcif.gz", "1acb"),
+                ("2uzi", "https://models.rcsb.org/2uzi.bcif.gz", "2uzi")
+            ],
+            src_location=SrcLocation.stream,
+            src_from=SrcProteinFrom.structure,
+            structure_location=StructureLocation.remote,
+            structure_format=StructureFormat.bciff,
+            accelerator=Accelerator.cpu
+        )
+
+        self.assertEqual(len(esm_embeddings), 5)
+        shapes = ((243, 1536), (65, 1536), (116, 1536), (106, 1536), (168, 1536))
+        for idx, shape in enumerate(shapes):
+            self.assertEqual(tuple(esm_embeddings[idx][0][0].shape), shape)
+
+
+    def test_esm_inference_from_cif_gz(self):
+        from rcsb_embedding_model.inference.esm_inference import predict
+
+        esm_embeddings = predict(
+            src_stream=[
+                ("1acb", "https://files.rcsb.org/download/1acb.cif.gz", "1acb"),
+                ("2uzi", "https://files.rcsb.org/download/2uzi.cif.gz", "2uzi")
+            ],
+            src_location=SrcLocation.stream,
+            src_from=SrcProteinFrom.structure,
+            structure_location=StructureLocation.remote,
+            structure_format=StructureFormat.mmcif,
+            accelerator=Accelerator.cpu
+        )
+
+        self.assertEqual(len(esm_embeddings), 5)
+        shapes = ((243, 1536), (65, 1536), (116, 1536), (106, 1536), (168, 1536))
+        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