From e3cb00be94c255f8e36fba1a0ea94006310e8a17 Mon Sep 17 00:00:00 2001
From: Alex Rogozhnikov <iamfullofspam@gmail.com>
Date: Fri, 20 Dec 2024 19:11:19 +0000
Subject: [PATCH] v.0.5.1 (release not out of main)

---
 README.md                             |  4 +-
 chai_lab/__init__.py                  |  2 +-
 examples/predict_structure.private.py | 67 +++++++++++++++++++++++++++
 3 files changed, 70 insertions(+), 3 deletions(-)
 create mode 100644 examples/predict_structure.private.py

diff --git a/README.md b/README.md
index 1b9f2e1..ebc388e 100644
--- a/README.md
+++ b/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
diff --git a/chai_lab/__init__.py b/chai_lab/__init__.py
index 2487274..e0b8595 100644
--- a/chai_lab/__init__.py
+++ b/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"
diff --git a/examples/predict_structure.private.py b/examples/predict_structure.private.py
new file mode 100644
index 0000000..d8f5a1d
--- /dev/null
+++ b/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"))