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

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

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

@@ -1,6 +1,6 @@
 [project]
 name = "rcsb-embedding-model"
-version = "0.0.2"
+version = "0.0.3"
 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.3}/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:
@@ -146,3 +143,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
+    )