update: model weight urls definition

jmisilo · Oct 29, 2023 · 4663547 · 4663547
1 parent 21f7165
commit 4663547
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/src/utils/downloads.py b/src/utils/downloads.py
@@ -4,17 +4,18 @@
 
 import gdown
 
+MODEL_WEIGHTS = {
+    "L": "1Gh32arzhW06C1ZJyzcJSSfdJDi3RgWoG",
+    "S": "1pSQruQyg8KJq6VmzhMLFbT_VaHJMdlWF",
+}
+
 
 def download_weights(checkpoint_fpath, model_size="L"):
     """
     Downloads weights from Google Drive.
     """
 
-    download_id = (
-        "1Gh32arzhW06C1ZJyzcJSSfdJDi3RgWoG"
-        if model_size.strip().upper() == "L"
-        else "1pSQruQyg8KJq6VmzhMLFbT_VaHJMdlWF"
-    )
+    download_id = MODEL_WEIGHTS[model_size.strip().upper()]
 
     gdown.download(
         f"https://drive.google.com/uc?id={download_id}", checkpoint_fpath, quiet=False