From e828658332e01dd49e8ae25a0cd1d99fda6ee4b8 Mon Sep 17 00:00:00 2001
From: Piti Cookie <piti_cookie@helilot.com>
Date: Thu, 31 Mar 2022 07:21:20 -0700
Subject: [PATCH] Added models include option in docker build

---
 Dockerfile        | 10 +++++++---
 install_models.py |  8 ++++++--
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/Dockerfile b/Dockerfile
index 82601ce..9aae39e 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -1,6 +1,7 @@
 FROM python:3.8.12-slim-bullseye
 
 ARG with_models=false
+ARG models=
 
 WORKDIR /app
 
@@ -14,14 +15,17 @@ RUN pip install --upgrade pip
 
 COPY . .
 
-# check for offline build
+
 RUN if [ "$with_models" = "true" ]; then  \
         # install only the dependencies first
         pip install -e .;  \
         # initialize the language models
-        ./install_models.py;  \
+        if [ ! -z "$models" ]; then \
+                  ./install_models.py --load_only_lang_codes "$models";   \
+        else \
+                  ./install_models.py;  \
+        fi \
     fi
-
 # Install package from source code
 RUN pip install . \
   && pip cache purge
diff --git a/install_models.py b/install_models.py
index 23631f4..a6aa2ef 100755
--- a/install_models.py
+++ b/install_models.py
@@ -1,7 +1,11 @@
 #!/usr/bin/env python
-
+import argparse
 from app.init import check_and_install_models, check_and_install_transliteration
 
 if __name__ == "__main__":
-    check_and_install_models(force=True)
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--load_only_lang_codes", type=str, default="")
+    args = parser.parse_args()
+    lang_codes = args.load_only_lang_codes.split(",") or None
+    check_and_install_models(force=True, load_only_lang_codes=lang_codes)
     check_and_install_transliteration(force=True)