Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,245 @@
+---
+frameworks:
+- ""
+language:
+- en
+license: mit
+tags:
+- OneScience
+- protein-ligand scoring
+- graph neural network
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">GenScore</span>
+  </strong>
+</p>
+
+# Model Overview
+
+GenScore is a **protein-ligand scoring framework** that can predict protein-small molecule binding affinity and evaluate docking pose quality. It demonstrates balanced scoring, ranking, docking, and virtual screening capabilities across multiple datasets.
+
+# Model Description
+
+GenScore is based on graph neural networks and extends RTMScore.
+
+Weights and datasets are not provided at the moment. They will be uploaded to Hugging Face soon, and command-line downloads will be supported later.
+
+# Use Cases
+
+| Use case | Description |
+| :---: | :---: |
+| Protein-ligand scoring | Predicts binding scores for a given protein, or an extracted binding pocket, and ligand conformation. |
+| Automatic pocket generation | Automatically extracts the binding pocket from a full protein structure based on the reference ligand position. |
+| Contribution analysis | Outputs atom-level and residue-level contributions to the final score for interpretability analysis. |
+| Model training | Trains the GenScore scoring network using preprocessed PDBbind protein-ligand graph data. |
+| CASF-2016 benchmark evaluation | Supports three standard tests: scoring/ranking, docking, and virtual screening. |
+
+# Usage
+
+## 1. Using OneCode
+
+You can try intelligent one-click AI4S programming through the OneCode online environment:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- Running on a GPU or DCU is recommended.
+- CPU can be used for connectivity checks, but it is relatively slow.
+- DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version that matches the current cluster.
+
+**Software Requirements**
+
+DCU users who want to learn more about adaptation details can contact liubiao@sugon.com.
+
+**Environment Checks**
+
+- NVIDIA GPU:
+
+```bash
+nvidia-smi
+```
+
+- Hygon DCU:
+
+```bash
+hy-smi
+```
+
+## Quick Start
+
+### 1. Install the Runtime Environment
+
+```bash
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+If the following code cannot find required libraries at runtime, activate CUDA as shown below.
+
+```bash
+source ${ROCM_PATH}/cuda/env.sh
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"
+```
+
+### 2. Download the Model Package
+
+```bash
+# By default, the package is downloaded to the GenScore folder under the current path. To change this, adjust the path after local_dir.
+hf download --model OneScience-Sugon/GenScore --local-dir ./model
+
+cd model
+```
+
+### Training Weights and Datasets
+
+Weights and datasets are not provided at the moment. They will be uploaded to Hugging Face soon, and command-line downloads will be supported later.
+
+### 3. Usage Modes
+
+#### 1. Protein-Ligand Scoring Inference (`run_genscore.sh`)
+
+```bash
+cd ./scripts
+bash run_genscore.sh
+```
+
+This script runs four inference examples in sequence, covering common use cases:
+
+1. **GT model + automatic pocket generation**: inputs a full protein, a reference ligand, and a decoy ligand, then automatically generates the binding pocket and scores it.
+2. **GatedGCN model + pre-extracted pocket**: inputs a pre-extracted pocket PDB and a decoy ligand for scoring.
+3. **Atom contribution analysis**: uses the GatedGCN model to compute each atom's contribution to the score.
+4. **Residue contribution analysis**: uses the GatedGCN model to compute each residue's contribution to the score.
+
+#### 2. Model Training
+
+##### 2.1 Small-Scale Smoke Test (`train_genscore_smoke.sh`)
+
+```bash
+cd ./scripts
+bash train_genscore_smoke.sh
+```
+
+Default configuration:
+
+| Parameter | Default value |
+|------|--------|
+| Number of training epochs | 100 |
+| Batch size | 16 |
+| Number of validation samples | 1500 |
+| Encoder | `gatedgcn` |
+| Output model | `genscore_smoke_bs16.pth` |
+
+##### 2.2 Full Training (`train_genscore_full.sh`)
+
+```bash
+cd ./scripts
+bash train_genscore_full.sh
+```
+
+Key parameters can be overridden through environment variables:
+
+| Environment variable | Description | Default value |
+|----------|------|--------|
+| `GENSCORE_DATA_DIR` | Training data directory | `${ONESCIENCE_DATASETS_DIR}/GenScore/genscore_data/rtmscore_s` |
+| `GENSCORE_DATA_PREFIX` | Data file prefix | `v2020_train` |
+| `GENSCORE_ENCODER` | Graph encoder type | `gatedgcn` |
+| `GENSCORE_MODEL_PATH` | Output model path | `examples/biosciences/genscore/genscore_${GENSCORE_ENCODER}_full_3000.pth` |
+| `GENSCORE_NUM_EPOCHS` | Number of training epochs | `3000` |
+| `GENSCORE_BATCH_SIZE` | Batch size | `64` |
+| `GENSCORE_NUM_WORKERS` | Number of data loading workers | `8` |
+| `GENSCORE_VALNUM` | Number of validation samples | `1500` |
+| `GENSCORE_PATIENCE` | Early-stopping patience | `70` |
+---
+
+#### 3. CASF-2016 Benchmark Evaluation (`run_genscore_benchmarks.sh`)
+
+```bash
+cd ./scripts
+bash run_genscore_benchmarks.sh all
+```
+
+You can also run a single evaluation task:
+
+```bash
+bash run_genscore_benchmarks.sh scoring
+bash run_genscore_benchmarks.sh docking
+bash run_genscore_benchmarks.sh screening
+```
+
+Evaluation task descriptions:
+
+| Task | Description |
+|------|------|
+| `scoring` | Scoring capability evaluation (scoring/ranking) |
+| `docking` | Docking capability evaluation (docking power) |
+| `screening` | Virtual screening capability evaluation (screening power) |
+
+### Data Preprocessing
+
+Training requires preprocessed PDBbind graph data, including:
+
+```text
+<data_prefix>_ids.npy
+<data_prefix>_lig.pt
+<data_prefix>_prot.pt
+```
+
+You can preprocess the raw PDBbind data with the following command:
+
+```bash
+cd ./scripts
+export PYTHONPATH=../../../src:$PYTHONPATH
+python preprocess_pdbbind.py \
+  --dir /path/to/pdbbind \
+  --ref /path/to/pdbbind_2020_general.csv \
+  --cutoff 10.0 \
+  --outprefix /path/to/preprocessed/pdbbind/v2020_train
+```
+
+Main parameters:
+
+| Parameter | Description |
+|------|------|
+| `--dir` | Raw PDBbind data directory |
+| `--ref` | PDBbind index CSV file |
+| `--cutoff` | Protein-ligand distance cutoff, default `10.0` Å |
+| `--outprefix` | Output file prefix |
+
+### Notes
+
+- Make sure the `ONESCIENCE_DATASETS_DIR` environment variable is correctly set before running the scripts.
+- The scripts automatically set ROCm/DCU-related `LD_LIBRARY_PATH` values and can run directly on Hygon DCU platforms.
+- Automatic pocket generation depends on OpenBabel and ProDy. If OpenBabel requires explicit data paths, set `BABEL_LIBDIR` and `BABEL_DATADIR` in advance.
+- Training and inference use a single device by default (`HIP_VISIBLE_DEVICES=0` or `CUDA_VISIBLE_DEVICES=0`). For multi-device training, adjust the parallel strategy in the training scripts.
+- The training scripts check whether `<data_prefix>_ids.npy`, `<data_prefix>_lig.pt`, and `<data_prefix>_prot.pt` exist. If any are missing, the scripts will report an error and exit.
+- The `docking` and `screening` tasks in CASF-2016 evaluation require complete CASF-2016 and PDBbind v2020 data.
+
+# Official OneScience Information
+
+| Platform | OneScience main repository | Skills repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+If you use GenScore in your research, please cite the original work:
+
+```bibtex
+@article{genScore,
+  title={GenScore: a generalized protein-ligand scoring framework},
+  author={Shen, Chao and Hu, Yafeng and Wang, Zhe and Zhang, Xujun and Li, Jianxin and Wang, Guisheng and Wang, Tingjun and Chen, Yen-Wei and Pan, Peichen and Hou, Tingjun},
+  journal={Journal of Chemical Information and Modeling},
+  year={2023}
+}
+```
+
+For more information, see the official GenScore repository: https://github.com/sc8668/GenScore
+
+For license information, refer to the description in the [official GenScore repository](https://github.com/sc8668/GenScore).

benchmarks/_common.py

@@ -0,0 +1,102 @@
+import sys
+from pathlib import Path
+
+_DIR = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(_DIR))
+
+import argparse
+
+import torch as th
+import torch.multiprocessing
+from torch_geometric.loader import DataLoader
+
+from onescience.datapipes.genscore.data import PDBbindDataset
+from onescience.metrics.genscore.utils import run_an_eval_epoch
+from models.inference import _build_encoder, scoring
+from models.model.model import GenScore
+
+torch.multiprocessing.set_sharing_strategy("file_system")
+
+
+def add_model_args(parser):
+    parser.add_argument("--model-path", required=True, help="Path to a trained GenScore checkpoint.")
+    parser.add_argument("--encoder", choices=["gt", "gatedgcn"], default="gatedgcn")
+    parser.add_argument("--batch-size", type=int, default=128)
+    parser.add_argument("--num-workers", type=int, default=10)
+    parser.add_argument("--cutoff", type=float, default=10.0)
+    parser.add_argument("--outprefix", default="gatedgcn1x5")
+    parser.add_argument("--dist-threhold", type=float, default=5.0)
+    parser.add_argument("--hidden-dim0", type=int, default=128)
+    parser.add_argument("--hidden-dim", type=int, default=128)
+    parser.add_argument("--n-gaussians", type=int, default=10)
+    parser.add_argument("--dropout-rate", type=float, default=0.15)
+
+
+def runtime_kwargs(args):
+    return {
+        "batch_size": args.batch_size,
+        "dist_threhold": args.dist_threhold,
+        "device": "cuda" if th.cuda.is_available() else "cpu",
+        "num_workers": args.num_workers,
+        "num_node_featsp": 41,
+        "num_node_featsl": 41,
+        "num_edge_featsp": 5,
+        "num_edge_featsl": 10,
+        "hidden_dim0": args.hidden_dim0,
+        "hidden_dim": args.hidden_dim,
+        "n_gaussians": args.n_gaussians,
+        "dropout_rate": args.dropout_rate,
+    }
+
+
+def score_ligand_file(prot, lig, args, parallel=False):
+    return scoring(
+        prot=prot,
+        lig=lig,
+        modpath=args.model_path,
+        cut=args.cutoff,
+        gen_pocket=False,
+        reflig=None,
+        encoder=args.encoder,
+        explicit_H=False,
+        use_chirality=True,
+        parallel=parallel,
+        **runtime_kwargs(args),
+    )
+
+
+def score_preprocessed(ids, prots, ligs, args):
+    kwargs = runtime_kwargs(args)
+    data = PDBbindDataset(ids=ids, prots=prots, ligs=ligs)
+    loader = DataLoader(
+        dataset=data,
+        batch_size=kwargs["batch_size"],
+        shuffle=False,
+        num_workers=kwargs["num_workers"],
+    )
+
+    ligmodel, protmodel = _build_encoder(args.encoder, kwargs)
+    model = GenScore(
+        ligmodel,
+        protmodel,
+        in_channels=kwargs["hidden_dim0"],
+        hidden_dim=kwargs["hidden_dim"],
+        n_gaussians=kwargs["n_gaussians"],
+        dropout_rate=kwargs["dropout_rate"],
+        dist_threhold=kwargs["dist_threhold"],
+    ).to(kwargs["device"])
+
+    checkpoint = th.load(args.model_path, map_location=th.device(kwargs["device"]))
+    model.load_state_dict(checkpoint["model_state_dict"])
+    preds = run_an_eval_epoch(
+        model,
+        loader,
+        pred=True,
+        dist_threhold=kwargs["dist_threhold"],
+        device=kwargs["device"],
+    )
+    return data.pdbids, preds
+
+
+def formatter():
+    return argparse.ArgumentDefaultsHelpFormatter

benchmarks/casf2016_docking.py

@@ -0,0 +1,142 @@
+import argparse
+import os
+import pickle
+
+import numpy as np
+import pandas as pd
+from joblib import Parallel, delayed
+
+from _common import add_model_args, formatter, runtime_kwargs, score_ligand_file
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Run GenScore CASF-2016 docking benchmark.",
+        formatter_class=formatter(),
+    )
+    add_model_args(parser)
+    parser.add_argument("--casf-dir", required=True, help="CASF-2016 root directory.")
+    parser.add_argument("--pdbbind-dir", required=True, help="PDBbind root directory.")
+    parser.add_argument(
+        "--native-ligand-dir",
+        default=None,
+        help="Directory containing <pdbid>_ligand.mol2 files.",
+    )
+    parser.add_argument("--outdir", required=True, help="Directory for CASF docking score .dat files.")
+    parser.add_argument("--decoys-subdir", default="decoys_docking")
+    parser.add_argument("--coreset-subdir", default="coreset")
+    parser.add_argument("--refined-subdir", default="v2020-refined")
+    parser.add_argument("--other-pl-subdir", default="v2020-other-PL")
+    parser.add_argument("--parallel", action="store_true", default=False)
+    return parser.parse_args()
+
+
+def _complex_ids(path):
+    return [name for name in os.listdir(path) if os.path.isdir(os.path.join(path, name))]
+
+
+def _pocket_path(pdbbind_dir, subset, pdbid, cutoff):
+    return os.path.join(
+        pdbbind_dir,
+        subset,
+        pdbid,
+        f"{pdbid}_prot",
+        f"{pdbid}_p_pocket_{cutoff}.pdb",
+    )
+
+
+def _native_ligand_path(pdbbind_dir, native_ligand_dir, subset, pdbid):
+    candidates = []
+    if native_ligand_dir:
+        candidates.append(
+            os.path.join(native_ligand_dir, f"{pdbid}_ligand.mol2")
+        )
+    candidates.extend([
+        os.path.join(pdbbind_dir, "mol2", f"{pdbid}_ligand.mol2"),
+        os.path.join(
+            pdbbind_dir, subset, pdbid, f"{pdbid}_prot",
+            f"{pdbid}_l.mol2",
+        ),
+        os.path.join(
+            pdbbind_dir, subset, pdbid, f"{pdbid}_prot",
+            f"{pdbid}_l.sdf",
+        ),
+    ])
+    for candidate in candidates:
+        if os.path.isfile(candidate):
+            return candidate
+    raise FileNotFoundError(
+        f"No native ligand found for {pdbid}. Checked: "
+        + ", ".join(candidates)
+    )
+
+
+def _decoys_path(casf_dir, decoys_subdir, pdbid):
+    root = os.path.join(casf_dir, decoys_subdir)
+    return _ligand_path(root, f"{pdbid}_decoys")
+
+
+def _ligand_path(root, stem):
+    candidates = [
+        os.path.join(root, f"{stem}.{extension}")
+        for extension in ("sdf", "mol2")
+    ]
+    for candidate in candidates:
+        if os.path.isfile(candidate):
+            return candidate
+    raise FileNotFoundError(
+        f"No ligand file found for {stem}. Checked: {', '.join(candidates)}"
+    )
+
+
+def score_compound(pdbid, subset, args):
+    prot = _pocket_path(args.pdbbind_dir, subset, pdbid, args.cutoff)
+    decoys = _decoys_path(args.casf_dir, args.decoys_subdir, pdbid)
+    ids_decoy, scores_decoy = score_ligand_file(prot, decoys, args, parallel=True)
+
+    native = _native_ligand_path(
+        args.pdbbind_dir,
+        args.native_ligand_dir,
+        subset,
+        pdbid,
+    )
+    ids_native, scores_native = score_ligand_file(prot, native, args, parallel=False)
+    ids_native = list(ids_native)
+    ids_native.pop(-1)
+    ids_native.append(f"{pdbid}_ligand")
+    return ids_decoy + ids_native, np.append(scores_decoy, scores_native)
+
+
+def run_for_subset(pdbids, subset, args):
+    if args.parallel and runtime_kwargs(args)["device"] == "cpu":
+        return Parallel(n_jobs=-1, backend="threading")(
+            delayed(score_compound)(pdbid, subset, args) for pdbid in pdbids
+        )
+    return [score_compound(pdbid, subset, args) for pdbid in pdbids]
+
+
+def main():
+    args = parse_args()
+    coreset = _complex_ids(os.path.join(args.casf_dir, args.coreset_subdir))
+    refined_ids = set(_complex_ids(os.path.join(args.pdbbind_dir, args.refined_subdir)))
+    other_ids = set(_complex_ids(os.path.join(args.pdbbind_dir, args.other_pl_subdir)))
+
+    ids_refined = [pdbid for pdbid in coreset if pdbid in refined_ids]
+    ids_other = [pdbid for pdbid in coreset if pdbid in other_ids]
+
+    results = run_for_subset(ids_refined, args.refined_subdir, args)
+    results += run_for_subset(ids_other, args.other_pl_subdir, args)
+
+    os.makedirs(args.outdir, exist_ok=True)
+    for ids, scores in results:
+        pdbid = ids[0].split("_")[0]
+        df = pd.DataFrame(zip(ids, scores), columns=["#code", "score"])
+        df["#code"] = df["#code"].str.split("-").apply(lambda item: item[0])
+        df.to_csv(os.path.join(args.outdir, f"{pdbid}_score.dat"), index=False, sep="\t")
+
+    with open(f"{args.outprefix}_docking.pkl", "wb") as handle:
+        pickle.dump(results, handle)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/casf2016_scoring_ranking.py

@@ -0,0 +1,62 @@
+import argparse
+import os
+
+import numpy as np
+import pandas as pd
+from scipy.stats import pearsonr
+from sklearn import linear_model
+from sklearn.metrics import mean_squared_error
+
+from _common import add_model_args, formatter, score_preprocessed
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Run GenScore CASF-2016 scoring/ranking benchmark.",
+        formatter_class=formatter(),
+    )
+    add_model_args(parser)
+    parser.add_argument("--data-dir", required=True, help="Directory containing preprocessed CASF tensors.")
+    parser.add_argument("--test-prefix", default="v2020_casf")
+    parser.add_argument("--coreset-file", required=True, help="Path to CASF-2016 CoreSet.dat.")
+    parser.add_argument("--outdir", required=True, help="Output directory for CASF ranking .dat file.")
+    return parser.parse_args()
+
+
+def obtain_metrics(df):
+    regr = linear_model.LinearRegression()
+    regr.fit(df.score.values.reshape(-1, 1), df.logKa.values.reshape(-1, 1))
+    preds = regr.predict(df.score.values.reshape(-1, 1))
+    rp = pearsonr(df.logKa, df.score)[0]
+    mse = mean_squared_error(df.logKa, preds)
+    num = df.shape[0]
+    sd = np.sqrt((mse * num) / (num - 1))
+    print("The regression equation: logKa = %.2f + %.2f * Score" % (float(regr.coef_), float(regr.intercept_)))
+    print("Number of favorable sample (N): %d" % num)
+    print("Pearson correlation coefficient (R): %.3f" % rp)
+    print("Standard deviation in fitting (SD): %.2f" % sd)
+
+
+def main():
+    args = parse_args()
+    prots = os.path.join(args.data_dir, f"{args.test_prefix}_prot.pt")
+    ligs = os.path.join(args.data_dir, f"{args.test_prefix}_lig.pt")
+    ids = os.path.join(args.data_dir, f"{args.test_prefix}_ids.npy")
+
+    _, preds = score_preprocessed(ids, prots, ligs, args)
+
+    core = pd.read_csv(args.coreset_file, sep=r"[,,\t, ]+", header=0, engine="python")
+    df_score = pd.DataFrame(zip(np.load(ids, allow_pickle=True)[0], preds), columns=["#code", "score"])
+    testdf = pd.merge(core, df_score, on="#code")
+
+    os.makedirs(args.outdir, exist_ok=True)
+    testdf[["#code", "score"]].to_csv(
+        os.path.join(args.outdir, f"{args.outprefix}.dat"),
+        index=False,
+        sep="\t",
+    )
+    obtain_metrics(testdf)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/casf2016_screening.py

@@ -0,0 +1,104 @@
+import argparse
+import os
+import pickle
+
+import pandas as pd
+from joblib import Parallel, delayed
+
+from _common import add_model_args, formatter, runtime_kwargs, score_ligand_file
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Run GenScore CASF-2016 screening benchmark.",
+        formatter_class=formatter(),
+    )
+    add_model_args(parser)
+    parser.add_argument("--casf-dir", required=True, help="CASF-2016 root directory.")
+    parser.add_argument("--pdbbind-dir", required=True, help="PDBbind root directory.")
+    parser.add_argument("--outdir", required=True, help="Directory for CASF screening score .dat files.")
+    parser.add_argument("--decoys-subdir", default="decoys_screening")
+    parser.add_argument("--coreset-subdir", default="coreset")
+    parser.add_argument("--refined-subdir", default="v2020-refined")
+    parser.add_argument("--other-pl-subdir", default="v2020-other-PL")
+    parser.add_argument("--parallel", action="store_true", default=False)
+    return parser.parse_args()
+
+
+def _complex_ids(path):
+    return [name for name in os.listdir(path) if os.path.isdir(os.path.join(path, name))]
+
+
+def _pocket_path(pdbbind_dir, subset, pdbid, cutoff):
+    return os.path.join(
+        pdbbind_dir,
+        subset,
+        pdbid,
+        f"{pdbid}_prot",
+        f"{pdbid}_p_pocket_{cutoff}.pdb",
+    )
+
+
+def _ligand_path(root, stem):
+    candidates = [os.path.join(root, f"{stem}.{extension}") for extension in ("mol2", "sdf")]
+    for path in candidates:
+        if os.path.isfile(path):
+            return path
+    raise FileNotFoundError(
+        f"No ligand file found for {stem}. Checked: {', '.join(candidates)}"
+    )
+
+
+def score_compound(pdbid, ligid, subset, args):
+    prot = _pocket_path(args.pdbbind_dir, subset, pdbid, args.cutoff)
+    root = os.path.join(args.casf_dir, args.decoys_subdir, pdbid)
+    lig = _ligand_path(root, f"{pdbid}_{ligid}")
+    return score_ligand_file(prot, lig, args, parallel=True)
+
+
+def score_target(pdbid, subset, ligids, args):
+    print(f"{pdbid} started.....")
+    ids_list = []
+    scores_list = []
+    for ligid in ligids:
+        ids, scores = score_compound(pdbid, ligid, subset, args)
+        if ids is not None and scores is not None:
+            ids_list.extend(ids)
+            scores_list.extend(scores)
+    print(f"{pdbid} finished.....")
+    return pdbid, [ids_list, scores_list]
+
+
+def run_for_subset(pdbids, subset, ligids, args):
+    if args.parallel and runtime_kwargs(args)["device"] == "cpu":
+        return Parallel(n_jobs=-1, backend="threading")(
+            delayed(score_target)(pdbid, subset, ligids, args) for pdbid in pdbids
+        )
+    return [score_target(pdbid, subset, ligids, args) for pdbid in pdbids]
+
+
+def main():
+    args = parse_args()
+    ligids = _complex_ids(os.path.join(args.casf_dir, args.coreset_subdir))
+    pdbids = _complex_ids(os.path.join(args.casf_dir, args.decoys_subdir))
+    refined_ids = set(_complex_ids(os.path.join(args.pdbbind_dir, args.refined_subdir)))
+    other_ids = set(_complex_ids(os.path.join(args.pdbbind_dir, args.other_pl_subdir)))
+
+    ids_refined = [pdbid for pdbid in pdbids if pdbid in refined_ids]
+    ids_other = [pdbid for pdbid in pdbids if pdbid in other_ids]
+
+    results = run_for_subset(ids_refined, args.refined_subdir, ligids, args)
+    results += run_for_subset(ids_other, args.other_pl_subdir, ligids, args)
+
+    os.makedirs(args.outdir, exist_ok=True)
+    for pdbid, payload in results:
+        df = pd.DataFrame(zip(*payload), columns=["#code_ligand_num", "score"])
+        df["#code_ligand_num"] = df["#code_ligand_num"].str.split("-").apply(lambda item: item[0])
+        df.to_csv(os.path.join(args.outdir, f"{pdbid}_score.dat"), index=False, sep="\t")
+
+    with open(f"{args.outprefix}_screening.pkl", "wb") as handle:
+        pickle.dump(results, handle)
+
+
+if __name__ == "__main__":
+    main()

models/__init__.py

@@ -0,0 +1,2 @@
+"""GenScore protein-ligand scoring model integration."""
+

models/inference.py

@@ -0,0 +1,298 @@
+import argparse
+import os
+
+import MDAnalysis as mda
+import numpy as np
+import pandas as pd
+import torch as th
+import torch.multiprocessing
+from torch_geometric.loader import DataLoader
+
+from onescience.datapipes.genscore.data import VSDataset
+from models.model.model import GatedGCN, GenScore, GraphTransformer
+from onescience.metrics.genscore.utils import run_an_eval_epoch
+
+torch.multiprocessing.set_sharing_strategy("file_system")
+
+
+def _default_model_path() -> str:
+    return os.path.join(
+        os.path.dirname(__file__),
+        "..",
+        "..",
+        "..",
+        "..",
+        "examples",
+        "biosciences",
+        "genscore",
+        "trained_models",
+        "GT_0.0_1.pth",
+    )
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Run GenScore protein-ligand scoring.")
+    parser.add_argument("-p", "--prot", required=True, help="Input protein or pocket file (.pdb).")
+    parser.add_argument("-l", "--lig", required=True, help="Input ligand file (.sdf/.mol2).")
+    parser.add_argument(
+        "-m",
+        "--model",
+        default=_default_model_path(),
+        help="Path to a trained GenScore checkpoint.",
+    )
+    parser.add_argument(
+        "-e",
+        "--encoder",
+        default="gt",
+        choices=["gt", "gatedgcn"],
+        help="Protein and ligand graph encoder.",
+    )
+    parser.add_argument("-o", "--outprefix", default="out", help="Output file prefix.")
+    parser.add_argument(
+        "-gen_pocket",
+        "--gen_pocket",
+        action="store_true",
+        default=False,
+        help="Generate a pocket from the input protein.",
+    )
+    parser.add_argument(
+        "-c",
+        "--cutoff",
+        default=10.0,
+        type=float,
+        help="Pocket and interaction cutoff distance.",
+    )
+    parser.add_argument("-rl", "--reflig", default=None, help="Reference ligand for pocket generation.")
+    parser.add_argument(
+        "-pl",
+        "--parallel",
+        default=False,
+        action="store_true",
+        help="Build ligand graphs in parallel.",
+    )
+    parser.add_argument(
+        "-ac",
+        "--atom_contribution",
+        default=False,
+        action="store_true",
+        help="Compute atom-level score contributions.",
+    )
+    parser.add_argument(
+        "-rc",
+        "--res_contribution",
+        default=False,
+        action="store_true",
+        help="Compute residue-level score contributions.",
+    )
+    parser.add_argument("--batch_size", type=int, default=128)
+    parser.add_argument("--num_workers", type=int, default=10)
+    args = parser.parse_args()
+
+    if args.gen_pocket and args.reflig is None:
+        raise ValueError("If pocket generation is enabled, --reflig must be provided.")
+    if args.atom_contribution and args.res_contribution:
+        raise ValueError("Only one of atom_contribution and res_contribution is supported.")
+    return args
+
+
+def _build_encoder(encoder, kwargs):
+    if encoder == "gt":
+        ligmodel = GraphTransformer(
+            in_channels=kwargs["num_node_featsl"],
+            edge_features=kwargs["num_edge_featsl"],
+            num_hidden_channels=kwargs["hidden_dim0"],
+            activ_fn=th.nn.SiLU(),
+            transformer_residual=True,
+            num_attention_heads=4,
+            norm_to_apply="batch",
+            dropout_rate=0.15,
+            num_layers=6,
+        )
+        protmodel = GraphTransformer(
+            in_channels=kwargs["num_node_featsp"],
+            edge_features=kwargs["num_edge_featsp"],
+            num_hidden_channels=kwargs["hidden_dim0"],
+            activ_fn=th.nn.SiLU(),
+            transformer_residual=True,
+            num_attention_heads=4,
+            norm_to_apply="batch",
+            dropout_rate=0.15,
+            num_layers=6,
+        )
+    else:
+        ligmodel = GatedGCN(
+            in_channels=kwargs["num_node_featsl"],
+            edge_features=kwargs["num_edge_featsl"],
+            num_hidden_channels=kwargs["hidden_dim0"],
+            residual=True,
+            dropout_rate=0.15,
+            equivstable_pe=False,
+            num_layers=6,
+        )
+        protmodel = GatedGCN(
+            in_channels=kwargs["num_node_featsp"],
+            edge_features=kwargs["num_edge_featsp"],
+            num_hidden_channels=kwargs["hidden_dim0"],
+            residual=True,
+            dropout_rate=0.15,
+            equivstable_pe=False,
+            num_layers=6,
+        )
+    return ligmodel, protmodel
+
+
+def scoring(
+    prot,
+    lig,
+    modpath,
+    cut=10.0,
+    gen_pocket=False,
+    reflig=None,
+    encoder="gt",
+    atom_contribution=False,
+    res_contribution=False,
+    explicit_H=False,
+    use_chirality=True,
+    parallel=False,
+    **kwargs,
+):
+    data = VSDataset(
+        ligs=lig,
+        prot=prot,
+        cutoff=cut,
+        gen_pocket=gen_pocket,
+        reflig=reflig,
+        explicit_H=explicit_H,
+        use_chirality=use_chirality,
+        parallel=parallel,
+    )
+    test_loader = DataLoader(
+        dataset=data,
+        batch_size=kwargs["batch_size"],
+        shuffle=False,
+        num_workers=kwargs["num_workers"],
+    )
+
+    ligmodel, protmodel = _build_encoder(encoder, kwargs)
+    model = GenScore(
+        ligmodel,
+        protmodel,
+        in_channels=kwargs["hidden_dim0"],
+        hidden_dim=kwargs["hidden_dim"],
+        n_gaussians=kwargs["n_gaussians"],
+        dropout_rate=kwargs["dropout_rate"],
+        dist_threhold=kwargs["dist_threhold"],
+    ).to(kwargs["device"])
+
+    checkpoint = th.load(modpath, map_location=th.device(kwargs["device"]))
+    model.load_state_dict(checkpoint["model_state_dict"])
+
+    if atom_contribution:
+        preds, at_contrs, _ = run_an_eval_epoch(
+            model,
+            test_loader,
+            pred=True,
+            atom_contribution=True,
+            res_contribution=False,
+            dist_threhold=kwargs["dist_threhold"],
+            device=kwargs["device"],
+        )
+        atids = [f"{a.GetSymbol()}{a.GetIdx()}" for a in data.ligs[0].GetAtoms()]
+        return data.ids, preds, atids, at_contrs
+
+    if res_contribution:
+        preds, _, res_contrs = run_an_eval_epoch(
+            model,
+            test_loader,
+            pred=True,
+            atom_contribution=False,
+            res_contribution=True,
+            dist_threhold=kwargs["dist_threhold"],
+            device=kwargs["device"],
+        )
+        universe = mda.Universe(data.prot)
+        resids = [
+            f"{chain_id}_{resname}{resid}"
+            for chain_id, resname, resid in zip(
+                universe.residues.chainIDs,
+                universe.residues.resnames,
+                universe.residues.resids,
+            )
+        ]
+        return data.ids, preds, resids, res_contrs
+
+    preds = run_an_eval_epoch(
+        model,
+        test_loader,
+        pred=True,
+        dist_threhold=kwargs["dist_threhold"],
+        device=kwargs["device"],
+    )
+    return data.ids, preds
+
+
+def _runtime_args(args):
+    return {
+        "batch_size": args.batch_size,
+        "dist_threhold": 5.0,
+        "device": "cuda" if th.cuda.is_available() else "cpu",
+        "num_workers": args.num_workers,
+        "num_node_featsp": 41,
+        "num_node_featsl": 41,
+        "num_edge_featsp": 5,
+        "num_edge_featsl": 10,
+        "hidden_dim0": 128,
+        "hidden_dim": 128,
+        "n_gaussians": 10,
+        "dropout_rate": 0.15,
+    }
+
+
+def main():
+    inargs = parse_args()
+    runtime = _runtime_args(inargs)
+    common = {
+        "prot": inargs.prot,
+        "lig": inargs.lig,
+        "modpath": inargs.model,
+        "cut": inargs.cutoff,
+        "gen_pocket": inargs.gen_pocket,
+        "reflig": inargs.reflig,
+        "encoder": inargs.encoder,
+        "explicit_H": False,
+        "use_chirality": True,
+        "parallel": inargs.parallel,
+        **runtime,
+    }
+
+    if inargs.atom_contribution:
+        ids, scores, atids, at_contrs = scoring(atom_contribution=True, **common)
+        df = pd.DataFrame(at_contrs).T
+        df.columns = ids
+        df.index = atids
+        df = df[df.apply(np.sum, axis=1) != 0].T
+        dfx = pd.DataFrame(zip(*(ids, scores)), columns=["id", "score"])
+        dfx.index = dfx.id
+        df = pd.concat([dfx["score"], df], axis=1)
+        df.sort_values("score", ascending=False, inplace=True)
+        df.to_csv(f"{inargs.outprefix}_at.csv")
+    elif inargs.res_contribution:
+        ids, scores, resids, res_contrs = scoring(res_contribution=True, **common)
+        df = pd.DataFrame(res_contrs).T
+        df.columns = ids
+        df.index = resids
+        df = df[df.apply(np.sum, axis=1) != 0].T
+        dfx = pd.DataFrame(zip(*(ids, scores)), columns=["id", "score"])
+        dfx.index = dfx.id
+        df = pd.concat([dfx["score"], df], axis=1)
+        df.sort_values("score", ascending=False, inplace=True)
+        df.to_csv(f"{inargs.outprefix}_res.csv")
+    else:
+        ids, scores = scoring(**common)
+        df = pd.DataFrame(zip(*(ids, scores)), columns=["id", "score"])
+        df.sort_values("score", ascending=False, inplace=True)
+        df.to_csv(f"{inargs.outprefix}.csv", index=False)
+
+
+if __name__ == "__main__":
+    main()

models/model/__init__.py

@@ -0,0 +1 @@
+

models/model/layer/__init__.py

@@ -0,0 +1 @@
+

models/model/layer/gatedgcn_layer.py

@@ -0,0 +1,151 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch_geometric.nn as pyg_nn
+from torch_geometric.graphgym.models.layer import LayerConfig
+from torch_scatter import scatter
+
+from torch_geometric.graphgym.config import cfg
+from torch_geometric.graphgym.register import register_layer
+
+
+class GatedGCNLayer(pyg_nn.conv.MessagePassing):
+    """
+        GatedGCN layer
+        Residual Gated Graph ConvNets
+        https://arxiv.org/pdf/1711.07553.pdf
+    """
+    def __init__(self, in_dim, out_dim, dropout, residual,
+                 equivstable_pe=False, **kwargs):
+        super().__init__(**kwargs)
+        self.A = pyg_nn.Linear(in_dim, out_dim, bias=True)
+        self.B = pyg_nn.Linear(in_dim, out_dim, bias=True)
+        self.C = pyg_nn.Linear(in_dim, out_dim, bias=True)
+        self.D = pyg_nn.Linear(in_dim, out_dim, bias=True)
+        self.E = pyg_nn.Linear(in_dim, out_dim, bias=True)
+
+        # Handling for Equivariant and Stable PE using LapPE
+        # ICLR 2022 https://openreview.net/pdf?id=e95i1IHcWj
+        self.EquivStablePE = equivstable_pe
+        if self.EquivStablePE:
+            self.mlp_r_ij = nn.Sequential(
+                nn.Linear(1, out_dim), nn.ReLU(),
+                nn.Linear(out_dim, 1),
+                nn.Sigmoid())
+
+        self.bn_node_x = nn.BatchNorm1d(out_dim)
+        self.bn_edge_e = nn.BatchNorm1d(out_dim)
+        self.dropout = dropout
+        self.residual = residual
+        self.e = None
+
+    def forward(self, batch):
+        x, e, edge_index = batch.x, batch.edge_attr, batch.edge_index
+
+        """
+        x               : [n_nodes, in_dim]
+        e               : [n_edges, in_dim]
+        edge_index      : [2, n_edges]
+        """
+        if self.residual:
+            x_in = x
+            e_in = e
+
+        Ax = self.A(x)
+        Bx = self.B(x)
+        Ce = self.C(e)
+        Dx = self.D(x)
+        Ex = self.E(x)
+
+        # Handling for Equivariant and Stable PE using LapPE
+        # ICLR 2022 https://openreview.net/pdf?id=e95i1IHcWj
+        pe_LapPE = batch.pe_EquivStableLapPE if self.EquivStablePE else None
+
+        x, e = self.propagate(edge_index,
+                              Bx=Bx, Dx=Dx, Ex=Ex, Ce=Ce,
+                              e=e, Ax=Ax,
+                              PE=pe_LapPE)
+
+        x = self.bn_node_x(x)
+        e = self.bn_edge_e(e)
+
+        x = F.relu(x)
+        e = F.relu(e)
+
+        x = F.dropout(x, self.dropout, training=self.training)
+        e = F.dropout(e, self.dropout, training=self.training)
+
+        if self.residual:
+            x = x_in + x
+            e = e_in + e
+
+        batch.x = x
+        batch.edge_attr = e
+
+        return batch
+
+    def message(self, Dx_i, Ex_j, PE_i, PE_j, Ce):
+        """
+        {}x_i           : [n_edges, out_dim]
+        {}x_j           : [n_edges, out_dim]
+        {}e             : [n_edges, out_dim]
+        """
+        e_ij = Dx_i + Ex_j + Ce
+        sigma_ij = torch.sigmoid(e_ij)
+
+        # Handling for Equivariant and Stable PE using LapPE
+        # ICLR 2022 https://openreview.net/pdf?id=e95i1IHcWj
+        if self.EquivStablePE:
+            r_ij = ((PE_i - PE_j) ** 2).sum(dim=-1, keepdim=True)
+            r_ij = self.mlp_r_ij(r_ij)  # the MLP is 1 dim --> hidden_dim --> 1 dim
+            sigma_ij = sigma_ij * r_ij
+
+        self.e = e_ij
+        return sigma_ij
+
+    def aggregate(self, sigma_ij, index, Bx_j, Bx):
+        """
+        sigma_ij        : [n_edges, out_dim]  ; is the output from message() function
+        index           : [n_edges]
+        {}x_j           : [n_edges, out_dim]
+        """
+        dim_size = Bx.shape[0]  # or None ??   <--- Double check this
+
+        sum_sigma_x = sigma_ij * Bx_j
+        numerator_eta_xj = scatter(sum_sigma_x, index, 0, None, dim_size,
+                                   reduce='sum')
+
+        sum_sigma = sigma_ij
+        denominator_eta_xj = scatter(sum_sigma, index, 0, None, dim_size,
+                                     reduce='sum')
+
+        out = numerator_eta_xj / (denominator_eta_xj + 1e-6)
+        return out
+
+    def update(self, aggr_out, Ax):
+        """
+        aggr_out        : [n_nodes, out_dim] ; is the output from aggregate() function after the aggregation
+        {}x             : [n_nodes, out_dim]
+        """
+        x = Ax + aggr_out
+        e_out = self.e
+        del self.e
+        return x, e_out
+
+
+@register_layer('gatedgcnconv')
+class GatedGCNGraphGymLayer(nn.Module):
+    """GatedGCN layer.
+    Residual Gated Graph ConvNets
+    https://arxiv.org/pdf/1711.07553.pdf
+    """
+    def __init__(self, layer_config: LayerConfig, **kwargs):
+        super().__init__()
+        self.model = GatedGCNLayer(in_dim=layer_config.dim_in,
+                                   out_dim=layer_config.dim_out,
+                                   dropout=0.,  # Dropout is handled by GraphGym's `GeneralLayer` wrapper
+                                   residual=False,  # Residual connections are handled by GraphGym's `GNNStackStage` wrapper
+                                   **kwargs)
+
+    def forward(self, batch):
+        return self.model(batch)

models/model/model.py

@@ -0,0 +1,563 @@
+import torch as th
+import torch.nn.functional as F
+import numpy as np
+from torch_scatter import scatter_add
+from torch_geometric.utils import to_dense_batch
+from torch import nn
+
+
+def glorot_orthogonal(tensor, scale):
+	"""Initialize a tensor's values according to an orthogonal Glorot initialization scheme."""
+	if tensor is not None:
+		th.nn.init.orthogonal_(tensor.data)
+		scale /= ((tensor.size(-2) + tensor.size(-1)) * tensor.var())
+		tensor.data *= scale.sqrt()
+
+
+class MultiHeadAttentionLayer(nn.Module):
+	"""Compute attention scores with a DGLGraph's node and edge (geometric) features."""
+	def __init__(self, num_input_feats, num_output_feats,
+				num_heads, using_bias=False, update_edge_feats=True):
+		super(MultiHeadAttentionLayer, self).__init__()
+		
+        # Declare shared variables
+		self.num_output_feats = num_output_feats
+		self.num_heads = num_heads
+		self.using_bias = using_bias
+		self.update_edge_feats = update_edge_feats
+		
+		# Define node features' query, key, and value tensors, and define edge features' projection tensors
+		self.Q = nn.Linear(num_input_feats, self.num_output_feats * self.num_heads, bias=using_bias)
+		self.K = nn.Linear(num_input_feats, self.num_output_feats * self.num_heads, bias=using_bias)
+		self.V = nn.Linear(num_input_feats, self.num_output_feats * self.num_heads, bias=using_bias)
+		self.edge_feats_projection = nn.Linear(num_input_feats, self.num_output_feats * self.num_heads, bias=using_bias)
+		
+		self.reset_parameters()
+		
+	def reset_parameters(self):
+		"""Reinitialize learnable parameters."""
+		scale = 2.0
+		if self.using_bias:
+			glorot_orthogonal(self.Q.weight, scale=scale)
+			self.Q.bias.data.fill_(0)
+			
+			glorot_orthogonal(self.K.weight, scale=scale)
+			self.K.bias.data.fill_(0)
+			
+			glorot_orthogonal(self.V.weight, scale=scale)
+			self.V.bias.data.fill_(0)
+			
+			glorot_orthogonal(self.edge_feats_projection.weight, scale=scale)
+			self.edge_feats_projection.bias.data.fill_(0)
+		else:
+			glorot_orthogonal(self.Q.weight, scale=scale)
+			glorot_orthogonal(self.K.weight, scale=scale)
+			glorot_orthogonal(self.V.weight, scale=scale)
+			glorot_orthogonal(self.edge_feats_projection.weight, scale=scale)
+	
+	def propagate_attention(self, edge_index, node_feats_q, node_feats_k, node_feats_v, edge_feats_projection):
+		row, col = edge_index
+		e_out = None
+		# Compute attention scores
+		alpha = node_feats_k[row] * node_feats_q[col]
+		# Scale and clip attention scores
+		alpha = (alpha / np.sqrt(self.num_output_feats)).clamp(-5.0,5.0)
+		# Use available edge features to modify the attention scores
+		alpha = alpha * edge_feats_projection
+		# Copy edge features as e_out to be passed to edge_feats_MLP
+		if self.update_edge_feats:
+			e_out = alpha
+		
+		# Apply softmax to attention scores, followed by clipping	
+		alphax = th.exp((alpha.sum(-1, keepdim=True)).clamp(-5.0,5.0))
+		# Send weighted values to target nodes
+		wV = scatter_add(node_feats_v[row]*alphax, col, dim=0, dim_size=node_feats_q.size(0))
+		z = scatter_add(alphax, col, dim=0, dim_size=node_feats_q.size(0))
+		return wV, z, e_out
+	
+	def forward(self, x, edge_attr, edge_index):
+		node_feats_q = self.Q(x).view(-1, self.num_heads, self.num_output_feats)
+		node_feats_k = self.K(x).view(-1, self.num_heads, self.num_output_feats)
+		node_feats_v = self.V(x).view(-1, self.num_heads, self.num_output_feats)
+		edge_feats_projection = self.edge_feats_projection(edge_attr).view(-1, self.num_heads, self.num_output_feats)	
+		wV, z, e_out = self.propagate_attention(edge_index, node_feats_q, node_feats_k, node_feats_v, edge_feats_projection)
+		
+		h_out = wV / (z + th.full_like(z, 1e-6))
+		return h_out, e_out	
+
+
+class GraphTransformerModule(nn.Module):
+	"""A Graph Transformer module (equivalent to one layer of graph convolutions)."""
+	def __init__(
+			self,
+			num_hidden_channels,
+			activ_fn=nn.SiLU(),
+			residual=True,
+			num_attention_heads=4,
+			norm_to_apply='batch',
+			dropout_rate=0.1,
+			num_layers=4,
+			):
+		super(GraphTransformerModule, self).__init__()
+		
+		# Record parameters given
+		self.activ_fn = activ_fn
+		self.residual = residual
+		self.num_attention_heads = num_attention_heads
+		self.norm_to_apply = norm_to_apply
+		self.dropout_rate = dropout_rate
+		self.num_layers = num_layers
+		
+		# --------------------
+		# Transformer Module
+		# --------------------
+		# Define all modules related to a Geometric Transformer module
+		self.apply_layer_norm = 'layer' in self.norm_to_apply.lower()
+		
+		self.num_hidden_channels, self.num_output_feats = num_hidden_channels, num_hidden_channels
+		if self.apply_layer_norm:
+			self.layer_norm1_node_feats = nn.LayerNorm(self.num_output_feats)
+			self.layer_norm1_edge_feats = nn.LayerNorm(self.num_output_feats)
+		else:  # Otherwise, default to using batch normalization
+			self.batch_norm1_node_feats = nn.BatchNorm1d(self.num_output_feats)
+			self.batch_norm1_edge_feats = nn.BatchNorm1d(self.num_output_feats)
+		
+		self.mha_module = MultiHeadAttentionLayer(
+			self.num_hidden_channels,
+			self.num_output_feats // self.num_attention_heads,
+			self.num_attention_heads,
+			self.num_hidden_channels != self.num_output_feats,  # Only use bias if a Linear() has to change sizes
+			update_edge_feats=True
+		)
+		
+		self.O_node_feats = nn.Linear(self.num_output_feats, self.num_output_feats)
+		self.O_edge_feats = nn.Linear(self.num_output_feats, self.num_output_feats)
+		
+		# MLP for node features
+		dropout = nn.Dropout(p=self.dropout_rate) if self.dropout_rate > 0.0 else nn.Identity()
+		self.node_feats_MLP = nn.ModuleList([
+			nn.Linear(self.num_output_feats, self.num_output_feats * 2, bias=False),
+			self.activ_fn,
+			dropout,
+			nn.Linear(self.num_output_feats * 2, self.num_output_feats, bias=False)
+		])
+		
+		if self.apply_layer_norm:
+			self.layer_norm2_node_feats = nn.LayerNorm(self.num_output_feats)
+			self.layer_norm2_edge_feats = nn.LayerNorm(self.num_output_feats)
+		else:  # Otherwise, default to using batch normalization
+			self.batch_norm2_node_feats = nn.BatchNorm1d(self.num_output_feats)
+			self.batch_norm2_edge_feats = nn.BatchNorm1d(self.num_output_feats)
+		
+		# MLP for edge features
+		self.edge_feats_MLP = nn.ModuleList([
+			nn.Linear(self.num_output_feats, self.num_output_feats * 2, bias=False),
+			self.activ_fn,
+			dropout,
+			nn.Linear(self.num_output_feats * 2, self.num_output_feats, bias=False)
+		])
+		
+		self.reset_parameters()
+	
+	def reset_parameters(self):
+		"""Reinitialize learnable parameters."""
+		scale = 2.0
+		glorot_orthogonal(self.O_node_feats.weight, scale=scale)
+		self.O_node_feats.bias.data.fill_(0)
+		glorot_orthogonal(self.O_edge_feats.weight, scale=scale)
+		self.O_edge_feats.bias.data.fill_(0)
+		
+		for layer in self.node_feats_MLP:
+			if hasattr(layer, 'weight'):  # Skip initialization for activation functions
+				glorot_orthogonal(layer.weight, scale=scale)
+		
+		for layer in self.edge_feats_MLP:
+			if hasattr(layer, 'weight'):
+				glorot_orthogonal(layer.weight, scale=scale)
+	
+	def run_gt_layer(self, data, node_feats, edge_feats):
+		"""Perform a forward pass of geometric attention using a multi-head attention (MHA) module."""
+		node_feats_in1 = node_feats  # Cache node representations for first residual connection
+		edge_feats_in1 = edge_feats  # Cache edge representations for first residual connection
+			
+		# Apply first round of normalization before applying geometric attention, for performance enhancement
+		if self.apply_layer_norm:
+			node_feats = self.layer_norm1_node_feats(node_feats)
+			edge_feats = self.layer_norm1_edge_feats(edge_feats)
+		else:  # Otherwise, default to using batch normalization
+			node_feats = self.batch_norm1_node_feats(node_feats)
+			edge_feats = self.batch_norm1_edge_feats(edge_feats)
+		
+		# Get multi-head attention output using provided node and edge representations
+		node_attn_out, edge_attn_out = self.mha_module(node_feats, edge_feats, data.edge_index)
+		
+		node_feats = node_attn_out.view(-1, self.num_output_feats)
+		edge_feats = edge_attn_out.view(-1, self.num_output_feats)
+		
+		node_feats = F.dropout(node_feats, self.dropout_rate, training=self.training)
+		edge_feats = F.dropout(edge_feats, self.dropout_rate, training=self.training)
+		
+		node_feats = self.O_node_feats(node_feats)
+		edge_feats = self.O_edge_feats(edge_feats)
+		
+		# Make first residual connection
+		if self.residual:
+			node_feats = node_feats_in1 + node_feats  # Make first node residual connection
+			edge_feats = edge_feats_in1 + edge_feats  # Make first edge residual connection
+		
+		node_feats_in2 = node_feats  # Cache node representations for second residual connection
+		edge_feats_in2 = edge_feats  # Cache edge representations for second residual connection
+		
+		# Apply second round of normalization after first residual connection has been made
+		if self.apply_layer_norm:
+			node_feats = self.layer_norm2_node_feats(node_feats)
+			edge_feats = self.layer_norm2_edge_feats(edge_feats)
+		else:  # Otherwise, default to using batch normalization
+			node_feats = self.batch_norm2_node_feats(node_feats)
+			edge_feats = self.batch_norm2_edge_feats(edge_feats)
+		
+		# Apply MLPs for node and edge features
+		for layer in self.node_feats_MLP:
+			node_feats = layer(node_feats)
+		for layer in self.edge_feats_MLP:
+			edge_feats = layer(edge_feats)
+		
+		# Make second residual connection
+		if self.residual:
+			node_feats = node_feats_in2 + node_feats  # Make second node residual connection
+			edge_feats = edge_feats_in2 + edge_feats  # Make second edge residual connection
+		
+		# Return edge representations along with node representations (for tasks other than interface prediction)
+		return node_feats, edge_feats
+	
+	def forward(self, data, node_feats, edge_feats):
+		"""Perform a forward pass of a Geometric Transformer to get intermediate node and edge representations."""
+		node_feats, edge_feats = self.run_gt_layer(data, node_feats, edge_feats)
+		return node_feats, edge_feats
+
+
+class FinalGraphTransformerModule(nn.Module):
+	"""A (final layer) Graph Transformer module that combines node and edge representations using self-attention."""	
+	def __init__(self,
+				num_hidden_channels,
+				activ_fn=nn.SiLU(),
+				residual=True,
+				num_attention_heads=4,
+				norm_to_apply='batch',
+				dropout_rate=0.1,
+				num_layers=4):
+		super(FinalGraphTransformerModule, self).__init__()
+		
+		# Record parameters given
+		self.activ_fn = activ_fn
+		self.residual = residual
+		self.num_attention_heads = num_attention_heads
+		self.norm_to_apply = norm_to_apply
+		self.dropout_rate = dropout_rate
+		self.num_layers = num_layers
+		
+		# --------------------
+		# Transformer Module
+		# --------------------
+		# Define all modules related to a Geometric Transformer module
+		self.apply_layer_norm = 'layer' in self.norm_to_apply.lower()
+		
+		self.num_hidden_channels, self.num_output_feats = num_hidden_channels, num_hidden_channels
+		if self.apply_layer_norm:
+			self.layer_norm1_node_feats = nn.LayerNorm(self.num_output_feats)
+			self.layer_norm1_edge_feats = nn.LayerNorm(self.num_output_feats)
+		else:  # Otherwise, default to using batch normalization
+			self.batch_norm1_node_feats = nn.BatchNorm1d(self.num_output_feats)
+			self.batch_norm1_edge_feats = nn.BatchNorm1d(self.num_output_feats)
+		
+		self.mha_module = MultiHeadAttentionLayer(
+					self.num_hidden_channels,
+					self.num_output_feats // self.num_attention_heads,
+					self.num_attention_heads,
+					self.num_hidden_channels != self.num_output_feats,  # Only use bias if a Linear() has to change sizes
+					update_edge_feats=False)
+		
+		self.O_node_feats = nn.Linear(self.num_output_feats, self.num_output_feats)
+		
+		# MLP for node features
+		dropout = nn.Dropout(p=self.dropout_rate) if self.dropout_rate > 0.0 else nn.Identity()
+		self.node_feats_MLP = nn.ModuleList([
+					nn.Linear(self.num_output_feats, self.num_output_feats * 2, bias=False),
+					self.activ_fn,
+					dropout,
+					nn.Linear(self.num_output_feats * 2, self.num_output_feats, bias=False)
+					])
+		
+		if self.apply_layer_norm:
+			self.layer_norm2_node_feats = nn.LayerNorm(self.num_output_feats)
+		else:  # Otherwise, default to using batch normalization
+			self.batch_norm2_node_feats = nn.BatchNorm1d(self.num_output_feats)
+		
+		self.reset_parameters()
+	
+	def reset_parameters(self):
+		"""Reinitialize learnable parameters."""
+		scale = 2.0
+		glorot_orthogonal(self.O_node_feats.weight, scale=scale)
+		self.O_node_feats.bias.data.fill_(0)
+		
+		for layer in self.node_feats_MLP:
+			if hasattr(layer, 'weight'):  # Skip initialization for activation functions
+				glorot_orthogonal(layer.weight, scale=scale)
+		
+		#glorot_orthogonal(self.conformation_module.weight, scale=scale)
+	
+	def run_gt_layer(self, data, node_feats, edge_feats):
+		"""Perform a forward pass of geometric attention using a multi-head attention (MHA) module."""
+		node_feats_in1 = node_feats  # Cache node representations for first residual connection
+		#edge_feats = self.conformation_module(edge_feats)
+		
+		# Apply first round of normalization before applying geometric attention, for performance enhancement
+		if self.apply_layer_norm:
+			node_feats = self.layer_norm1_node_feats(node_feats)
+			edge_feats = self.layer_norm1_edge_feats(edge_feats)
+		else:  # Otherwise, default to using batch normalization
+			node_feats = self.batch_norm1_node_feats(node_feats)
+			edge_feats = self.batch_norm1_edge_feats(edge_feats)
+		
+		# Get multi-head attention output using provided node and edge representations
+		node_attn_out, _ = self.mha_module(node_feats, edge_feats, data.edge_index)
+		node_feats = node_attn_out.view(-1, self.num_output_feats)		
+		node_feats = F.dropout(node_feats, self.dropout_rate, training=self.training)
+		node_feats = self.O_node_feats(node_feats)
+		
+		# Make first residual connection
+		if self.residual:
+			node_feats = node_feats_in1 + node_feats  # Make first node residual connection
+		
+		node_feats_in2 = node_feats  # Cache node representations for second residual connection
+		
+		# Apply second round of normalization after first residual connection has been made
+		if self.apply_layer_norm:
+			node_feats = self.layer_norm2_node_feats(node_feats)
+		else:  # Otherwise, default to using batch normalization
+			node_feats = self.batch_norm2_node_feats(node_feats)
+		
+		# Apply MLP for node features
+		for layer in self.node_feats_MLP:
+			node_feats = layer(node_feats)
+		
+		# Make second residual connection
+		if self.residual:
+			node_feats = node_feats_in2 + node_feats  # Make second node residual connection
+		
+		# Return node representations
+		return node_feats
+	
+	def forward(self, data, node_feats, edge_feats):
+		"""Perform a forward pass of a Geometric Transformer to get final node representations."""
+		node_feats = self.run_gt_layer(data, node_feats, edge_feats)
+		return node_feats
+
+
+class GraphTransformer(nn.Module):
+	"""A graph transformer
+	"""
+	def __init__(
+			self,
+			in_channels, 
+			edge_features=10,
+			num_hidden_channels=128,
+            activ_fn=nn.SiLU(),
+			transformer_residual=True,
+			num_attention_heads=4,
+			norm_to_apply='batch',
+			dropout_rate=0.1,
+			num_layers=4,
+			**kwargs
+			):
+		super(GraphTransformer, self).__init__()
+		
+		# Initialize model parameters
+		self.activ_fn = activ_fn
+		self.transformer_residual = transformer_residual
+		self.num_attention_heads = num_attention_heads
+		self.norm_to_apply = norm_to_apply
+		self.dropout_rate = dropout_rate
+		self.num_layers = num_layers
+		
+		# --------------------
+		# Initializer Modules
+		# --------------------
+		# Define all modules related to edge and node initialization
+		self.node_encoder = nn.Linear(in_channels, num_hidden_channels)
+		self.edge_encoder = nn.Linear(edge_features, num_hidden_channels) 
+        # --------------------
+		# Transformer Module
+		# --------------------
+		# Define all modules related to a variable number of Geometric Transformer modules
+		num_intermediate_layers = max(0, num_layers - 1)
+		gt_block_modules = [GraphTransformerModule(
+										num_hidden_channels=num_hidden_channels,
+										activ_fn=activ_fn,
+										residual=transformer_residual,
+										num_attention_heads=num_attention_heads,
+										norm_to_apply=norm_to_apply,
+										dropout_rate=dropout_rate,
+										num_layers=num_layers) for _ in range(num_intermediate_layers)]
+		if num_layers > 0:
+			gt_block_modules.extend([
+							FinalGraphTransformerModule(
+										num_hidden_channels=num_hidden_channels,
+										activ_fn=activ_fn,
+										residual=transformer_residual,
+										num_attention_heads=num_attention_heads,
+										norm_to_apply=norm_to_apply,
+										dropout_rate=dropout_rate,
+										num_layers=num_layers)])
+		self.gt_block = nn.ModuleList(gt_block_modules)
+	
+	def forward(self, data):		
+		node_feats = self.node_encoder(data.x)
+		edge_feats = self.edge_encoder(data.edge_attr)
+			
+		# Apply a given number of intermediate geometric attention layers to the node and edge features given
+		for gt_layer in self.gt_block[:-1]:
+			node_feats, edge_feats = gt_layer(data, node_feats, edge_feats)
+		
+		# Apply final layer to update node representations by merging current node and edge representations
+		node_feats = self.gt_block[-1](data, node_feats, edge_feats)
+		data.x = node_feats
+		data.edge_attr = edge_feats
+		#return node_feats
+		return data
+
+#==============================================
+from .layer.gatedgcn_layer import GatedGCNLayer
+
+
+class GatedGCN(nn.Module):
+	"""A graph transformer
+	"""
+	def __init__(
+			self,
+			in_channels, 
+			edge_features=10,
+			num_hidden_channels=128,
+			dropout_rate=0.1,
+			num_layers=4,
+			residual=True,
+			equivstable_pe=False,
+			**kwargs
+			):
+		super(GatedGCN, self).__init__()
+		
+		# Initialize model parameters
+		self.residual = residual
+		self.dropout_rate = dropout_rate
+		self.num_layers = num_layers
+		
+		self.node_encoder = nn.Linear(in_channels, num_hidden_channels)
+		self.edge_encoder = nn.Linear(edge_features, num_hidden_channels) 
+		
+		gt_block_modules = [GatedGCNLayer(
+							num_hidden_channels,
+							num_hidden_channels,
+							dropout_rate, 
+							residual,
+							equivstable_pe=equivstable_pe) for _ in range(num_layers)]
+		
+		self.gt_block = nn.ModuleList(gt_block_modules)
+	
+	def forward(self, data):		
+		data.x = self.node_encoder(data.x)
+		data.edge_attr = self.edge_encoder(data.edge_attr)
+			
+		# Apply a given number of intermediate geometric attention layers to the node and edge features given
+		for gt_layer in self.gt_block:
+			data = gt_layer(data)
+		
+		# Apply final layer to update node representations by merging current node and edge representations
+		#data.x = node_feats
+		#data.edge_attr = edge_feats
+		#return node_feats
+		return data
+
+#==============================================
+class GenScore(nn.Module):
+	def __init__(self, ligand_model, target_model, in_channels, hidden_dim, n_gaussians, dropout_rate=0.15,
+					dist_threhold=1000):
+		super(GenScore, self).__init__()
+		
+		self.ligand_model = ligand_model
+		self.target_model = target_model
+		self.MLP = nn.Sequential(nn.Linear(in_channels*2, hidden_dim), 
+								nn.BatchNorm1d(hidden_dim), 
+								nn.ELU(), 
+								nn.Dropout(p=dropout_rate)) 
+		
+		self.z_pi = nn.Linear(hidden_dim, n_gaussians)
+		self.z_sigma = nn.Linear(hidden_dim, n_gaussians)
+		self.z_mu = nn.Linear(hidden_dim, n_gaussians)
+		self.atom_types = nn.Linear(in_channels, 17)
+		self.bond_types = nn.Linear(in_channels*2, 4)
+		
+		#self.device = 'cuda' if th.cuda.is_available() else 'cpu'
+		self.dist_threhold = dist_threhold	
+	
+	def forward(self, data_ligand, data_target):
+		
+		h_l = self.ligand_model(data_ligand)
+		h_t = self.target_model(data_target)
+		
+		h_l_x, l_mask = to_dense_batch(h_l.x, h_l.batch, fill_value=0)
+		h_t_x, t_mask = to_dense_batch(h_t.x, h_t.batch, fill_value=0)
+		h_l_pos, _ = to_dense_batch(h_l.pos, h_l.batch, fill_value=0)
+		h_t_pos, _ = to_dense_batch(h_t.pos, h_t.batch, fill_value=0)
+		
+		#assert h_l_x.size(0) == h_t_x.size(0), 'Encountered unequal batch-sizes'
+		(B, N_l, C_out), N_t = h_l_x.size(), h_t_x.size(1)
+		self.B = B
+		self.N_l = N_l
+		self.N_t = N_t
+				
+		# Combine and mask
+		h_l_x = h_l_x.unsqueeze(-2)
+		h_l_x = h_l_x.repeat(1, 1, N_t, 1) # [B, N_l, N_t, C_out]
+		
+		h_t_x = h_t_x.unsqueeze(-3)
+		h_t_x = h_t_x.repeat(1, N_l, 1, 1) # [B, N_l, N_t, C_out]
+		
+		C = th.cat((h_l_x, h_t_x), -1)
+		self.C_mask = C_mask = l_mask.view(B, N_l, 1) & t_mask.view(B, 1, N_t)
+		self.C = C = C[C_mask]
+		C = self.MLP(C)
+		
+		# Get batch indexes for ligand-target combined features
+		C_batch = th.tensor(range(B)).unsqueeze(-1).unsqueeze(-1)
+                # LL 2026 UPDATE
+		C_mask = C_mask.to(C_batch.device)
+		C_batch = C_batch.repeat(1, N_l, N_t)[C_mask]
+		
+ 		# Outputs
+		pi = F.softmax(self.z_pi(C), -1)
+		sigma = F.elu(self.z_sigma(C))+1.1
+		mu = F.elu(self.z_mu(C))+1
+		atom_types = self.atom_types(h_l.x)
+		bond_types = self.bond_types(th.cat([h_l.x[h_l.edge_index[0]], h_l.x[h_l.edge_index[1]]], axis=1))      
+		
+		
+		dist = self.compute_euclidean_distances_matrix(h_l_pos, h_t_pos.view(B,-1,3))[C_mask]
+		return pi, sigma, mu, dist.unsqueeze(1).detach(), atom_types, bond_types, C_batch
+	
+	def compute_euclidean_distances_matrix(self, X, Y):
+		# Based on: https://medium.com/@souravdey/l2-distance-matrix-vectorization-trick-26aa3247ac6c
+		# (X-Y)^2 = X^2 + Y^2 -2XY
+		X = X.double()
+		Y = Y.double()
+		
+		dists = -2 * th.bmm(X, Y.permute(0, 2, 1)) + th.sum(Y**2,    axis=-1).unsqueeze(1) + th.sum(X**2, axis=-1).unsqueeze(-1)	
+		return th.nan_to_num((dists**0.5).view(self.B, self.N_l,-1,24),10000).min(axis=-1)[0]
+		
+
+
+
+
+
+
+

models/train.py

@@ -0,0 +1,217 @@
+import argparse
+
+import numpy as np
+import torch as th
+import torch.multiprocessing
+from torch_geometric.loader import DataLoader
+
+from onescience.datapipes.genscore.data import PDBbindDataset
+from models.model.model import GatedGCN, GenScore, GraphTransformer
+from onescience.metrics.genscore.utils import (
+    EarlyStopping,
+    run_a_train_epoch,
+    run_an_eval_epoch,
+    set_random_seed,
+)
+
+torch.multiprocessing.set_sharing_strategy("file_system")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Train GenScore.")
+    parser.add_argument("--num_epochs", type=int, default=5000)
+    parser.add_argument("--batch_size", type=int, default=64)
+    parser.add_argument("--aux_weight", type=float, default=0.001)
+    parser.add_argument("--affi_weight", type=float, default=-0.5)
+    parser.add_argument("--patience", type=int, default=70)
+    parser.add_argument("--num_workers", type=int, default=8)
+    parser.add_argument("--model_path", type=str, default="genscore.pth")
+    parser.add_argument("--encoder", type=str, choices=["gt", "gatedgcn"], default="gt")
+    parser.add_argument("--mode", type=str, choices=["lower", "higher"], default="lower")
+    parser.add_argument("--finetune", action="store_true", default=False)
+    parser.add_argument("--original_model_path", type=str, default=None)
+    parser.add_argument("--lr", type=int, default=3)
+    parser.add_argument("--weight_decay", type=int, default=5)
+    parser.add_argument("--data_dir", type=str, required=True)
+    parser.add_argument("--data_prefix", type=str, default="v2020_train")
+    parser.add_argument("--valnum", type=int, default=1500)
+    parser.add_argument("--seeds", type=int, default=126)
+    parser.add_argument("--hidden_dim0", type=int, default=128)
+    parser.add_argument("--hidden_dim", type=int, default=128)
+    parser.add_argument("--n_gaussians", type=int, default=10)
+    parser.add_argument("--dropout_rate", type=float, default=0.15)
+    parser.add_argument("--dist_threhold", type=float, default=7.0)
+    parser.add_argument("--dist_threhold2", type=float, default=5.0)
+    return parser.parse_args()
+
+
+def _build_encoder(args):
+    if args.encoder == "gt":
+        ligmodel = GraphTransformer(
+            in_channels=41,
+            edge_features=10,
+            num_hidden_channels=args.hidden_dim0,
+            activ_fn=th.nn.SiLU(),
+            transformer_residual=True,
+            num_attention_heads=4,
+            norm_to_apply="batch",
+            dropout_rate=0.15,
+            num_layers=6,
+        )
+        protmodel = GraphTransformer(
+            in_channels=41,
+            edge_features=5,
+            num_hidden_channels=args.hidden_dim0,
+            activ_fn=th.nn.SiLU(),
+            transformer_residual=True,
+            num_attention_heads=4,
+            norm_to_apply="batch",
+            dropout_rate=0.15,
+            num_layers=6,
+        )
+    else:
+        ligmodel = GatedGCN(
+            in_channels=41,
+            edge_features=10,
+            num_hidden_channels=args.hidden_dim0,
+            residual=True,
+            dropout_rate=0.15,
+            equivstable_pe=False,
+            num_layers=6,
+        )
+        protmodel = GatedGCN(
+            in_channels=41,
+            edge_features=5,
+            num_hidden_channels=args.hidden_dim0,
+            residual=True,
+            dropout_rate=0.15,
+            equivstable_pe=False,
+            num_layers=6,
+        )
+    return ligmodel, protmodel
+
+
+def main():
+    args = parse_args()
+    args.device = "cuda" if th.cuda.is_available() else "cpu"
+
+    data = PDBbindDataset(
+        ids=f"{args.data_dir}/{args.data_prefix}_ids.npy",
+        ligs=f"{args.data_dir}/{args.data_prefix}_lig.pt",
+        prots=f"{args.data_dir}/{args.data_prefix}_prot.pt",
+    )
+    train_inds, val_inds = data.train_and_test_split(valnum=args.valnum, seed=args.seeds)
+    train_data = PDBbindDataset(
+        ids=data.pdbids[train_inds],
+        ligs=data.gls[train_inds],
+        prots=data.gps[train_inds],
+        labels=data.labels[train_inds],
+    )
+    val_data = PDBbindDataset(
+        ids=data.pdbids[val_inds],
+        ligs=data.gls[val_inds],
+        prots=data.gps[val_inds],
+        labels=data.labels[val_inds],
+    )
+
+    ligmodel, protmodel = _build_encoder(args)
+    model = GenScore(
+        ligmodel,
+        protmodel,
+        in_channels=args.hidden_dim0,
+        hidden_dim=args.hidden_dim,
+        n_gaussians=args.n_gaussians,
+        dropout_rate=args.dropout_rate,
+        dist_threhold=args.dist_threhold,
+    ).to(args.device)
+
+    if args.finetune:
+        if args.original_model_path is None:
+            raise ValueError('--original_model_path is required when --finetune is used.')
+        checkpoint = th.load(args.original_model_path, map_location=th.device(args.device))
+        model.load_state_dict(checkpoint["model_state_dict"])
+
+    optimizer = th.optim.Adam(
+        model.parameters(),
+        lr=10**-args.lr,
+        weight_decay=10**-args.weight_decay,
+    )
+    train_loader = DataLoader(
+        dataset=train_data,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+    )
+    val_loader = DataLoader(
+        dataset=val_data,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=args.num_workers,
+    )
+    stopper = EarlyStopping(patience=args.patience, mode=args.mode, filename=args.model_path)
+
+    set_random_seed(args.seeds)
+    for epoch in range(args.num_epochs):
+        total_loss_train, mdn_loss_train, affi_loss_train, atom_loss_train, bond_loss_train = run_a_train_epoch(
+            epoch,
+            model,
+            train_loader,
+            optimizer,
+            affi_weight=args.affi_weight,
+            aux_weight=args.aux_weight,
+            dist_threhold=args.dist_threhold2,
+            device=args.device,
+        )
+        if np.isinf(mdn_loss_train) or np.isnan(mdn_loss_train):
+            print("Inf ERROR")
+            break
+
+        total_loss_val, mdn_loss_val, affi_loss_val, atom_loss_val, bond_loss_val = run_an_eval_epoch(
+            model,
+            val_loader,
+            dist_threhold=args.dist_threhold2,
+            affi_weight=args.affi_weight,
+            aux_weight=args.aux_weight,
+            device=args.device,
+        )
+        early_stop = stopper.step(total_loss_val, model)
+        print(
+            "epoch {:d}/{:d}, total_loss_val {:.4f}, mdn_loss_val {:.4f}, "
+            "affi_loss_val {:.4f}, atom_loss_val {:.4f}, bond_loss_val {:.4f}, "
+            "best validation {:.4f}".format(
+                epoch + 1,
+                args.num_epochs,
+                total_loss_val,
+                mdn_loss_val,
+                affi_loss_val,
+                atom_loss_val,
+                bond_loss_val,
+                stopper.best_score,
+            )
+        )
+        if early_stop:
+            break
+
+    stopper.load_checkpoint(model)
+    train_metrics = run_an_eval_epoch(
+        model,
+        train_loader,
+        dist_threhold=args.dist_threhold2,
+        affi_weight=args.affi_weight,
+        aux_weight=args.aux_weight,
+        device=args.device,
+    )
+    val_metrics = run_an_eval_epoch(
+        model,
+        val_loader,
+        dist_threhold=args.dist_threhold2,
+        affi_weight=args.affi_weight,
+        aux_weight=args.aux_weight,
+        device=args.device,
+    )
+    print("train metrics:", train_metrics)
+    print("validation metrics:", val_metrics)
+
+
+if __name__ == "__main__":
+    main()

scripts/genscore.py

@@ -0,0 +1,11 @@
+import sys
+from pathlib import Path
+
+GENSCORE_DIR = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(GENSCORE_DIR))
+
+from models.inference import main
+
+
+if __name__ == "__main__":
+    main()

scripts/preprocess_pdbbind.py

@@ -0,0 +1,117 @@
+import argparse
+import os
+
+import numpy as np
+import pandas as pd
+import torch as th
+from joblib import Parallel, delayed
+
+from onescience.datapipes.genscore.feats.mol2graph_rdmda_res import mol_to_graph2
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Preprocess PDBbind complexes for GenScore training.")
+    parser.add_argument(
+        "-d",
+        "--dir",
+        default=".",
+        help="Directory containing PDBbind-style protein-ligand complex folders.",
+    )
+    parser.add_argument(
+        "-c",
+        "--cutoff",
+        default=10.0,
+        type=float,
+        help="Pocket cutoff used in pocket file names and graph construction.",
+    )
+    parser.add_argument(
+        "-o",
+        "--outprefix",
+        default="out",
+        help="Output prefix for generated _ids.npy, _prot.pt, and _lig.pt files.",
+    )
+    parser.add_argument(
+        "-r",
+        "--ref",
+        default="pdbbind_2020_general.csv",
+        help="CSV containing labels indexed by PDB id with a labels column.",
+    )
+    parser.add_argument(
+        "-usH",
+        "--useH",
+        default=False,
+        action="store_true",
+        help="Use explicit hydrogen atoms.",
+    )
+    parser.add_argument(
+        "-uschi",
+        "--use_chirality",
+        default=False,
+        action="store_true",
+        help="Use chirality features.",
+    )
+    parser.add_argument(
+        "-p",
+        "--parallel",
+        default=False,
+        action="store_true",
+        help="Build graphs in parallel.",
+    )
+    return parser.parse_args()
+
+
+def _label_query(pdbid, labels):
+    return labels.loc[pdbid, "labels"]
+
+
+def _pdbbind_paths(root_dir, pdbid, cutoff):
+    complex_dir = os.path.join(root_dir, pdbid, f"{pdbid}_prot")
+    prot_path = os.path.join(complex_dir, f"{pdbid}_p_pocket_{cutoff}.pdb")
+    lig_path = os.path.join(complex_dir, f"{pdbid}_l.sdf")
+    return prot_path, lig_path
+
+
+def _pdbbind_handle(pdbid, args, labels):
+    prot_path, lig_path = _pdbbind_paths(args.dir, pdbid, args.cutoff)
+    try:
+        graph_prot, graph_lig = mol_to_graph2(
+            prot_path,
+            lig_path,
+            cutoff=args.cutoff,
+            explicit_H=args.useH,
+            use_chirality=args.use_chirality,
+        )
+    except Exception as exc:
+        print(f"{pdbid} failed to generate graph: {exc}")
+        return None
+    return pdbid, graph_prot, graph_lig, _label_query(pdbid, labels)
+
+
+def main():
+    args = parse_args()
+    labels = pd.read_csv(args.ref, index_col=0, header=0)
+    pdbids = [
+        name
+        for name in os.listdir(args.dir)
+        if os.path.isdir(os.path.join(args.dir, name))
+    ]
+
+    if args.parallel:
+        results = Parallel(n_jobs=-1)(
+            delayed(_pdbbind_handle)(pdbid, args, labels) for pdbid in pdbids
+        )
+    else:
+        results = [_pdbbind_handle(pdbid, args, labels) for pdbid in pdbids]
+
+    results = [item for item in results if item is not None]
+    if not results:
+        raise RuntimeError("No valid PDBbind complexes were converted.")
+
+    ids, graphs_p, graphs_l, label_values = list(zip(*results))
+    np.save(f"{args.outprefix}_ids", (ids, label_values))
+    th.save(graphs_p, f"{args.outprefix}_prot.pt")
+    th.save(graphs_l, f"{args.outprefix}_lig.pt")
+
+
+if __name__ == "__main__":
+    main()

scripts/run_genscore.sh

@@ -0,0 +1,88 @@
+#!/usr/bin/env bash
+set -euo pipefail
+source ${ROCM_PATH}/cuda/env.sh
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"
+export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT_PREFIX="${SCRIPT_DIR}/out"
+ROOT_DIR="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+
+cd "${ROOT_DIR}"
+source ${ROOT_DIR}/env.sh
+
+export PYTHONPATH="${ROOT_DIR}/src:${PYTHONPATH:-}"
+
+MODEL_DIR="${GENSCORE_MODEL_DIR:-${ONESCIENCE_DATASETS_DIR}/GenScore/trained_models}"
+DATA_DIR="${GENSCORE_DATA_DIR:-${ONESCIENCE_DATASETS_DIR}/GenScore/genscore_data/inferdata}"
+BATCH_SIZE="${GENSCORE_BATCH_SIZE:-8}"
+NUM_WORKERS="${GENSCORE_NUM_WORKERS:-0}"
+
+for input_file in 1qkt_decoys.sdf 1qkt_l.sdf 1qkt_p.pdb 1qkt_p_pocket_10.0.pdb; do
+  if [[ ! -f "${DATA_DIR}/${input_file}" ]]; then
+    echo "Missing required inference file: ${DATA_DIR}/${input_file}" >&2
+    exit 1
+  fi
+done
+
+run_step() {
+  local name="$1"
+  local output="$2"
+  shift 2
+
+  echo "[run] ${name}"
+  "$@"
+  echo "[done] ${name}: ${output}"
+}
+
+# input is protein (needs to be converted to pocket)
+run_step "GT scoring with generated pocket" "${OUTPUT_PREFIX}_gt.csv" \
+  python "${SCRIPT_DIR}/genscore.py" \
+  -p "${DATA_DIR}/1qkt_p.pdb" \
+  -l "${DATA_DIR}/1qkt_decoys.sdf" \
+  -rl "${DATA_DIR}/1qkt_l.sdf" \
+  -gen_pocket \
+  -c 10.0 \
+  -e gt \
+  -m "${MODEL_DIR}/GT_0.0_1.pth" \
+  -o "${OUTPUT_PREFIX}" \
+  --batch_size "${BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}"
+
+# input is pocket
+run_step "GatedGCN scoring with prepared pocket" "${OUTPUT_PREFIX}_out_gatedgcn.csv" \
+  python "${SCRIPT_DIR}/genscore.py" \
+  -p "${DATA_DIR}/1qkt_p_pocket_10.0.pdb" \
+  -l "${DATA_DIR}/1qkt_decoys.sdf" \
+  -e gatedgcn \
+  -m "${MODEL_DIR}/GatedGCN_0.5_1.pth" \
+  -o "${OUTPUT_PREFIX}" \
+  --batch_size "${BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}"
+
+# calculate the atom contributions of the score
+run_step "GatedGCN atom contribution scoring" "${OUTPUT_PREFIX}_out_at.csv" \
+  python "${SCRIPT_DIR}/genscore.py" \
+  -p "${DATA_DIR}/1qkt_p_pocket_10.0.pdb" \
+  -l "${DATA_DIR}/1qkt_decoys.sdf" \
+  -e gatedgcn \
+  -ac \
+  -m "${MODEL_DIR}/GatedGCN_ft_1.0_1.pth" \
+  -o "${OUTPUT_PREFIX}" \
+  --batch_size "${BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}"
+
+# calculate the residue contributions of the score
+run_step "GatedGCN residue contribution scoring" "${OUTPUT_PREFIX}_out_res.csv" \
+  python "${SCRIPT_DIR}/genscore.py" \
+  -p "${DATA_DIR}/1qkt_p_pocket_10.0.pdb" \
+  -l "${DATA_DIR}/1qkt_decoys.sdf" \
+  -e gatedgcn \
+  -rc \
+  -m "${MODEL_DIR}/GatedGCN_ft_1.0_1.pth" \
+  -o "${OUTPUT_PREFIX}" \
+  --batch_size "${BATCH_SIZE}" \
+  --num_workers "${NUM_WORKERS}"
+
+echo "[done] All GenScore inference examples completed."

scripts/run_genscore_benchmarks.sh

@@ -0,0 +1,151 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+ROOT_DIR="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+cd "${ROOT_DIR}"
+
+if [[ -n "${ROCM_PATH:-}" && -f "${ROCM_PATH}/cuda/env.sh" ]]; then
+  source "${ROCM_PATH}/cuda/env.sh"
+fi
+
+source "${ROOT_DIR}/env.sh"
+export PYTHONPATH="${ROOT_DIR}/src:${PYTHONPATH:-}"
+
+if [[ -n "${CONDA_PREFIX:-}" ]]; then
+  export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib:${CONDA_PREFIX}/lib/python3.11/site-packages/torch/lib:${CONDA_PREFIX}/lib/python3.11/site-packages/fastpt/torch/lib:${LD_LIBRARY_PATH:-}"
+fi
+if [[ -n "${ROCM_PATH:-}" && -d "${ROCM_PATH}/opencl/lib" ]]; then
+  export LD_LIBRARY_PATH="${ROCM_PATH}/opencl/lib:${LD_LIBRARY_PATH:-}"
+fi
+
+TASK="${1:-all}"
+
+GENSCORE_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)"
+BENCHMARK_DIR="${GENSCORE_DIR}/benchmarks"
+MODEL_PATH="${MODEL_PATH:-${ONESCIENCE_DATASETS_DIR}/GenScore/trained_models/GatedGCN_ft_1.0_1.pth}"
+OUT_ROOT="${OUT_ROOT:-${GENSCORE_DIR}/benchmark_outputs}"
+
+GENSCORE_DATA_ROOT="${GENSCORE_DATA_ROOT:-${ONESCIENCE_DATASETS_DIR}/GenScore/genscore_data}"
+CASF_DIR="${CASF_DIR:-${GENSCORE_DATA_ROOT}/CASF-2016}"
+PDBBIND_DIR="${PDBBIND_DIR:-${GENSCORE_DATA_ROOT}/PDBbind_v2020}"
+PDBBIND_NATIVE_LIGAND_DIR="${PDBBIND_NATIVE_LIGAND_DIR:-${PDBBIND_DIR}/mol2}"
+PDBBIND_REFINED_SUBDIR="${PDBBIND_REFINED_SUBDIR:-refined-set}"
+PDBBIND_OTHER_PL_SUBDIR="${PDBBIND_OTHER_PL_SUBDIR:-v2020-other-PL}"
+CASF_PREPROCESSED_DIR="${CASF_PREPROCESSED_DIR:-${GENSCORE_DATA_ROOT}/rtmscore_s}"
+CASF_TEST_PREFIX="${CASF_TEST_PREFIX:-v2020_casf}"
+CASF_CORESET_FILE="${CASF_CORESET_FILE:-${CASF_PREPROCESSED_DIR}/v2020_casf_coreset.csv}"
+
+ENCODER="${ENCODER:-gatedgcn}"
+CUTOFF="${CUTOFF:-10.0}"
+BATCH_SIZE="${BATCH_SIZE:-128}"
+NUM_WORKERS="${NUM_WORKERS:-0}"
+OUTPREFIX="${OUTPREFIX:-gatedgcn1x5}"
+
+require_dir() {
+  local name="$1"
+  local path="$2"
+  if [[ ! -d "${path}" ]]; then
+    echo "Missing required directory ${name}: ${path}" >&2
+    exit 1
+  fi
+}
+
+require_file() {
+  local name="$1"
+  local path="$2"
+  if [[ ! -f "${path}" ]]; then
+    echo "Missing required file ${name}: ${path}" >&2
+    exit 1
+  fi
+}
+
+run_scoring_ranking() {
+  require_dir "CASF_PREPROCESSED_DIR" "${CASF_PREPROCESSED_DIR}"
+  require_file "CASF_CORESET_FILE" "${CASF_CORESET_FILE}"
+  require_file "MODEL_PATH" "${MODEL_PATH}"
+
+  echo "[run] CASF-2016 scoring/ranking"
+  python "${BENCHMARK_DIR}/casf2016_scoring_ranking.py" \
+    --data-dir "${CASF_PREPROCESSED_DIR}" \
+    --test-prefix "${CASF_TEST_PREFIX}" \
+    --coreset-file "${CASF_CORESET_FILE}" \
+    --model-path "${MODEL_PATH}" \
+    --encoder "${ENCODER}" \
+    --cutoff "${CUTOFF}" \
+    --batch-size "${BATCH_SIZE}" \
+    --num-workers "${NUM_WORKERS}" \
+    --outprefix "${OUTPREFIX}" \
+    --outdir "${OUT_ROOT}/power_ranking/examples"
+  echo "[done] CASF-2016 scoring/ranking"
+}
+
+run_docking() {
+  require_dir "CASF_DIR" "${CASF_DIR}"
+  require_dir "PDBBIND_DIR" "${PDBBIND_DIR}"
+  require_dir "PDBBIND_NATIVE_LIGAND_DIR" "${PDBBIND_NATIVE_LIGAND_DIR}"
+  require_dir "PDBbind refined subset" "${PDBBIND_DIR}/${PDBBIND_REFINED_SUBDIR}"
+  require_dir "PDBbind other-PL subset" "${PDBBIND_DIR}/${PDBBIND_OTHER_PL_SUBDIR}"
+  require_file "MODEL_PATH" "${MODEL_PATH}"
+
+  echo "[run] CASF-2016 docking"
+  python "${BENCHMARK_DIR}/casf2016_docking.py" \
+    --casf-dir "${CASF_DIR}" \
+    --pdbbind-dir "${PDBBIND_DIR}" \
+    --native-ligand-dir "${PDBBIND_NATIVE_LIGAND_DIR}" \
+    --refined-subdir "${PDBBIND_REFINED_SUBDIR}" \
+    --other-pl-subdir "${PDBBIND_OTHER_PL_SUBDIR}" \
+    --model-path "${MODEL_PATH}" \
+    --encoder "${ENCODER}" \
+    --cutoff "${CUTOFF}" \
+    --batch-size "${BATCH_SIZE}" \
+    --num-workers "${NUM_WORKERS}" \
+    --outprefix "${OUTPREFIX}" \
+    --outdir "${OUT_ROOT}/power_docking/examples/${OUTPREFIX}"
+  echo "[done] CASF-2016 docking"
+}
+
+run_screening() {
+  require_dir "CASF_DIR" "${CASF_DIR}"
+  require_dir "PDBBIND_DIR" "${PDBBIND_DIR}"
+  require_dir "PDBbind refined subset" "${PDBBIND_DIR}/${PDBBIND_REFINED_SUBDIR}"
+  require_dir "PDBbind other-PL subset" "${PDBBIND_DIR}/${PDBBIND_OTHER_PL_SUBDIR}"
+  require_file "MODEL_PATH" "${MODEL_PATH}"
+
+  echo "[run] CASF-2016 screening"
+  python "${BENCHMARK_DIR}/casf2016_screening.py" \
+    --casf-dir "${CASF_DIR}" \
+    --pdbbind-dir "${PDBBIND_DIR}" \
+    --refined-subdir "${PDBBIND_REFINED_SUBDIR}" \
+    --other-pl-subdir "${PDBBIND_OTHER_PL_SUBDIR}" \
+    --model-path "${MODEL_PATH}" \
+    --encoder "${ENCODER}" \
+    --cutoff "${CUTOFF}" \
+    --batch-size "${BATCH_SIZE}" \
+    --num-workers "${NUM_WORKERS}" \
+    --outprefix "${OUTPREFIX}" \
+    --outdir "${OUT_ROOT}/power_screening/examples/${OUTPREFIX}"
+  echo "[done] CASF-2016 screening"
+}
+
+case "${TASK}" in
+  scoring|scoring_ranking)
+    run_scoring_ranking
+    ;;
+  docking)
+    run_docking
+    ;;
+  screening)
+    run_screening
+    ;;
+  all)
+    run_scoring_ranking
+    run_docking
+    run_screening
+    ;;
+  *)
+    echo "Usage: $0 [scoring|docking|screening|all]" >&2
+    exit 1
+    ;;
+esac

scripts/train_genscore.py

@@ -0,0 +1,11 @@
+import sys
+from pathlib import Path
+
+GENSCORE_DIR = Path(__file__).resolve().parent.parent
+sys.path.insert(0, str(GENSCORE_DIR))
+
+from models.train import main
+
+
+if __name__ == "__main__":
+    main()

scripts/train_genscore_full.sh

@@ -0,0 +1,62 @@
+#!/usr/bin/env bash
+
+set -euo pipefail
+source ${ROCM_PATH}/cuda/env.sh
+export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+ROOT_DIR="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+echo ${ROOT_DIR}
+source ${ROOT_DIR}/env.sh
+
+cd "${ROOT_DIR}"
+
+export PYTHONPATH="${ROOT_DIR}/src:${PYTHONPATH:-}"
+
+if [[ -n "${CONDA_PREFIX:-}" ]]; then
+  export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib:${CONDA_PREFIX}/lib/python3.11/site-packages/torch/lib:${CONDA_PREFIX}/lib/python3.11/site-packages/fastpt/torch/lib:/opt/dtk/.hyhal/rocm_smi/lib:${LD_LIBRARY_PATH:-}"
+fi
+
+GENSCORE_DATA_DIR="${GENSCORE_DATA_DIR:-${ONESCIENCE_DATASETS_DIR}/GenScore/genscore_data/rtmscore_s}"
+GENSCORE_DATA_PREFIX="${GENSCORE_DATA_PREFIX:-v2020_train}"
+GENSCORE_ENCODER="${GENSCORE_ENCODER:-gatedgcn}"
+GENSCORE_MODEL_PATH="${GENSCORE_MODEL_PATH:-${SCRIPT_DIR}/genscore_${GENSCORE_ENCODER}_full_3000.pth}"
+
+GENSCORE_NUM_EPOCHS="${GENSCORE_NUM_EPOCHS:-3000}"
+GENSCORE_BATCH_SIZE="${GENSCORE_BATCH_SIZE:-64}"
+GENSCORE_NUM_WORKERS="${GENSCORE_NUM_WORKERS:-8}"
+GENSCORE_VALNUM="${GENSCORE_VALNUM:-1500}"
+GENSCORE_PATIENCE="${GENSCORE_PATIENCE:-70}"
+
+for suffix in ids.npy lig.pt prot.pt; do
+  path="${GENSCORE_DATA_DIR}/${GENSCORE_DATA_PREFIX}_${suffix}"
+  if [[ ! -f "${path}" ]]; then
+    echo "Missing required training file: ${path}" >&2
+    exit 1
+  fi
+done
+
+mkdir -p "$(dirname "${GENSCORE_MODEL_PATH}")"
+
+echo "GenScore full training"
+echo "  data_dir: ${GENSCORE_DATA_DIR}"
+echo "  data_prefix: ${GENSCORE_DATA_PREFIX}"
+echo "  model_path: ${GENSCORE_MODEL_PATH}"
+echo "  encoder: ${GENSCORE_ENCODER}"
+echo "  num_epochs: ${GENSCORE_NUM_EPOCHS}"
+echo "  batch_size: ${GENSCORE_BATCH_SIZE}"
+echo "  num_workers: ${GENSCORE_NUM_WORKERS}"
+echo "  valnum: ${GENSCORE_VALNUM}"
+echo "  patience: ${GENSCORE_PATIENCE}"
+
+python ${SCRIPT_DIR}/train_genscore.py \
+  --data_dir "${GENSCORE_DATA_DIR}" \
+  --data_prefix "${GENSCORE_DATA_PREFIX}" \
+  --model_path "${GENSCORE_MODEL_PATH}" \
+  --num_epochs "${GENSCORE_NUM_EPOCHS}" \
+  --batch_size "${GENSCORE_BATCH_SIZE}" \
+  --num_workers "${GENSCORE_NUM_WORKERS}" \
+  --valnum "${GENSCORE_VALNUM}" \
+  --patience "${GENSCORE_PATIENCE}" \
+  --encoder "${GENSCORE_ENCODER}"

scripts/train_genscore_smoke.sh

@@ -0,0 +1,31 @@
+#!/usr/bin/env bash
+set -euo pipefail
+source ${ROCM_PATH}/cuda/env.sh
+export LD_LIBRARY_PATH=${ROCM_PATH}/opencl/lib:$LD_LIBRARY_PATH
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+ROOT_DIR="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+echo ${ROOT_DIR}
+source ${ROOT_DIR}/env.sh
+cd "${ROOT_DIR}"
+
+export PYTHONPATH="${ROOT_DIR}/src:${PYTHONPATH:-}"
+
+if [[ -n "${CONDA_PREFIX:-}" ]]; then
+  export LD_LIBRARY_PATH="${CONDA_PREFIX}/lib:${CONDA_PREFIX}/lib/python3.11/site-packages/torch/lib:${CONDA_PREFIX}/lib/python3.11/site-packages/fastpt/torch/lib:/opt/dtk/.hyhal/rocm_smi/lib:${LD_LIBRARY_PATH:-}"
+fi
+
+GENSCORE_DATA_DIR="${GENSCORE_DATA_DIR:-${ONESCIENCE_DATASETS_DIR}/GenScore/genscore_data/rtmscore_s}"
+GENSCORE_DATA_PREFIX="${GENSCORE_DATA_PREFIX:-v2020_train}"
+GENSCORE_MODEL_PATH="${GENSCORE_MODEL_PATH:-${SCRIPT_DIR}/genscore_smoke_bs16.pth}"
+
+python ${SCRIPT_DIR}/train_genscore.py \
+  --data_dir "${GENSCORE_DATA_DIR}" \
+  --data_prefix "${GENSCORE_DATA_PREFIX}" \
+  --model_path "${GENSCORE_MODEL_PATH}" \
+  --num_epochs 100 \
+  --batch_size 16 \
+  --num_workers 0 \
+  --valnum 1500 \
+  --encoder "${GENSCORE_ENCODER:-gatedgcn}"