rcsb-embedding-model 0.0.2__tar.gz → 0.0.4__tar.gz

{rcsb_embedding_model-0.0.2 → rcsb_embedding_model-0.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.2
+Version: 0.0.4
 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
@@ -11,7 +11,7 @@ Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.10
 Requires-Dist: esm>=3.2.0
-Requires-Dist: torch>=2.6.0
+Requires-Dist: torch>=2.2.0
 Description-Content-Type: text/markdown
 
 # RCSB Embedding Model: A Deep Learning Approach for 3D Structure Embeddings
@@ -70,7 +70,7 @@ res_embedding = model.residue_embedding(
 )
 ```
 
-### **Generating Protein Structure**
+### **Generating Protein Structure Embeddings**
 Protein 3D structure embedding can be calculated as:
 
 ```python

{rcsb_embedding_model-0.0.2 → rcsb_embedding_model-0.0.4}/README.md

@@ -54,7 +54,7 @@ res_embedding = model.residue_embedding(
 )
 ```
 
-### **Generating Protein Structure**
+### **Generating Protein Structure Embeddings**
 Protein 3D structure embedding can be calculated as:
 
 ```python

{rcsb_embedding_model-0.0.2 → rcsb_embedding_model-0.0.4}/examples/esm_embeddings.py

@@ -4,7 +4,7 @@ import torch
 from biotite.structure import chain_iter, get_residues, filter_amino_acids
 from biotite.structure.io.pdb import PDBFile
 from biotite.structure.io.pdbx import CIFFile, get_structure, BinaryCIFFile
-from esm.models.esm3 import ESM3, ESM3_OPEN_SMALL, ESM3InferenceClient
+from esm.models.esm3 import ESM3, ESM3_OPEN_SMALL
 from esm.sdk.api import ESMProtein, SamplingConfig
 from esm.utils.structure.protein_chain import ProteinChain
 

{rcsb_embedding_model-0.0.2 → rcsb_embedding_model-0.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.2"
+version = "0.0.4"
 authors = [
     { name="Joan Segura", email="joan.segura@rcsb.org" },
 ]
@@ -15,7 +15,7 @@ license = "BSD-3-Clause"
 license-files = ["LICEN[CS]E*"]
 dependencies=[
     "esm >= 3.2.0",
-    "torch >= 2.6.0"
+    "torch >= 2.2.0"
 ]
 [project.urls]
 Homepage = "https://github.com/rcsb/rcsb-embedding-model"

{rcsb_embedding_model-0.0.2 → rcsb_embedding_model-0.0.4}/src/rcsb_embedding_model/rcsb_structure_embedding.py

@@ -29,10 +29,7 @@ class RcsbStructureEmbedding:
     def load_residue_embedding(self, device=None):
         if not device:
             device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.__residue_embedding = ESM3.from_pretrained(
-            ESM3_OPEN_SMALL,
-            device
-        )
+        self.__residue_embedding = _load_res_model(device)
 
     def load_aggregator_embedding(self, device=None):
         if not device:
@@ -69,6 +66,25 @@ class RcsbStructureEmbedding:
             dim=0
         )
 
+    def sequence_embedding(self, sequence):
+        self.__check_residue_embedding()
+
+        if sequence.startswith(">"):
+            sequence = "".join(line.strip() for line in sequence.splitlines() if not line.startswith(">"))
+
+        if len(sequence) < RcsbStructureEmbedding.MIN_RES:
+            raise ValueError(f"Sequence too short for embedding (min {RcsbStructureEmbedding.MIN_RES} residues)")
+
+        protein = ESMProtein(sequence=sequence)
+        protein_tensor = self.__residue_embedding.encode(protein)
+
+        result = self.__residue_embedding.forward_and_sample(
+            protein_tensor,
+            SamplingConfig(return_per_residue_embeddings=True)
+        )
+
+        return result.per_residue_embedding
+
     def aggregator_embedding(self, residue_embedding):
         self.__check_aggregator_embedding()
         return self.__aggregator_embedding(residue_embedding)
@@ -146,3 +162,10 @@ def _load_model(model_path, device=None):
     aggregator_model.to(device)
     aggregator_model.eval()
     return aggregator_model
+
+
+def _load_res_model(device=None):
+    return ESM3.from_pretrained(
+        ESM3_OPEN_SMALL,
+        device
+    )

{rcsb_embedding_model-0.0.2 → rcsb_embedding_model-0.0.4}/tests/test_model.py

@@ -18,6 +18,16 @@ class TestEmbeddingModel(unittest.TestCase):
         )
         self.assertEqual(list(res_embedding.shape), [243, 1536])
 
+    def test_sequence_embedding(self):
+
+        model = RcsbStructureEmbedding()
+        res_embedding = model.sequence_embedding(
+            sequence="CGVPAIQPVLSGLSRIVNGEEAVPGSWPWQVSLQDKTGFHFCGGSLINENWVVTAAHCGVTTSDVVVAGEFDQGSSSEKIQKLKIAKVFKNSK"
+                     "YNSLTINNDITLLKLSTAASFSQTVSAVCLPSASDDFAAGTTCVTTGWGLTRYTNANTPDRLQQASLPLLSNTNCKKYWGTKIKDAMICAGAS"
+                     "GVSSCMGDSGGPLVCKKNGAWTLVGIVSWGSSTCSTSTPGVYARVTALVNWVQQTLAAN"
+        )
+        self.assertEqual(list(res_embedding.shape), [247, 1536])
+
     def test_aggregator_embedding(self):
 
         model = RcsbStructureEmbedding()