Download embedding model in Dockerfile

openradx · Feb 11, 2024 · 7883b2e · 7883b2e
1 parent 9b2ba00
commit 7883b2e
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 6 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -43,6 +43,10 @@ RUN  mkdir /tmp/vespa-cli \
     && cp -r /tmp/vespa-cli/share/* /usr/local/share/ \
     && rm -rf /tmp/vespa-cli.tar.gz /tmp/vespa-cli
 
+# Download embedding model, which gets deployed by setup_vespa.py
+ADD "https://github.com/vespa-engine/sample-apps/raw/master/simple-semantic-search/model/e5-small-v2-int8.onnx" /opt/models/model.onnx
+ADD "https://raw.githubusercontent.com/vespa-engine/sample-apps/master/simple-semantic-search/model/tokenizer.json" /opt/models/tokenizer.json
+
 
 # `builder-base` stage is used to build deps + create our virtual environment
 FROM python-base as builder-base

diff --git a/radis/vespa/management/commands/setup_vespa.py b/radis/vespa/management/commands/setup_vespa.py
@@ -1,3 +1,4 @@
+import shutil
 import subprocess
 import tempfile
 from pathlib import Path
@@ -85,6 +86,11 @@ def handle(self, *args, **options):
             configurator = VespaConfigurator(app_folder)
             configurator.apply()
 
+            models_folder = app_folder / "files" / "models"
+            models_folder.mkdir()
+            shutil.copy("/opt/models/model.onnx", models_folder)
+            shutil.copy("/opt/models/tokenizer.json", models_folder)
+
         if options["deploy"]:
             vespa_host: str
             if options["host"]:

diff --git a/radis/vespa/vespa_app.py b/radis/vespa/vespa_app.py
@@ -19,10 +19,6 @@
 )
 
 REPORT_SCHEMA_NAME = "report"
-EMBEDDER_MODEL_URL = "https://github.com/vespa-engine/sample-apps/raw/master/\
-simple-semantic-search/model/e5-small-v2-int8.onnx"
-TOKENIZER_MODEL_URL = "https://raw.githubusercontent.com/vespa-engine/sample-apps/master/\
-simple-semantic-search/model/tokenizer.json"
 
 
 def _create_report_schema():
@@ -128,11 +124,11 @@ def _create_app_package(schemas: list[Schema]):
                 parameters=[
                     Parameter(
                         "transformer-model",
-                        {"url": EMBEDDER_MODEL_URL},
+                        {"path": "files/models/model.onnx"},
                     ),
                     Parameter(
                         "tokenizer-model",
-                        {"url": TOKENIZER_MODEL_URL},
+                        {"path": "files/models/tokenizer.json"},
                     ),
                 ],
             )