v.0.5.1 (release not out of main)

Author: arogozhnikov

README.md

@@ -12,7 +12,7 @@ For more information on the model's performance and capabilities, see our [techn
 
 ```shell
 # version on pypi:
-pip install chai_lab==0.5.0
+pip install chai_lab==0.5.1
 
 # newest available version (updates daily to test features that weren't released yet):
 pip install git+https://github.com/chaidiscovery/chai-lab.git
@@ -123,7 +123,7 @@ Devcontainers work on local Linux setup, and on remote machines over an SSH conn
 Since this is an initial release, we expect to make some breaking changes to the API and are not guaranteeing backwards compatibility. We recommend pinning the current version in your requirements, i.e.:
 
 ```
-chai_lab==0.5.0
+chai_lab==0.5.1
 ```
 
 ## Citations

chai_lab/__init__.py

@@ -2,4 +2,4 @@
 # Licensed under the Apache License, Version 2.0.
 # See the LICENSE file for details.
 
-__version__ = "0.5.0"
+__version__ = "0.5.1"

examples/predict_structure.private.py

@@ -0,0 +1,67 @@
+import datetime
+from pathlib import Path
+
+import numpy as np
+import torch
+
+import chai_lab
+from chai_lab.chai1 import run_inference
+
+# We use fasta-like format for inputs.
+# - each entity encodes protein, ligand, RNA or DNA
+# - each entity is labeled with unique name;
+# - ligands are encoded with SMILES; modified residues encoded like AAA(SEP)AAA
+
+# Example given below, just modify it
+
+
+example_fasta = (
+    """
+>protein|name=example-of-long-protein
+"""
+    + "A" * 1700
+).strip()
+
+fasta_path = Path("/tmp/example.fasta")
+fasta_path.write_text(example_fasta)
+
+output_dir = (
+    Path(__file__)
+    .parent.parent.joinpath(
+        f"outputs/{chai_lab.__version__}/{datetime.datetime.now().isoformat()}"
+    )
+    .absolute()
+)
+
+
+if False:
+    from chai_lab.dispatch import MemSourceTrackingDispatchMode
+
+    with MemSourceTrackingDispatchMode():
+        pass
+if True:
+    x = torch.zeros(85174583296 // 2, dtype=torch.uint8, device=torch.device("cuda:0"))
+    # x = torch.zeros(
+    #     85174583296 // 2 - 1_000_000_000,
+    #     dtype=torch.uint8,
+    #     device=torch.device("cuda:0"),
+    # )
+    # assert 0 == 1
+    candidates = run_inference(
+        fasta_file=fasta_path,
+        output_dir=output_dir,
+        # 'default' setup
+        num_trunk_recycles=2,
+        num_diffn_timesteps=2,
+        seed=42,
+        device="cuda:0",
+        use_esm_embeddings=True,
+        low_memory=True,
+    )
+
+cif_paths = candidates.cif_paths
+scores = [rd.aggregate_score for rd in candidates.ranking_data]
+
+
+# Load pTM, ipTM, pLDDTs and clash scores for sample 2
+scores = np.load(output_dir.joinpath("scores.model_idx_2.npz"))