fix: Change class names and entrypoints to 0pResolution

ecmwf-lab · Oct 7, 2024 · faea8f5 · faea8f5
1 parent dc1a0cb
commit faea8f5
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -60,8 +60,8 @@ optional-dependencies.dev = [
 urls.Repository = "https://github.com/ecmwf-lab/ai-models-aurora"
 entry-points."ai_models.model".aurora = "ai_models_aurora.model:model"
 entry-points."ai_models.model"."aurora-0.1-finetuned" = "ai_models_aurora.model:Aurora0p1FineTuned"
-entry-points."ai_models.model"."aurora-2.5-finetuned" = "ai_models_aurora.model:Aurora2p5FineTuned"
-entry-points."ai_models.model"."aurora-2.5-pretrained" = "ai_models_aurora.model:Aurora2p5Pretrained"
+entry-points."ai_models.model"."aurora-2.5-finetuned" = "ai_models_aurora.model:Aurora0p25FineTuned"
+entry-points."ai_models.model"."aurora-2.5-pretrained" = "ai_models_aurora.model:Aurora0p25Pretrained"
 
 [tool.setuptools_scm]
 version_file = "src/ai_models_aurora/_version.py"
diff --git a/src/ai_models_aurora/model.py b/src/ai_models_aurora/model.py
@@ -193,7 +193,7 @@ class Aurora0p25Pretrained(Aurora0p25):
 
 
 # https://microsoft.github.io/aurora/models.html#aurora-0-25-fine-tuned
-class Aurora025FineTuned(Aurora0p25):
+class Aurora0p25FineTuned(Aurora0p25):
     use_lora = True
     checkpoint = "aurora-0.25-finetuned.ckpt"
 
@@ -241,7 +241,7 @@ def model(model_version, **kwargs):
 
     models = {
         "0.25-pretrained": Aurora0p25Pretrained,
-        "0.25-finetuned": Aurora025FineTuned,
+        "0.25-finetuned": Aurora0p25FineTuned,
         "0.1-finetuned": Aurora0p1FineTuned,
         "default": Aurora0p1FineTuned,
         "latest": Aurora0p1FineTuned,  # Backward compatibility