rcsb-embedding-model 0.0.18__py3-none-any.whl → 0.0.20__py3-none-any.whl

rcsb_embedding_model/modules/esm_module.py

@@ -1,5 +1,4 @@
 from esm.sdk.api import SamplingConfig
-from esm.sdk import batch_executor
 from lightning import LightningModule
 
 from rcsb_embedding_model.utils.model import get_residue_model
@@ -14,14 +13,9 @@ class EsmModule(LightningModule):
         self.esm3 = get_residue_model(self.device)
 
     def predict_step(self, prot_batch, batch_idx):
-        prot_embeddings = []
-        def __batch_embedding(esm_prot):
-            return self.esm3.forward_and_sample(
-                self.esm3.encode(esm_prot), SamplingConfig(return_per_residue_embeddings=True)
-            ).per_residue_embedding
-        with batch_executor() as executor:
-            prot_embeddings = executor.execute_batch(
-                user_func=__batch_embedding,
-                esm_prot=[esm_prot for esm_prot, name in prot_batch]
-            )
-        return tuple(prot_embeddings), tuple([name for esm_prot, name in prot_batch])
+        return tuple([self.__compute_embeddings(esm_prot) for esm_prot, name in prot_batch]), tuple([name for esm_prot, name in prot_batch])
+
+    def __compute_embeddings(self, esm_prot):
+        return self.esm3.forward_and_sample(
+            self.esm3.encode(esm_prot), SamplingConfig(return_per_residue_embeddings=True)
+        ).per_residue_embedding

rcsb_embedding_model/utils/data.py

@@ -76,10 +76,10 @@ def concatenate_tensors(file_list, max_residues, dim=0):
         tensor_cat = torch.cat(tensors, dim=dim)
         return tensor_cat
     else:
-        raise ValueError("No valid tensors were loaded to concatenate.")
+        raise ValueError(f"No valid tensors were loaded to concatenate. {', '.join(file_list)}")
 
 def adapt_csv_to_embedding_chain_stream(src_file, res_embedding_location):
     def __parse_row(row):
         r = row.split(",")
         return os.path.join(res_embedding_location, f"{r[0]}.{r[2]}.pt"), f"{r[0]}.{r[2]}"
-    return tuple([__parse_row(r) for r in open(src_file)])
+    return tuple([__parse_row(r.strip()) for r in open(src_file)])

{rcsb_embedding_model-0.0.18.dist-info → rcsb_embedding_model-0.0.20.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.18
+Version: 0.0.20
 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
@@ -17,7 +17,7 @@ Description-Content-Type: text/markdown
 
 # RCSB Embedding Model
 
-**Version** 0.0.18
+**Version** 0.0.20
 
 
 ## Overview

{rcsb_embedding_model-0.0.18.dist-info → rcsb_embedding_model-0.0.20.dist-info}/RECORD

@@ -15,16 +15,16 @@ rcsb_embedding_model/inference/structure_inference.py,sha256=lqbDBPSea8IoNyQXl83
 rcsb_embedding_model/model/layers.py,sha256=lhKaWC4gTS_T5lHOP0mgnnP8nKTPEOm4MrjhESA4hE8,743
 rcsb_embedding_model/model/residue_embedding_aggregator.py,sha256=k3UW63Ax8DtjCMdD3O5xNxtyAu28l2n3-Ab6nS0atm0,1967
 rcsb_embedding_model/modules/chain_module.py,sha256=sDSPXJmWuU2C3lt1NorlbUVWZvRSLzumPdFQk01h3VI,403
-rcsb_embedding_model/modules/esm_module.py,sha256=4IQgrNQlGThxl0PhobVzyp7N3FcyAbvek_KxJozGImQ,945
+rcsb_embedding_model/modules/esm_module.py,sha256=MzE-9_25262m80yW7Pb-Pkf8L-R_xb3FPcz9gOlNgAI,705
 rcsb_embedding_model/modules/structure_module.py,sha256=dEtDNdWo1j2sSDa0JiOHQfEfQzIWqSLEKpvOX0GrXZ4,1048
 rcsb_embedding_model/types/api_types.py,sha256=SCwALwvEb0KRKaoWKbuN7JyfOH-1whsI0Z4ki41dht8,1235
-rcsb_embedding_model/utils/data.py,sha256=FVb6tzoX4SrJf3Fr6UFbxZJQsUr9xp5RbkK6nqXhcuQ,3222
+rcsb_embedding_model/utils/data.py,sha256=ATgHC6d7xgc2eYtCsEsgMxhAFvYp4R9C0TlNStc1jG0,3254
 rcsb_embedding_model/utils/model.py,sha256=rpZa-gfm3cEtbBd7UXMHrZv3x6f0AC8TJT3gtrSxr5I,852
 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.18.dist-info/METADATA,sha256=PzSnwGeAeUbYxhpRBgEiZZdj2bGdLrT8QAy0uB_BxNQ,5310
-rcsb_embedding_model-0.0.18.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-rcsb_embedding_model-0.0.18.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
-rcsb_embedding_model-0.0.18.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
-rcsb_embedding_model-0.0.18.dist-info/RECORD,,
+rcsb_embedding_model-0.0.20.dist-info/METADATA,sha256=LBfyPJYo7iUu2XmBC0mzoX-Qo1_FHu8BPWw8au1VkQI,5310
+rcsb_embedding_model-0.0.20.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+rcsb_embedding_model-0.0.20.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
+rcsb_embedding_model-0.0.20.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
+rcsb_embedding_model-0.0.20.dist-info/RECORD,,