environment.yml

# Conda environment for ModelLens.
#
# Usage:
#   conda env create -f environment.yml
#   conda activate modellens
#
# Notes:
#   * The PyTorch wheel below targets CUDA 12.1. If your driver is older, swap
#     `pytorch-cuda=12.1` for `pytorch-cuda=11.8` (or use the CPU build by
#     removing pytorch-cuda and the `pytorch` channel).
#   * `wandb` is optional; set `use_wandb: False` in the YAML config if you
#     do not want to log to Weights & Biases.
name: modellens
channels:
  - pytorch
  - nvidia
  - conda-forge
  - defaults
dependencies:
  - python=3.10
  - pip>=24.0

  # Numerics
  - numpy>=1.24,<2.0
  - pandas>=2.0
  - scipy>=1.11
  - scikit-learn>=1.3

  # Deep learning
  - pytorch>=2.1
  - pytorch-cuda=12.1
  - torchvision

  # Utilities
  - pyyaml>=6.0
  - tqdm>=4.65

  - pip:
      - wandb>=0.16
      - huggingface-hub>=0.20