{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tuning for downstream tasks\n",
    "\n",
    "This tutorial demonstrates how to fine-tune DreaMS for a downstream task involving 10 molecular properties, including logP, quantitative estimation of drug-likeness, and synthetic accessibility, among others. We'll use the MassSpecGym dataset, which has been split using Murcko histograms (as detailed in the [previous tutorial](https://dreams-docs.readthedocs.io/en/latest/tutorials/murcko_hist_split.html); can be downloaded from [Hugging Face Hub](https://huggingface.co/datasets/roman-bushuiev/GeMS/resolve/main/data/auxiliary/MassSpecGym_MurckoHist_split.hdf5)). To fine-tune DreaMS, one needs to run the `fine_tune.sh` script located in the `DreaMS/dreams/training` folder. The content of this script is shown in the following code snippet.\n",
    "\n",
    "```bash\n",
    "#!/bin/bash\n",
    "#SBATCH --job-name DreaMS_fine-tuning\n",
    "#SBATCH --account OPEN-29-57\n",
    "#SBATCH --partition qgpu\n",
    "#SBATCH --nodes 1\n",
    "#SBATCH --gpus 8\n",
    "#SBATCH --time 10:00:00\n",
    "\n",
    "# Activate conda environment\n",
    "eval \"$(conda shell.bash hook)\"\n",
    "conda activate dreams\n",
    "\n",
    "# Export project definitions\n",
    "$(python -c \"from dreams.definitions import export; export()\")\n",
    "\n",
    "# Move to running dir\n",
    "cd \"${DREAMS_DIR}\" || exit 3\n",
    "\n",
    "# Run the training script\n",
    "# Replace `python3 training/train.py` with `srun --export=ALL --preserve-env python3 training/train.py \\`\n",
    "# when executing on a SLURM cluster via `sbatch`.\n",
    "python3 training/train.py \\\n",
    " --project_name MolecularProperties \\\n",
    " --job_key \"my_run_name\" \\\n",
    " --run_name \"my_run_name\" \\\n",
    " --train_objective mol_props \\\n",
    " --train_regime fine-tuning \\\n",
    " --dataset_pth \"${DATA_DIR}/MassSpecGym_MurckoHist_split.hdf5\" \\\n",
    " --dformat A \\\n",
    " --model DreaMS \\\n",
    " --lr 3e-5 \\\n",
    " --batch_size 64 \\\n",
    " --prec_intens 1.1 \\\n",
    " --num_devices 8 \\\n",
    " --max_epochs 103 \\\n",
    " --log_every_n_steps 5 \\\n",
    " --head_depth 1 \\\n",
    " --seed 3407 \\\n",
    " --train_precision 64   \\\n",
    " --pre_trained_pth \"${PRETRAINED}/ssl_model.ckpt\" \\\n",
    " --val_check_interval 0.1 \\\n",
    " --max_peaks_n 100 \\\n",
    " --save_top_k -1\n",
    "```\n",
    "\n",
    "There are several important points to note here.\n",
    "\n",
    "- The header of the file specifies the job name, account, partition, nodes, GPUs, and time. This is used by the SLURM scheduler to allocate resources for the job when submitted to a SLURM cluster via `sbatch training/pre_train.sh`. Importantly, to run the training script on a SLURM cluster, you need to replace `python3 training/train.py` with `srun --export=ALL --preserve-env python3 training/train.py \\` when executing the script via `sbatch`. If your cluster uses a different scheduler (e.g., PBS), you can modify the header to suit your needs. Note, that the `#SBATCH` commands are ignored when the script is run locally or on a non-SLURM cluster.\n",
    "\n",
    "- The script activates the `dreams` conda environment and exports the project definitions. This ensures that the training script can access the project's configurations. Therefore, the `dreams` conda environment needs to be installed beforehand (according to [Getting started](https://dreams-docs.readthedocs.io/en/latest/) section of the documentation).\n",
    "\n",
    "- The `--dataset_pth` specifies the path to the training dataset. In this tutorial, we use the dataset, which can be downloaded from the [GeMS Hugging Face Hub repository](https://huggingface.co/datasets/roman-bushuiev/GeMS/resolve/main/data/GeMS_A/GeMS_A10.hdf5).\n",
    "\n",
    "- The `--train_objective mol_props` specifies the training objective as molecular properties prediction. This tells the training script to generate molecular properties for each molecule as trainign labels. To train DreaMS on other tasks, one needs to either specify a different available objective (e.g., `has_F` for fluorine detection) or to implement a custom objective. The implementation of custom objectives necessiates one to write a custom class inherited from the [`FineTuningHead`](https://dreams-docs.readthedocs.io/en/latest/dreams.models.heads.html).\n",
    "\n",
    "- It is recommended to sign up for a [WandB](https://wandb.ai/site) account and log in via the command line using `wandb login`. This allows you to monitor the training progress and inspect the model's performance.\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dreams",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}