Merge pull request #88 from jmisilo/jmisilo/fix/model-weights-url-mis…

…match Model weights url mismatch
jmisilo · Oct 29, 2023 · 4f74c7a · 4f74c7a
2 parents 8afd2a6 + 4663547
commit 4f74c7a
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 = (
-        "1pSQruQyg8KJq6VmzhMLFbT_VaHJMdlWF"
-        if model_size.strip().upper() == "L"
-        else "1Gh32arzhW06C1ZJyzcJSSfdJDi3RgWoG"
-    )
+    download_id = MODEL_WEIGHTS[model_size.strip().upper()]
 
     gdown.download(
         f"https://drive.google.com/uc?id={download_id}", checkpoint_fpath, quiet=False