services:
  tabnet-bench:
    build:
      context: .
      args:
        INSTALL_TABFM: "true"
        TABFM_EXTRA: pytorch
    volumes:
      - ./results:/workspace/results
      # TabFM's first load() per model_type downloads weights from Hugging
      # Face (~22 min unauthenticated in testing); persist the HF cache so
      # repeated `docker compose run` calls don't re-download them.
      - hf-cache:/root/.cache/huggingface
    command: ["--dataset", "adult", "--methods", "xgboost", "--quick"]

  tabnet-bench-gpu:
    extends: tabnet-bench
    build:
      context: .
      args:
        INSTALL_TABFM: "true"
        TABFM_EXTRA: pytorch
        # Matches a driver reporting "CUDA Version: 12.8" in `nvidia-smi`.
        # Adjust if your driver supports a different max CUDA version.
        TORCH_INDEX_URL: https://download.pytorch.org/whl/cu128
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]

volumes:
  hf-cache: