environment.yml

# Conda environment shared across the molecule, peptide, and language experiments.
# Create with:
#     conda env create -f environment.yml
#     conda activate a2d2
#
# NOTE: flash-attn is hardware-specific and must be built against your installed torch
# and CUDA, so it is not listed below. It is imported by the shared transformer backbone
# (model/casual_transformer.py, model/rotary.py) and is required for all experiments.
# After creating the env, install it with:
#     pip install flash-attn==2.8.3 --no-build-isolation
# Adjust pytorch-cuda below to match your CUDA toolkit / GPU.
name: a2d2
channels:
  - pytorch
  - nvidia
  - conda-forge
dependencies:
  - python=3.11
  - pip
  - pytorch
  - pytorch-cuda=12.1
  - rdkit=2023.9.6
  - jupyterlab          # for demo/quality_inference_demo.ipynb
  - pip:
      # --- core scientific / DL stack ---
      - numpy==1.26.4
      - scipy==1.17.1
      - pandas==2.1.4
      - scikit-learn==1.8.0
      - pytorch-lightning==2.6.0
      - lightning==2.6.1
      - transformers==4.55.4
      - tokenizers==0.21.4
      - safetensors==0.7.0
      - accelerate==0.33.0
      - peft==0.15.1            # LoRA adapters (language experiment)
      - datasets==2.19.2
      - huggingface-hub==0.36.2
      - einops==0.8.2
      - timm==1.0.26
      - omegaconf==2.3.0
      - wandb==0.26.1
      # --- molecule experiment ---
      - safe-mol==0.1.14
      - datamol==0.12.5
      - PyTDC==1.1.15
      # --- peptide experiment ---
      - SmilesPE==0.0.3
      - fair-esm==2.0.0
      - xgboost==3.2.0
      # --- plotting / utilities ---
      - matplotlib==3.10.6
      - seaborn==0.13.2
      - tqdm==4.67.1
      - joblib==1.5.3
      - loguru==0.7.3
      - fsspec==2024.3.1