From 6f0182eac23333b1c4712435a4782af32cff280e Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 30 Jan 2024 17:10:18 -0800
Subject: [PATCH 1/8] separate dependencies

---
 .github/actions/setup-venv/action.yml |  4 +--
 .github/workflows/main.yml            | 15 +++++----
 README.md                             |  6 ++--
 dev-requirements.txt                  | 10 ------
 docker/Dockerfile.gantry              |  7 ++--
 docker/Dockerfile.lumi                |  8 +++--
 pyproject.toml                        | 48 ++++++++++++++++++++++++---
 requirements.txt                      | 19 -----------
 8 files changed, 69 insertions(+), 48 deletions(-)
 delete mode 100644 dev-requirements.txt
 delete mode 100644 requirements.txt

diff --git a/.github/actions/setup-venv/action.yml b/.github/actions/setup-venv/action.yml
index 9d227cc25..3d4dc6a7e 100644
--- a/.github/actions/setup-venv/action.yml
+++ b/.github/actions/setup-venv/action.yml
@@ -30,7 +30,7 @@ runs:
       id: virtualenv-cache
       with:
         path: .venv
-        key: ${{ inputs.cache-prefix }}-${{ runner.os }}-${{ env.PYTHON_VERSION }}-${{ hashFiles('requirements.txt', 'dev-requirements.txt', 'hf_olmo/requirements.txt') }}
+        key: ${{ inputs.cache-prefix }}-${{ runner.os }}-${{ env.PYTHON_VERSION }}-${{ hashFiles('*requirements.txt', '*pyproject.toml') }}
 
     - if: steps.virtualenv-cache.outputs.cache-hit != 'true'
       shell: bash
@@ -39,7 +39,7 @@ runs:
         test -d .venv || virtualenv -p $(which python) --copies --reset-app-data .venv
         . .venv/bin/activate
         pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
-        pip install -e .[dev]
+        pip install -e .[all]
         pip install -e hf_olmo
 
     - if: steps.virtualenv-cache.outputs.cache-hit == 'true'
diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 087b9926b..6142d4ad8 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -78,10 +78,10 @@ jobs:
         uses: actions/cache@v3
         with:
           path: .mypy_cache
-          key: mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt') }}-${{ github.ref }}-${{ github.sha }}
+          key: mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt', '*pyproject.toml') }}-${{ github.ref }}-${{ github.sha }}
           restore-keys: |
-            mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt') }}-${{ github.ref }}
-            mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt') }}
+            mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt', '*pyproject.toml') }}-${{ github.ref }}
+            mypy-${{ env.CACHE_PREFIX }}-${{ runner.os }}-${{ matrix.python }}-${{ hashFiles('*requirements.txt', '*pyproject.toml') }}
 
       - name: ${{ matrix.task.name }}
         run: |
@@ -172,10 +172,11 @@ jobs:
         with:
           python-version: '3.10'
 
-      - name: Install requirements
-        run: |
-          pip install --upgrade pip setuptools wheel build
-          pip install -r dev-requirements.txt
+      - name: Setup Python environment
+        uses: ./.github/actions/setup-venv
+        with:
+          python-version: '3.10'
+          cache-prefix: ${{ env.CACHE_PREFIX }}
 
       - name: Prepare environment
         run: |
diff --git a/README.md b/README.md
index d84f5a660..bcf077e50 100644
--- a/README.md
+++ b/README.md
@@ -25,15 +25,17 @@ To install from source (recommended for training/fine-tuning) run:
 ```bash
 git clone https://github.com/allenai/OLMo.git
 cd OLMo
-pip install -e .
+pip install -e .[all]
 ```
 
 Otherwise you can install the model code by itself directly from PyPI with:
 
 ```bash
-pip install ai2-olmo
+pip install ai2-olmo[hf]
 ```
 
+The `hf` extra enables the HuggingFace integration.
+
 ## Models overview
 
 The core models in the OLMo family released so far are (all trained on the [Dolma dataset](https://huggingface.co/datasets/allenai/dolma)): 
diff --git a/dev-requirements.txt b/dev-requirements.txt
deleted file mode 100644
index 3c17627e5..000000000
--- a/dev-requirements.txt
+++ /dev/null
@@ -1,10 +0,0 @@
-ruff
-mypy>=1.0,<1.4
-black>=23.1,<24.0
-isort>=5.12,<5.13
-pytest
-pytest-sphinx
-twine>=1.11.0
-setuptools
-wheel
-build
diff --git a/docker/Dockerfile.gantry b/docker/Dockerfile.gantry
index 18fd894b0..fad65043f 100644
--- a/docker/Dockerfile.gantry
+++ b/docker/Dockerfile.gantry
@@ -8,7 +8,10 @@ FROM olmo-torch2-base
 
 WORKDIR /stage
 
-COPY requirements.txt .
-RUN pip install --no-cache-dir -r requirements.txt
+COPY pyproject.toml .
+RUN mkdir olmo && touch olmo/__init__.py && \
+    pip install --no-cache-dir .[all] && \
+    pip uninstall -y ai2-olmo && \
+    rm -rf olmo/
 
 WORKDIR /app/olmo
diff --git a/docker/Dockerfile.lumi b/docker/Dockerfile.lumi
index 74d036d00..9bcf84a62 100644
--- a/docker/Dockerfile.lumi
+++ b/docker/Dockerfile.lumi
@@ -88,8 +88,12 @@ RUN cd /opt && \
     sed -i 's/hostname -I/hostname -s/g' /usr/local/lib/python3.10/dist-packages/deepspeed/comm/comm.py
 
 # Install more dependencies
-COPY requirements.txt requirements.txt
-RUN pip install --no-cache-dir -r requirements.txt
+COPY pyproject.toml .
+RUN mkdir olmo && touch olmo/__init__.py && \
+    pip install --no-cache-dir .[all] && \
+    pip uninstall -y ai2-olmo && \
+    rm -rf olmo/
+
 RUN pip install --no-cache-dir py-spy
 RUN pip install --no-cache-dir wandb --upgrade
 
diff --git a/pyproject.toml b/pyproject.toml
index a1592d8ee..3887351a1 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,14 +4,56 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ai2-olmo"
-dynamic = ["version", "dependencies", "optional-dependencies"]
+dynamic = ["version"]
 readme = "README.md"
 description = "Open Language Model (OLMo)"
 authors = [
     { name = "Allen Institute for Artificial Intelligence", email = "olmo@allenai.org" }
 ]
 requires-python = ">=3.8"
-license = {file = "LICENSE"}
+license = { file = "LICENSE" }
+dependencies = [
+    "numpy",
+    "torch>=2.0",
+    "omegaconf",
+    "rich",
+    "boto3",
+    "google-cloud-storage",
+    "tokenizers",
+    "packaging",
+    "cached_path",
+]
+
+[project.optional-dependencies]
+dev = [
+    "ruff",
+    "mypy>=1.0,<1.4",
+    "black>=23.1,<24.0",
+    "isort>=5.12,<5.13",
+    "pytest",
+    "pytest-sphinx",
+    "twine>=1.11.0",
+    "setuptools",
+    "wheel",
+    "build",
+]
+train = [
+    "wandb",
+    "beaker-gantry",
+    "click",
+    "torchmetrics",
+    "smashed[remote]>=0.21.1"
+    "safetensors",
+    "datasets",
+    "scikit-learn",
+    "msgspec>=0.14.0",
+]
+hf = [
+    "transformers",
+]
+all = [
+    "ai2-olmo[dev,train,hf]",
+]
 
 [project.urls]
 Homepage = "https://github.com/allenai/OLMo"
@@ -25,8 +67,6 @@ olmo = ["py.typed"]
 
 [tool.setuptools.dynamic]
 version = { attr = "olmo.version.VERSION" }
-dependencies = { file = ["requirements.txt", "hf_olmo/requirements.txt"] }
-optional-dependencies = { dev = { file = ["dev-requirements.txt"] } }
 
 [tool.setuptools.packages.find]
 include = ["olmo*", "hf_olmo*"]
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 9256a2663..000000000
--- a/requirements.txt
+++ /dev/null
@@ -1,19 +0,0 @@
-numpy
-torch>=2.0
-torchmetrics
-tokenizers
-click
-rich
-cached-path
-beaker-gantry
-omegaconf
-wandb
-boto3
-google-cloud-storage
-packaging
-datasets
-scikit-learn
-smashed[remote]>=0.21.1
-msgspec>=0.14.0
-cached_path
-safetensors
\ No newline at end of file

From 9c829750325d8089156ccff92438af4d95263de4 Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 30 Jan 2024 17:12:37 -0800
Subject: [PATCH 2/8] TOML fix

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 3887351a1..05a7e8be6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -42,7 +42,7 @@ train = [
     "beaker-gantry",
     "click",
     "torchmetrics",
-    "smashed[remote]>=0.21.1"
+    "smashed[remote]>=0.21.1",
     "safetensors",
     "datasets",
     "scikit-learn",

From 983df9944e810b15b6403e15505c12b783de9a84 Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 30 Jan 2024 17:17:25 -0800
Subject: [PATCH 3/8] Fixes

---
 pyproject.toml                       | 2 +-
 scripts/beaker/beaker_interactive.sh | 2 +-
 scripts/test_entrypoint.sh           | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 05a7e8be6..ea496b9df 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,7 @@ requires-python = ">=3.8"
 license = { file = "LICENSE" }
 dependencies = [
     "numpy",
-    "torch>=2.0",
+    "torch>=2.0,<2.2",
     "omegaconf",
     "rich",
     "boto3",
diff --git a/scripts/beaker/beaker_interactive.sh b/scripts/beaker/beaker_interactive.sh
index 0a2f40593..c6456dfe8 100644
--- a/scripts/beaker/beaker_interactive.sh
+++ b/scripts/beaker/beaker_interactive.sh
@@ -49,4 +49,4 @@ gh repo clone allenai/LLM
 cd LLM
 
 # Install other dependencies.
-pip install -e '.[dev]'
+pip install -e '.[all]'
diff --git a/scripts/test_entrypoint.sh b/scripts/test_entrypoint.sh
index 4f289eba1..5dfcaea52 100644
--- a/scripts/test_entrypoint.sh
+++ b/scripts/test_entrypoint.sh
@@ -29,7 +29,7 @@ git checkout --quiet "$COMMIT_SHA"
 
 # Install dependencies.
 pip install --upgrade pip
-pip install --no-cache-dir '.[dev]'
+pip install --no-cache-dir '.[all]'
 
 # Create directory for results.
 mkdir -p /results

From faff2c674ad5db4487fa5724de2f27a7f0bd3b5f Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 30 Jan 2024 17:19:57 -0800
Subject: [PATCH 4/8] Fix install action

---
 .github/actions/setup-venv/action.yml | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/.github/actions/setup-venv/action.yml b/.github/actions/setup-venv/action.yml
index 3d4dc6a7e..c9a16c222 100644
--- a/.github/actions/setup-venv/action.yml
+++ b/.github/actions/setup-venv/action.yml
@@ -8,6 +8,10 @@ inputs:
     description: Update this to invalidate the cache
     required: true
     default: v0
+  torch-version:
+    description: The PyTorch version to install
+    required: false
+    default: '<2.2'
 runs:
   using: composite
   steps:
@@ -38,7 +42,7 @@ runs:
         # Set up virtual environment without cache hit.
         test -d .venv || virtualenv -p $(which python) --copies --reset-app-data .venv
         . .venv/bin/activate
-        pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+        pip install torch${{inputs.torch-version}} torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
         pip install -e .[all]
         pip install -e hf_olmo
 
@@ -47,7 +51,7 @@ runs:
       run: |
         # Set up virtual environment from cache hit.
         . .venv/bin/activate
-        pip install --no-deps -e .[dev]
+        pip install --no-deps -e .[all]
         pip install --no-deps -e hf_olmo
 
     - shell: bash

From 7e5868549623c40fddc5a9442f63a0c3fe70be3b Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 30 Jan 2024 17:25:29 -0800
Subject: [PATCH 5/8] Clean up

---
 .github/actions/setup-venv/action.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/actions/setup-venv/action.yml b/.github/actions/setup-venv/action.yml
index c9a16c222..578bba17a 100644
--- a/.github/actions/setup-venv/action.yml
+++ b/.github/actions/setup-venv/action.yml
@@ -42,7 +42,7 @@ runs:
         # Set up virtual environment without cache hit.
         test -d .venv || virtualenv -p $(which python) --copies --reset-app-data .venv
         . .venv/bin/activate
-        pip install torch${{inputs.torch-version}} torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cpu
+        pip install torch${{inputs.torch-version}} --extra-index-url https://download.pytorch.org/whl/cpu
         pip install -e .[all]
         pip install -e hf_olmo
 

From 1e2a63d996b92cf5877c90a84fd13a314db97b37 Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 30 Jan 2024 17:26:07 -0800
Subject: [PATCH 6/8] New cache

---
 .github/workflows/main.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 6142d4ad8..754a222f1 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -16,7 +16,7 @@ on:
 
 env:
   # Change this to invalidate existing cache.
-  CACHE_PREFIX: v1
+  CACHE_PREFIX: v2
   PYTHONPATH: ./
   TOKENIZERS_PARALLELISM: 'false'
 

From 3020bdf5626b552ff8d5089d8c5c57c3d2b44a5a Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 30 Jan 2024 17:28:09 -0800
Subject: [PATCH 7/8] Fix?

---
 .github/actions/setup-venv/action.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/actions/setup-venv/action.yml b/.github/actions/setup-venv/action.yml
index 578bba17a..3aaef6516 100644
--- a/.github/actions/setup-venv/action.yml
+++ b/.github/actions/setup-venv/action.yml
@@ -42,7 +42,7 @@ runs:
         # Set up virtual environment without cache hit.
         test -d .venv || virtualenv -p $(which python) --copies --reset-app-data .venv
         . .venv/bin/activate
-        pip install torch${{inputs.torch-version}} --extra-index-url https://download.pytorch.org/whl/cpu
+        pip install 'torch${{ inputs.torch-version }}' --extra-index-url https://download.pytorch.org/whl/cpu
         pip install -e .[all]
         pip install -e hf_olmo
 

From 7c632d5853af2989f5e31e07cc7c66dc636962ed Mon Sep 17 00:00:00 2001
From: epwalsh <petew@allenai.org>
Date: Tue, 30 Jan 2024 17:57:20 -0800
Subject: [PATCH 8/8] Make hf deps core

---
 README.md      | 4 +---
 pyproject.toml | 6 ++----
 2 files changed, 3 insertions(+), 7 deletions(-)

diff --git a/README.md b/README.md
index bcf077e50..0575491db 100644
--- a/README.md
+++ b/README.md
@@ -31,11 +31,9 @@ pip install -e .[all]
 Otherwise you can install the model code by itself directly from PyPI with:
 
 ```bash
-pip install ai2-olmo[hf]
+pip install ai2-olmo
 ```
 
-The `hf` extra enables the HuggingFace integration.
-
 ## Models overview
 
 The core models in the OLMo family released so far are (all trained on the [Dolma dataset](https://huggingface.co/datasets/allenai/dolma)): 
diff --git a/pyproject.toml b/pyproject.toml
index ea496b9df..db9af8201 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -22,6 +22,7 @@ dependencies = [
     "tokenizers",
     "packaging",
     "cached_path",
+    "transformers",
 ]
 
 [project.optional-dependencies]
@@ -48,11 +49,8 @@ train = [
     "scikit-learn",
     "msgspec>=0.14.0",
 ]
-hf = [
-    "transformers",
-]
 all = [
-    "ai2-olmo[dev,train,hf]",
+    "ai2-olmo[dev,train]",
 ]
 
 [project.urls]