From 46635475cea16ffb2aeba5386fb5eafe5dec465c Mon Sep 17 00:00:00 2001
From: jmisilo <jmisilo@interia.pl>
Date: Sun, 29 Oct 2023 20:30:39 +0100
Subject: [PATCH] update: model weight urls definition

---
 src/utils/downloads.py | 11 ++++++-----
 1 file changed, 6 insertions(+), 5 deletions(-)

diff --git a/src/utils/downloads.py b/src/utils/downloads.py
index 75ef840..b0ef2bc 100644
--- 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