rcsb-embedding-model 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

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

@@ -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.4.dist-info/RECORD

@@ -0,0 +1,8 @@
+rcsb_embedding_model/__init__.py,sha256=r3gLdeBIXkQEQA_K6QcRPO-TtYuAQSutk6pXRUE_nas,120
+rcsb_embedding_model/rcsb_structure_embedding.py,sha256=W_Os_xa_ZkYXR_KPjvwfzW8raMIzjsMCdhfKktbVY3s,5983
+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-0.0.4.dist-info/METADATA,sha256=4wBBZfkyq5FY6PgMeAzBB4Aka2k5eheJpL6uXVn9DXA,5381
+rcsb_embedding_model-0.0.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+rcsb_embedding_model-0.0.4.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
+rcsb_embedding_model-0.0.4.dist-info/RECORD,,

rcsb_embedding_model-0.0.2.dist-info/RECORD

@@ -1,8 +0,0 @@
-rcsb_embedding_model/__init__.py,sha256=r3gLdeBIXkQEQA_K6QcRPO-TtYuAQSutk6pXRUE_nas,120
-rcsb_embedding_model/rcsb_structure_embedding.py,sha256=quxl9SnLPIYJE7yn7kXgJvCKR-W9kXIJ5fBFTTLV1YQ,5195
-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-0.0.2.dist-info/METADATA,sha256=O5JNOufpbfTUJznAWaWMwwicerku0ih7TjWfzAc3OJo,5370
-rcsb_embedding_model-0.0.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-rcsb_embedding_model-0.0.2.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
-rcsb_embedding_model-0.0.2.dist-info/RECORD,,