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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.28
+Version: 0.0.29
 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.28 → rcsb_embedding_model-0.0.29}/pyproject.toml

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

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

@@ -59,10 +59,10 @@ class EsmProtFromChain(Dataset):
         return len(self.data)
 
     def __getitem__(self, idx):
-        src_name = self.data.loc[idx, EsmProtFromChain.STREAM_NAME_ATTR]
-        src_structure = self.data.loc[idx, EsmProtFromChain.STREAM_ATTR]
-        chain_id = self.data.loc[idx, EsmProtFromChain.CH_ATTR]
-        item_name = self.data.loc[idx, EsmProtFromChain.ITEM_NAME_ATTR]
+        src_name = self.data.iloc[idx][EsmProtFromChain.STREAM_NAME_ATTR]
+        src_structure = self.data.iloc[idx][EsmProtFromChain.STREAM_ATTR]
+        chain_id = self.data.iloc[idx][EsmProtFromChain.CH_ATTR]
+        item_name = self.data.iloc[idx][EsmProtFromChain.ITEM_NAME_ATTR]
         structure = self.__structure_provider.get_structure(
             src_name=src_name,
             src_structure=stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure,
@@ -87,7 +87,7 @@ if __name__ == '__main__':
         src_stream=args.file_list,
         src_location=SrcLocation.file,
         structure_location=StructureLocation.remote,
-        structure_format=StructureFormat.mmcif,
+        structure_format=StructureFormat.bciff,
     )
 
     esm3 = ESM3.from_pretrained(

{rcsb_embedding_model-0.0.28 → rcsb_embedding_model-0.0.29}/src/rcsb_embedding_model/dataset/esm_prot_from_structure.py

@@ -40,7 +40,7 @@ class EsmProtFromStructure(EsmProtFromChain):
 
     def __get_chains(self, src_stream):
         chains = []
-        for idx, row in (pd.DataFrame(
+        data = pd.DataFrame(
             src_stream,
             dtype=str,
             columns=EsmProtFromStructure.COLUMNS
@@ -50,7 +50,9 @@ class EsmProtFromStructure(EsmProtFromChain):
             index_col=None,
             dtype=str,
             names=EsmProtFromStructure.COLUMNS
-        )).iterrows():
+        )
+        data = data.sort_values(by=data.columns[0])
+        for idx, row in data.iterrows():
             src_name = row[EsmProtFromStructure.STREAM_NAME_ATTR]
             src_structure = row[EsmProtFromStructure.STREAM_ATTR]
             item_name = row[EsmProtFromStructure.ITEM_NAME_ATTR]

{rcsb_embedding_model-0.0.28 → rcsb_embedding_model-0.0.29}/src/rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py

@@ -33,7 +33,6 @@ class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile
         self.structure_format = structure_format
         self.min_res_n = min_res_n
         self.max_res_n = max_res_n
-        self.__structure_provider = structure_provider
         super().__init__(
             src_stream=self.__get_assemblies(src_stream),
             res_embedding_location=res_embedding_location,
@@ -47,17 +46,19 @@ class ResidueAssemblyDatasetFromStructure(ResidueAssemblyEmbeddingFromTensorFile
 
     def __get_assemblies(self, src_stream):
         assemblies = []
-        for idx, row in (pd.DataFrame(
-                src_stream,
-                dtype=str,
-                columns=ResidueAssemblyDatasetFromStructure.COLUMNS
+        data = pd.DataFrame(
+            src_stream,
+            dtype=str,
+            columns=ResidueAssemblyDatasetFromStructure.COLUMNS
         ) if self.src_location == SrcLocation.stream else pd.read_csv(
             src_stream,
             header=None,
             index_col=None,
             dtype=str,
             names=ResidueAssemblyDatasetFromStructure.COLUMNS
-        )).iterrows():
+        )
+        data = data.sort_values(by=data.columns[0])
+        for idx, row in data.iterrows():
             src_name = row[ResidueAssemblyDatasetFromStructure.STREAM_NAME_ATTR]
             src_structure = row[ResidueAssemblyDatasetFromStructure.STREAM_ATTR]
             structure = stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure

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

@@ -58,11 +58,10 @@ class ResidueAssemblyEmbeddingFromTensorFile(Dataset):
         return len(self.data)
 
     def __getitem__(self, idx):
-        src_name = self.data.loc[idx, ResidueAssemblyEmbeddingFromTensorFile.STREAM_NAME_ATTR]
-        src_structure = self.data.loc[idx, ResidueAssemblyEmbeddingFromTensorFile.STREAM_ATTR]
-        assembly_id = self.data.loc[idx, ResidueAssemblyEmbeddingFromTensorFile.ASSEMBLY_ATTR]
-        item_name = self.data.loc[idx, ResidueAssemblyEmbeddingFromTensorFile.ITEM_NAME_ATTR]
-
+        src_name = self.data.iloc[idx][ResidueAssemblyEmbeddingFromTensorFile.STREAM_NAME_ATTR]
+        src_structure = self.data.iloc[idx][ResidueAssemblyEmbeddingFromTensorFile.STREAM_ATTR]
+        assembly_id = self.data.iloc[idx][ResidueAssemblyEmbeddingFromTensorFile.ASSEMBLY_ATTR]
+        item_name = self.data.iloc[idx][ResidueAssemblyEmbeddingFromTensorFile.ITEM_NAME_ATTR]
         structure = self.__structure_provider.get_structure(
             src_name=src_name,
             src_structure=stringio_from_url(src_structure) if self.structure_location == StructureLocation.remote else src_structure,

{rcsb_embedding_model-0.0.28 → rcsb_embedding_model-0.0.29}/src/rcsb_embedding_model/dataset/residue_embedding_from_structure.py

@@ -42,17 +42,19 @@ class ResidueEmbeddingFromStructure(ResidueEmbeddingFromTensorFile):
 
     def __get_chains(self, src_stream):
         chains = []
-        for idx, row in (pd.DataFrame(
-                src_stream,
-                dtype=str,
-                columns=ResidueEmbeddingFromStructure.COLUMNS
+        data = pd.DataFrame(
+            src_stream,
+            dtype=str,
+            columns=ResidueEmbeddingFromStructure.COLUMNS
         ) if self.src_location == SrcLocation.stream else pd.read_csv(
             src_stream,
             header=None,
             index_col=None,
             dtype=str,
             names=ResidueEmbeddingFromStructure.COLUMNS
-        )).iterrows():
+        )
+        data = data.sort_values(by=data.columns[0])
+        for idx, row in data.iterrows():
             src_name = row[ResidueEmbeddingFromStructure.STREAM_NAME_ATTR]
             src_structure = row[ResidueEmbeddingFromStructure.STREAM_ATTR]
             item_name = row[ResidueEmbeddingFromStructure.ITEM_NAME_ATTR]

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

@@ -39,6 +39,6 @@ class ResidueEmbeddingFromTensorFile(Dataset):
         return len(self.data)
 
     def __getitem__(self, idx):
-        embedding_src = self.data.loc[idx, ResidueEmbeddingFromTensorFile.FILE_ATTR]
-        item_name = self.data.loc[idx, ResidueEmbeddingFromTensorFile.ITEM_NAME_ATTR]
+        embedding_src = self.data.iloc[idx][ResidueEmbeddingFromTensorFile.FILE_ATTR]
+        item_name = self.data.iloc[idx][ResidueEmbeddingFromTensorFile.ITEM_NAME_ATTR]
         return torch.load(embedding_src, map_location=torch.device('cpu')), item_name

rcsb_embedding_model-0.0.29/tests/resources/src_stream/instance.csv

@@ -0,0 +1,2 @@
+1acb,https://models.rcsb.org/1acb.bcif.gz,A,1acb.A
+2uzi,https://models.rcsb.org/2uzi.bcif.gz,A,2uzi.A

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

@@ -50,6 +50,24 @@ class TestRemoteInference(unittest.TestCase):
             self.assertEqual(tuple(esm_embeddings[idx][0][0].shape), shape)
 
 
+    def test_esm_inference_from_csv_bcif_gz(self):
+        from rcsb_embedding_model.inference.esm_inference import predict
+
+        esm_embeddings = predict(
+            src_stream=f"{self.__test_path}/resources/src_stream/instance.csv",
+            src_location=SrcLocation.file,
+            src_from=SrcProteinFrom.chain,
+            structure_location=StructureLocation.remote,
+            structure_format=StructureFormat.bciff,
+            accelerator=Accelerator.cpu
+        )
+
+        self.assertEqual(len(esm_embeddings), 2)
+        shapes = ((243, 1536), (116, 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