Update data download

Dockerfile

@@ -6,12 +6,12 @@ COPY . .
 # Install requirements
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Expose Hugging Face Space secrets to environment variables
-RUN --mount=type=secret,id=models_folder_url,mode=0444,required=true echo "MODELS_FOLDER_URL=$(cat /run/secrets/models_folder_url)" >> .env
-RUN --mount=type=secret,id=item_embeddings_url,mode=0444,required=true echo "ITEM_EMBEDDINGS_URL=$(cat /run/secrets/item_embeddings_url)" >> .env
-
-# Download external data
-RUN python download_external_data.py
+# Expose Hugging Face Space secrets and download external data
+RUN --mount=type=secret,id=models_folder_url,mode=0444,required=true \
+    --mount=type=secret,id=item_embeddings_url,mode=0444,required=true \
+    python download_external_data.py \
+    $(cat /run/secrets/models_folder_url) \
+    $(cat /run/secrets/item_embeddings_url)
 
 EXPOSE 7860
 

download_external_data.py

@@ -1,5 +1,6 @@
 """Script to download external data for the project at build time."""
 
+import argparse
 import logging
 import os
 import tarfile
@@ -7,10 +8,13 @@ import tarfile
 import wget
 
 
-def download_and_extract_models() -> None:
-    """Downloads the models folder from the server and extracts it."""
+def download_and_extract_models(models_url: str) -> None:
+    """Downloads the models folder from the server and extracts it.
+
+    Args:
+        models_url: URL to download the models from.
+    """
     logging.debug("Downloading models folder.")
-    models_url = os.environ.get("MODELS_FOLDER_URL")
     models_targz = "models.tar.gz"
     models_folder = "data/models/"
     try:
@@ -27,10 +31,13 @@ def download_and_extract_models() -> None:
         logging.error(f"Error downloading models folder: {e}")
 
 
-def download_and_extract_item_embeddings() -> None:
-    """Downloads the item embeddings folder from the server and extracts it."""
+def download_and_extract_item_embeddings(item_embeddings_url: str) -> None:
+    """Downloads the item embeddings folder from the server and extracts it.
+
+    Args:
+        item_embeddings_url: URL to download the item embeddings from.
+    """
     logging.debug("Downloading item embeddings folder.")
-    item_embeddings_url = os.environ.get("ITEM_EMBEDDINGS_URL")
     item_embeddings_tarbz = "item_embeddings.tar.bz2"
     item_embeddings_folder = "data/"
 
@@ -50,12 +57,31 @@ def download_and_extract_item_embeddings() -> None:
         logging.error(f"Error downloading item embeddings folder: {e}")
 
 
+def parse_args() -> argparse.Namespace:
+    """Parses command line arguments."""
+    parser = argparse.ArgumentParser(
+        description="Download external data for the project."
+    )
+    parser.add_argument(
+        "models-url", type=str, help="URL to download the models folder."
+    )
+    parser.add_argument(
+        "item-embeddings-url",
+        type=str,
+        help="URL to download the item embeddings folder",
+    )
+    return parser.parse_args()
+
+
 if __name__ == "__main__":
     logging.basicConfig(level=logging.DEBUG)
+
+    args = parse_args()
+
     if not os.path.exists("data/models"):
         logging.info("Downloading models...")
-        download_and_extract_models()
+        download_and_extract_models(args.models_url)
 
     if not os.path.exists("data/embed_items"):
         logging.info("Downloading item embeddings...")
-        download_and_extract_item_embeddings()
+        download_and_extract_item_embeddings(args.item_embeddings_url)