rcsb-embedding-model 0.0.18__tar.gz → 0.0.19__tar.gz

{rcsb_embedding_model-0.0.18 → rcsb_embedding_model-0.0.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.18
+Version: 0.0.19
 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.19
 
 
 ## Overview

{rcsb_embedding_model-0.0.18 → rcsb_embedding_model-0.0.19}/README.md

@@ -1,6 +1,6 @@
 # RCSB Embedding Model
 
-**Version** 0.0.18
+**Version** 0.0.19
 
 
 ## Overview

{rcsb_embedding_model-0.0.18 → rcsb_embedding_model-0.0.19}/pyproject.toml

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

rcsb_embedding_model-0.0.19/src/rcsb_embedding_model/modules/esm_module.py

@@ -0,0 +1,21 @@
+from esm.sdk.api import SamplingConfig
+from lightning import LightningModule
+
+from rcsb_embedding_model.utils.model import get_residue_model
+
+
+class EsmModule(LightningModule):
+
+    def __init__(
+            self
+    ):
+        super().__init__()
+        self.esm3 = get_residue_model(self.device)
+
+    def predict_step(self, prot_batch, batch_idx):
+        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-0.0.18 → rcsb_embedding_model-0.0.19}/src/rcsb_embedding_model/utils/data.py

@@ -82,4 +82,4 @@ 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/src/rcsb_embedding_model/modules/esm_module.py

@@ -1,27 +0,0 @@
-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
-
-
-class EsmModule(LightningModule):
-
-    def __init__(
-            self
-    ):
-        super().__init__()
-        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])