rcsb-embedding-model 0.0.27__py3-none-any.whl → 0.0.29__py3-none-any.whl

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/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/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/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/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/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/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.27.dist-info → rcsb_embedding_model-0.0.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.27
+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.27.dist-info → rcsb_embedding_model-0.0.29.dist-info}/RECORD

@@ -2,12 +2,12 @@ rcsb_embedding_model/__init__.py,sha256=r3gLdeBIXkQEQA_K6QcRPO-TtYuAQSutk6pXRUE_
 rcsb_embedding_model/rcsb_structure_embedding.py,sha256=dKp9hXQO0JAnO4SEfjJ_mG_jHu3UxAPguv6jkOjp-BI,4487
 rcsb_embedding_model/cli/args_utils.py,sha256=7nP2q8pL5dWK_U7opxtWmoFcYVwasky6elHk-dASFaI,165
 rcsb_embedding_model/cli/inference.py,sha256=tfMvHAhkUIzJ2RbTtQjq7eWmOUrSyVfH5bjTkCCSIS8,19500
-rcsb_embedding_model/dataset/esm_prot_from_chain.py,sha256=9YVDONAUaIL0U0J_A_0gokN5sN6FgthJHRg_YXTIFLk,3956
-rcsb_embedding_model/dataset/esm_prot_from_structure.py,sha256=TeITPdi1uc3qLQ-Pgn807oH6eM0LYv-67RE50ZT4dLI,2551
-rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py,sha256=worRiNqOJRjyr693TaillsS65bdTdGOoHfwyT9yE1O4,2866
-rcsb_embedding_model/dataset/residue_assembly_embedding_from_tensor_file.py,sha256=U-XbLVC6d17viKWlSWMs80ztj9VsKjR167B8OvvOfWY,3516
-rcsb_embedding_model/dataset/residue_embedding_from_structure.py,sha256=dxfUNcVmdl8LrtQf1UJQ4E79e7R9LRsL0fjsq2GJQRk,2796
-rcsb_embedding_model/dataset/residue_embedding_from_tensor_file.py,sha256=qxSnbajARYvyl_KUzoDjqtuCI8YH_7ATpWqDSY0D6F0,1341
+rcsb_embedding_model/dataset/esm_prot_from_chain.py,sha256=l8pRgRpz-8gEzR6QsNcmdnX_KkgvgF1vdqNAsP8Lrc8,3960
+rcsb_embedding_model/dataset/esm_prot_from_structure.py,sha256=3HzXCCc-UqmZNbJaeXHyUsSIZZxMc2erbxAPGIxSmfE,2621
+rcsb_embedding_model/dataset/resdiue_assembly_embedding_from_structure.py,sha256=69h1VkrIXesHZi1cG3BOMMytSDeRzcBBP0_Z3Xz3dM8,2869
+rcsb_embedding_model/dataset/residue_assembly_embedding_from_tensor_file.py,sha256=R-_WY8YJ6qNBjWvY9IzCRThAA42GJ71hzwZBMnkMhOg,3519
+rcsb_embedding_model/dataset/residue_embedding_from_structure.py,sha256=1jmeEcCK41cAi2ZnqQkd667NWCAIGS3k6jGDF-WxtTk,2854
+rcsb_embedding_model/dataset/residue_embedding_from_tensor_file.py,sha256=4OPaw55yGKHjY2iPpCnemcfwfmTZ4j5VrGQ2oIMQw6A,1343
 rcsb_embedding_model/inference/assembly_inferece.py,sha256=8fPJjEXy1WsM5XB5U7KfdO5-Du6nEsawsaAjmWoXA9I,2329
 rcsb_embedding_model/inference/chain_inference.py,sha256=6f5wVzjtRtHU3BPMTe5k3nH_Nl440Am8BL8h1vmK1jI,2925
 rcsb_embedding_model/inference/esm_inference.py,sha256=rn6H43D8BYzMZbMu7UPsLYg2dgERmmpci5weNItrG5Q,2546
@@ -18,13 +18,13 @@ rcsb_embedding_model/modules/chain_module.py,sha256=KsZw2uagO4rpAKWv6ivqEMxIEzgt
 rcsb_embedding_model/modules/esm_module.py,sha256=otJRbCb319nCCob_4E1W_UClhkex9eDqcCyzWQO-vIs,740
 rcsb_embedding_model/modules/structure_module.py,sha256=4js02XzKvhc_G26ELsGhJ9SCi_wlvtVolObxfWt3BhE,1077
 rcsb_embedding_model/types/api_types.py,sha256=SCwALwvEb0KRKaoWKbuN7JyfOH-1whsI0Z4ki41dht8,1235
-rcsb_embedding_model/utils/data.py,sha256=ATgHC6d7xgc2eYtCsEsgMxhAFvYp4R9C0TlNStc1jG0,3254
+rcsb_embedding_model/utils/data.py,sha256=BOjYdIRHrFqk8qFuKGrgCtVyfDupzgOVmH_0C-ecMvg,3813
 rcsb_embedding_model/utils/model.py,sha256=xr3p02ohOgJ5UInwdIupN68Oq4yvNFhxobZRacS1adg,953
 rcsb_embedding_model/utils/structure_parser.py,sha256=IWMQ8brlEMe6_ND-DBESOli8vlqHxladTssjbM9RSKw,2751
 rcsb_embedding_model/utils/structure_provider.py,sha256=eWtxjkPpmRfmil_DKR1J6miaXR3lQ28DF5O0qrqSgGA,786
 rcsb_embedding_model/writer/batch_writer.py,sha256=rTFNasB0Xp4-XCNTXKeEWZxSrb7lvZytoRldJUWn9Jg,3312
-rcsb_embedding_model-0.0.27.dist-info/METADATA,sha256=rbZPxX85Fo7dBZK7eUUxm9etfChs09sWgbKWh4r7Dek,5310
-rcsb_embedding_model-0.0.27.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-rcsb_embedding_model-0.0.27.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
-rcsb_embedding_model-0.0.27.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
-rcsb_embedding_model-0.0.27.dist-info/RECORD,,
+rcsb_embedding_model-0.0.29.dist-info/METADATA,sha256=R3z_Non64xek0bqI40DAoqejy2BtSuHqYuAni2R-eFA,5310
+rcsb_embedding_model-0.0.29.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+rcsb_embedding_model-0.0.29.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
+rcsb_embedding_model-0.0.29.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
+rcsb_embedding_model-0.0.29.dist-info/RECORD,,