rcsb-embedding-model 0.0.36__py3-none-any.whl → 0.0.38__py3-none-any.whl

rcsb_embedding_model/utils/data.py

@@ -95,4 +95,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.strip()) for r in open(src_file)])
+    return tuple([__parse_row(r.strip()) for r in open(src_file) if len(r.split(",")) > 2])

rcsb_embedding_model/utils/esm/loaders.py

@@ -0,0 +1,65 @@
+from pathlib import Path
+import torch
+
+from esm.models.esm3 import ESM3
+from esm.models.vqvae import StructureTokenEncoder
+from esm.tokenization import TokenizerCollection, EsmSequenceTokenizer, StructureTokenizer, SecondaryStructureTokenizer, \
+    SASADiscretizingTokenizer, InterProQuantizedTokenizer, ResidueAnnotationsTokenizer
+
+from huggingface_hub import  snapshot_download
+
+def data_root():
+    path = Path(snapshot_download(repo_id="rcsb/rcsb-esm"))
+    return path
+
+
+def structure_encoder(device: torch.device | str = "cpu"):
+    with torch.device(device):
+        model = StructureTokenEncoder(
+            d_model=1024, n_heads=1, v_heads=128, n_layers=2, d_out=128, n_codes=4096
+        ).eval()
+    state_dict = torch.load(
+        data_root() / "data/weights/esm3_structure_encoder_v0.pth", map_location=device
+    )
+    model.load_state_dict(state_dict)
+    return model
+
+
+
+def get_model_tokenizers():
+
+    class CustomAnnotationsTokenizer(ResidueAnnotationsTokenizer):
+        def __init__(self, csv_path: str | None = None, max_annotations: int = 16):
+            from esm.utils.constants import esm3 as C
+            super().__init__("none", max_annotations)
+            if csv_path is None:
+                csv_path = str(data_root() / C.RESID_CSV)
+            self.csv_path = csv_path
+
+    return TokenizerCollection(
+        sequence=EsmSequenceTokenizer(),
+        structure=StructureTokenizer(),
+        secondary_structure=SecondaryStructureTokenizer(kind="ss8"),
+        sasa=SASADiscretizingTokenizer(),
+        function=InterProQuantizedTokenizer(),
+        residue_annotations=CustomAnnotationsTokenizer(),
+    )
+
+
+def esm_open(device: torch.device | str = "cpu"):
+    with torch.device(device):
+        model = ESM3(
+            d_model=1536,
+            n_heads=24,
+            v_heads=256,
+            n_layers=48,
+            structure_encoder_fn=structure_encoder,
+            structure_decoder_fn=lambda x: x,
+            function_decoder_fn=lambda x: x,
+            tokenizers=get_model_tokenizers(),
+        ).eval()
+    state_dict = torch.load(
+        data_root() / "data/weights/esm3_sm_open_v1.pth", map_location=device
+    )
+    model.load_state_dict(state_dict)
+    return model

rcsb_embedding_model/utils/model.py

@@ -1,9 +1,8 @@
 import torch
-from esm.models.esm3 import ESM3
-from esm.utils.constants.models import ESM3_OPEN_SMALL
-from huggingface_hub import hf_hub_download
 
+from huggingface_hub import hf_hub_download
 from rcsb_embedding_model.model.residue_embedding_aggregator import ResidueEmbeddingAggregator
+from rcsb_embedding_model.utils.esm.loaders import esm_open
 
 REPO_ID = "rcsb/rcsb-embedding-model"
 FILE_NAME = "rcsb-embedding-model.pt"
@@ -25,7 +24,5 @@ def get_aggregator_model(device=None):
 
 
 def get_residue_model(device=None):
-    return ESM3.from_pretrained(
-        ESM3_OPEN_SMALL,
-        device
-    )
+    return esm_open(device)
+

{rcsb_embedding_model-0.0.36.dist-info → rcsb_embedding_model-0.0.38.dist-info}/METADATA

@@ -1,16 +1,23 @@
 Metadata-Version: 2.4
 Name: rcsb-embedding-model
-Version: 0.0.36
+Version: 0.0.38
 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
 Author-email: Joan Segura <joan.segura@rcsb.org>
-License-Expression: BSD-3-Clause
+License: # Cambrian Non-Commercial License Agreement
+        
+        This project is licensed under the EvolutionaryScale Cambrian Non-Commercial License Agreement.
+        See: https://www.evolutionaryscale.ai/policies/cambrian-non-commercial-license-agreement
 License-File: LICENSE.md
 Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python :: 3
-Requires-Python: >=3.10
+Requires-Python: >=3.11
+Requires-Dist: biotite>=1.5.0
 Requires-Dist: esm>=3.2.0
+Requires-Dist: hf-xet>=1.1.10
+Requires-Dist: httpx>=0.28.1
+Requires-Dist: huggingface-hub>=0.30.2
 Requires-Dist: importlib-metadata>=8.7.0
 Requires-Dist: lightning>=2.5.0
 Requires-Dist: typer>=0.15.0
@@ -18,7 +25,7 @@ Description-Content-Type: text/markdown
 
 # RCSB Embedding Model
 
-**Version** 0.0.26
+**Version** 0.0.38
 
 
 ## Overview
@@ -125,4 +132,5 @@ Segura, J., Bittrich, S., et al. (2024). *Multi-scale structural similarity embe
 
 ## License
 
-This project is licensed under the BSD 3-Clause License. See [LICENSE.md](LICENSE.md) for details.
+This project uses the EvolutionaryScale ESM-3 model and is distributed under the
+[Cambrian Non-Commercial License Agreement](https://www.evolutionaryscale.ai/policies/cambrian-non-commercial-license-agreement).

{rcsb_embedding_model-0.0.36.dist-info → rcsb_embedding_model-0.0.38.dist-info}/RECORD

@@ -20,13 +20,14 @@ rcsb_embedding_model/modules/chain_module.py,sha256=KsZw2uagO4rpAKWv6ivqEMxIEzgt
 rcsb_embedding_model/modules/esm_module.py,sha256=otJRbCb319nCCob_4E1W_UClhkex9eDqcCyzWQO-vIs,740
 rcsb_embedding_model/modules/structure_module.py,sha256=4js02XzKvhc_G26ELsGhJ9SCi_wlvtVolObxfWt3BhE,1077
 rcsb_embedding_model/types/api_types.py,sha256=SCwALwvEb0KRKaoWKbuN7JyfOH-1whsI0Z4ki41dht8,1235
-rcsb_embedding_model/utils/data.py,sha256=ThrcYycIizsV_Ycn6PPxF12JRr1m2K-v8TsIaVqx10A,3816
-rcsb_embedding_model/utils/model.py,sha256=xr3p02ohOgJ5UInwdIupN68Oq4yvNFhxobZRacS1adg,953
+rcsb_embedding_model/utils/data.py,sha256=p7sbskLPBFtpZ-XM18wFY5Kei02Xso4wTWYTqHxJvVw,3841
+rcsb_embedding_model/utils/model.py,sha256=Xi6bSUsB2-IsQS9610gXnbAEvYlK2V7eJC-cDE-JBTA,875
 rcsb_embedding_model/utils/structure_parser.py,sha256=fSIbq_a_aEigCWY_1dUcW9d9Law0ZDOcZAxJlZL0Rt8,3377
 rcsb_embedding_model/utils/structure_provider.py,sha256=eWtxjkPpmRfmil_DKR1J6miaXR3lQ28DF5O0qrqSgGA,786
+rcsb_embedding_model/utils/esm/loaders.py,sha256=V7CADr7RReoztYmBQb2tjA8RBQIwFEjxBcocKAB_ea4,2221
 rcsb_embedding_model/writer/batch_writer.py,sha256=rTFNasB0Xp4-XCNTXKeEWZxSrb7lvZytoRldJUWn9Jg,3312
-rcsb_embedding_model-0.0.36.dist-info/METADATA,sha256=spFNxlrrwMORe5Su0559-997by2cgkuk9-yEQlhew60,5351
-rcsb_embedding_model-0.0.36.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-rcsb_embedding_model-0.0.36.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
-rcsb_embedding_model-0.0.36.dist-info/licenses/LICENSE.md,sha256=oUaHiKgfBkChth_Sm67WemEvatO1U0Go8LHjaskXY0w,1522
-rcsb_embedding_model-0.0.36.dist-info/RECORD,,
+rcsb_embedding_model-0.0.38.dist-info/METADATA,sha256=uBAVCA8bTzGUza-i1Dr6_WpDPVCSZam4pvqg_ZT37ms,5820
+rcsb_embedding_model-0.0.38.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+rcsb_embedding_model-0.0.38.dist-info/entry_points.txt,sha256=MK11jTIEmaV-x4CkPX5IymDaVs7Ky_f2xxU8BJVZ_9Q,69
+rcsb_embedding_model-0.0.38.dist-info/licenses/LICENSE.md,sha256=XyzxQe9PLJQlOmOOrqwmBaAfo0PAenOQ5NsgnApuVH4,230
+rcsb_embedding_model-0.0.38.dist-info/RECORD,,

rcsb_embedding_model-0.0.38.dist-info/licenses/LICENSE.md

@@ -0,0 +1,4 @@
+# Cambrian Non-Commercial License Agreement
+
+This project is licensed under the EvolutionaryScale Cambrian Non-Commercial License Agreement.
+See: https://www.evolutionaryscale.ai/policies/cambrian-non-commercial-license-agreement

rcsb_embedding_model-0.0.36.dist-info/licenses/LICENSE.md

@@ -1,28 +0,0 @@
-BSD 3-Clause License
-
-Copyright (c) 2024, RCSB Protein Data Bank, UC San Diego
-
-Redistribution and use in source and binary forms, with or without
-modification, are permitted provided that the following conditions are met:
-
-1. Redistributions of source code must retain the above copyright notice, this
-   list of conditions and the following disclaimer.
-
-2. Redistributions in binary form must reproduce the above copyright notice,
-   this list of conditions and the following disclaimer in the documentation
-   and/or other materials provided with the distribution.
-
-3. Neither the name of the copyright holder nor the names of its
-   contributors may be used to endorse or promote products derived from
-   this software without specific prior written permission.
-
-THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
-AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
-IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
-DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
-FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
-DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
-SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
-CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
-OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
-OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.