rcsb-embedding-model 0.0.3__tar.gz → 0.0.5__tar.gz

{rcsb_embedding_model-0.0.3 → rcsb_embedding_model-0.0.5}/PKG-INFO

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

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

{rcsb_embedding_model-0.0.3 → rcsb_embedding_model-0.0.5}/src/rcsb_embedding_model/rcsb_structure_embedding.py

@@ -16,7 +16,7 @@ class RcsbStructureEmbedding:
     MIN_RES = 10
     REPO_ID = "rcsb/rcsb-embedding-model"
     FILE_NAME = "rcsb-embedding-model.pt"
-    VERSION = "2d71cf6"
+    VERSION = "410606e40b1bb7968ce318c41009355c3ac32503"
 
     def __init__(self):
         self.__residue_embedding = None
@@ -66,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)

{rcsb_embedding_model-0.0.3 → rcsb_embedding_model-0.0.5}/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()