File size: 1,926 Bytes
dbb98e1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
"""Script to download external data for the project at build time."""
import logging
import os
import tarfile
import wget
def download_and_extract_models() -> None:
"""Downloads the models folder from the server and extracts it."""
logging.debug("Downloading models folder.")
models_url = os.environ.get("MODELS_FOLDER_URL")
models_targz = "models.tar.gz"
models_folder = "data/models/"
try:
wget.download(models_url, models_targz)
logging.debug("Extracting models folder.")
with tarfile.open(models_targz, "r:gz") as tar:
tar.extractall(models_folder)
os.remove(models_targz)
logging.debug("Models folder downloaded and extracted.")
except Exception as e:
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."""
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/"
try:
wget.download(item_embeddings_url, item_embeddings_tarbz)
logging.debug("Extracting item embeddings folder.")
with tarfile.open(item_embeddings_tarbz, "r:bz2") as tar:
tar.extractall(item_embeddings_folder)
os.remove(item_embeddings_tarbz)
logging.debug("Item embeddings folder downloaded and extracted.")
except Exception as e:
logging.error(f"Error downloading item embeddings folder: {e}")
if __name__ == "__main__":
if not os.path.exists("data/models"):
logging.info("Downloading models...")
download_and_extract_models()
if not os.path.exists("data/embed_items"):
logging.info("Downloading item embeddings...")
download_and_extract_item_embeddings()
|