requirements.txt

# --------- base --------- #
packaging~=23.2
absl-py

# --------- pytorch --------- #
torch~=2.3.1
torchvision~=0.18
lightning~=2.3.3
torchmetrics~=1.4.1

# --------- kernels --------- #
deepspeed~=0.14.4  # works in Colab
einops~=0.8.0
triton~=3.0.0
# flash-attn~=2.6.3

# --------- hydra --------- #
hydra-core
hydra-colorlog
hydra-optuna-sweeper

# --------- loggers --------- #
wandb
# neptune-client
# mlflow
# comet-ml
# aim>=3.16.2  # no lower than 3.16.2, see https://github.com/aimhubio/aim/issues/2550

# --------- others --------- #
rootutils~=1.0.7
pre-commit      # hooks for applying linters on commit
rich~=13.7.0
pytest~=8.0.0
matplotlib
tree~=0.2.4
dm-tree~=0.1.8
# sh            # for running bash commands in some tests (linux/macos only)
pandas
p-tqdm
proteinflow~=2.7.0
biopython~=1.79
omegaconf~=2.3.0
setuptools~=68.2.2
# numpy~=1.24.4
# scipy~=1.12.0
# jax~=0.4.23
# prettytable
tqdm~=4.66.1
pandas~=2.2.0
ml_collections