From 1cc20c43f7fa9f391902be5e922ae66371859227 Mon Sep 17 00:00:00 2001 From: Wessel Bruinsma Date: Wed, 21 Aug 2024 11:39:20 +0200 Subject: [PATCH] Move to public HF repo --- .github/workflows/ci.yaml | 6 ------ README.md | 12 ++---------- docs/example.ipynb | 2 +- docs/finetuning.md | 4 ++-- docs/usage.md | 4 ++-- tests/__init__.py | 5 ----- tests/test_model.py | 9 ++++----- tests/test_rollout.py | 3 +-- 8 files changed, 12 insertions(+), 33 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index b57e820..097fb18 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -13,7 +13,6 @@ jobs: fail-fast: false matrix: version: ["3.10", "3.11"] - environment: huggingface-access name: Test with Python ${{ matrix.version }} steps: @@ -26,11 +25,6 @@ jobs: run: | python -m pip install --upgrade pip python -m pip install --upgrade --no-cache-dir -e '.[dev]' - - name: Log into HuggingFace - run: | - huggingface-cli login --token ${{ secrets.HUGGINGFACE_TOKEN }} - name: Run tests run: | pytest -v --cov=aurora --cov-report term-missing - env: - HUGGINGFACE_REPO: wbruinsma/aurora diff --git a/README.md b/README.md index 386e406..b62467e 100644 --- a/README.md +++ b/README.md @@ -71,7 +71,7 @@ from aurora import AuroraSmall, Batch, Metadata model = AuroraSmall() -model.load_checkpoint("wbruinsma/aurora", "aurora-0.25-small-pretrained.ckpt") +model.load_checkpoint("microsoft/aurora", "aurora-0.25-small-pretrained.ckpt") batch = Batch( surf_vars={k: torch.randn(1, 2, 17, 32) for k in ("2t", "10u", "10v", "msl")}, @@ -90,8 +90,7 @@ prediction = model.forward(batch) print(prediction.surf_vars["2t"]) ``` -Note that this will incur a 500 MB download -and that you may need to authenticate with `huggingface-cli login`. +Note that this will incur a 500 MB download. Please read the [documentation](https://microsoft.github.io/aurora) for more detailed instructions and for which models are available. @@ -160,13 +159,6 @@ First, install the repository in editable mode and setup `pre-commit`: make install ``` -Then configure the HuggingFace repository where the weights can be found and log into HuggingFace: - -```bash -export HUGGINGFACE_REPO=wbruinsma/aurora -huggingface-cli login -``` - To run the tests and print coverage, run ```bash diff --git a/docs/example.ipynb b/docs/example.ipynb index 3867547..1c71af2 100644 --- a/docs/example.ipynb +++ b/docs/example.ipynb @@ -219,7 +219,7 @@ "from aurora import Aurora, rollout\n", "\n", "model = Aurora(use_lora=False) # The pretrained version does not use LoRA.\n", - "model.load_checkpoint(\"wbruinsma/aurora\", \"aurora-0.25-pretrained.ckpt\")\n", + "model.load_checkpoint(\"microsoft/aurora\", \"aurora-0.25-pretrained.ckpt\")\n", "\n", "model.eval()\n", "model = model.to(\"cuda\")\n", diff --git a/docs/finetuning.md b/docs/finetuning.md index b7ff913..ce043a3 100644 --- a/docs/finetuning.md +++ b/docs/finetuning.md @@ -7,7 +7,7 @@ you should use the pretrained version: from aurora import Aurora model = Aurora(use_lora=False) # Model is not fine-tuned. -model.load_checkpoint("wbruinsma/aurora", "aurora-0.25-pretrained.ckpt") +model.load_checkpoint("microsoft/aurora", "aurora-0.25-pretrained.ckpt") ``` You are also free to extend the model for your particular use case. @@ -24,7 +24,7 @@ model = Aurora(...) ... # Modify `model`. -model.load_checkpoint("wbruinsma/aurora", "aurora-0.25-pretrained.ckpt", strict=False) +model.load_checkpoint("microsoft/aurora", "aurora-0.25-pretrained.ckpt", strict=False) ``` More instructions coming soon! diff --git a/docs/usage.md b/docs/usage.md index b5f34d2..8dee01a 100755 --- a/docs/usage.md +++ b/docs/usage.md @@ -80,10 +80,10 @@ model = AuroraSmall() ``` A checkpoint can then be loaded with `model.load_checkpoint` by specifying a checkpoint file in a HuggingFace repository. -In this case, we use `aurora-0.25-small-pretrained.ckpt` from the repository `wbruinsma/aurora`: +In this case, we use `aurora-0.25-small-pretrained.ckpt` from the repository `microsoft/aurora`: ```python -model.load_checkpoint("wbruinsma/aurora", "aurora-0.25-small-pretrained.ckpt") +model.load_checkpoint("microsoft/aurora", "aurora-0.25-small-pretrained.ckpt") ``` Typically, you will want to set the model to evaluation mode, which disables e.g. drop-out: diff --git a/tests/__init__.py b/tests/__init__.py index 8ed8289..a679a52 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -1,6 +1 @@ """Copyright (c) Microsoft Corporation. Licensed under the MIT license.""" - -import os - -if "HUGGINGFACE_REPO" not in os.environ: - raise RuntimeError("The environment variable `HUGGINGFACE_REPO` must be set.") diff --git a/tests/test_model.py b/tests/test_model.py index 8e9f418..4e6cdb9 100644 --- a/tests/test_model.py +++ b/tests/test_model.py @@ -1,6 +1,5 @@ """Copyright (c) Microsoft Corporation. Licensed under the MIT license.""" -import os import pickle from datetime import datetime from typing import TypedDict @@ -36,7 +35,7 @@ def test_aurora_small() -> None: # Load test input. path = hf_hub_download( - repo_id=os.environ["HUGGINGFACE_REPO"], + repo_id="microsoft/aurora", filename="aurora-0.25-small-pretrained-test-input.pickle", ) with open(path, "rb") as f: @@ -44,7 +43,7 @@ def test_aurora_small() -> None: # Load test output. path = hf_hub_download( - repo_id=os.environ["HUGGINGFACE_REPO"], + repo_id="microsoft/aurora", filename="aurora-0.25-small-pretrained-test-output.pickle", ) with open(path, "rb") as f: @@ -52,7 +51,7 @@ def test_aurora_small() -> None: # Load static variables. path = hf_hub_download( - repo_id=os.environ["HUGGINGFACE_REPO"], + repo_id="microsoft/aurora", filename="aurora-0.25-static.pickle", ) with open(path, "rb") as f: @@ -94,7 +93,7 @@ def interpolate(v: np.ndarray) -> np.ndarray: # Load the checkpoint and run the model. model.load_checkpoint( - os.environ["HUGGINGFACE_REPO"], + "microsoft/aurora", "aurora-0.25-small-pretrained.ckpt", strict=False, # LoRA parameters not available. ) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 509af0d..6018c17 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -1,6 +1,5 @@ """Copyright (c) Microsoft Corporation. Licensed under the MIT license.""" -import os from datetime import datetime, timedelta import numpy as np @@ -14,7 +13,7 @@ def test_rollout(): # LoRA for every step and the other does not. model1 = AuroraSmall(use_lora=True, lora_mode="single") model1.load_checkpoint( - os.environ["HUGGINGFACE_REPO"], + "microsoft/aurora", "aurora-0.25-small-pretrained.ckpt", strict=False, # LoRA parameters not available. )