From 10b783ced344bc542c5d4b858c0cea31775bf87a Mon Sep 17 00:00:00 2001
From: Max Kuhmichel <kuhmichel.max@fh-swf.de>
Date: Tue, 24 Sep 2024 10:16:12 +0000
Subject: [PATCH] small util for downloading hf models

---
 images/downloadtool/Dockerfile        | 12 ++++++++++++
 images/downloadtool/README.md         |  1 +
 images/downloadtool/download_model.py | 15 +++++++++++++++
 images/downloadtool/requirements.txt  |  1 +
 4 files changed, 29 insertions(+)
 create mode 100644 images/downloadtool/Dockerfile
 create mode 100644 images/downloadtool/README.md
 create mode 100644 images/downloadtool/download_model.py
 create mode 100644 images/downloadtool/requirements.txt

diff --git a/images/downloadtool/Dockerfile b/images/downloadtool/Dockerfile
new file mode 100644
index 0000000..95485d8
--- /dev/null
+++ b/images/downloadtool/Dockerfile
@@ -0,0 +1,12 @@
+FROM python:3
+
+USER root
+
+WORKDIR /usr/src/app
+
+COPY requirements.txt ./
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY . .
+
+CMD [ "python", "./download_model.py" ]
\ No newline at end of file
diff --git a/images/downloadtool/README.md b/images/downloadtool/README.md
new file mode 100644
index 0000000..7d61511
--- /dev/null
+++ b/images/downloadtool/README.md
@@ -0,0 +1 @@
+# Downloads a model form huggingface
\ No newline at end of file
diff --git a/images/downloadtool/download_model.py b/images/downloadtool/download_model.py
new file mode 100644
index 0000000..177b17f
--- /dev/null
+++ b/images/downloadtool/download_model.py
@@ -0,0 +1,15 @@
+from huggingface_hub import snapshot_download
+import os
+
+model_id = os.environ.get("MODEL", None)
+model_folder = os.environ.get("MODEL_DIR", "/models")
+os.makedirs(model_folder, exist_ok=True)
+
+if model_folder is None:
+    print("INFO: Using default model dir: /models")
+
+if model_id is None:
+    print("WARNING: Not all envs where specified, can not download anything")
+else:
+    snapshot_download(repo_id=model_id, local_dir=os.path.join(model_folder, model_id), local_dir_use_symlinks=False, revision="main")
+    
\ No newline at end of file
diff --git a/images/downloadtool/requirements.txt b/images/downloadtool/requirements.txt
new file mode 100644
index 0000000..b2642b3
--- /dev/null
+++ b/images/downloadtool/requirements.txt
@@ -0,0 +1 @@
+huggingface_hub
\ No newline at end of file