From 6414293ffe7adff7cf8a83f0c25548de1986af0d Mon Sep 17 00:00:00 2001
From: Olly Welch <tools.kyr4g@aleeas.com>
Date: Tue, 28 Feb 2023 21:45:58 +0000
Subject: [PATCH] Use multi stage build to slim down ML image size

---
 machine-learning/Dockerfile | 29 ++++++++++++++++++-----------
 1 file changed, 18 insertions(+), 11 deletions(-)

diff --git a/machine-learning/Dockerfile b/machine-learning/Dockerfile
index acf92053ba..5c905295c8 100644
--- a/machine-learning/Dockerfile
+++ b/machine-learning/Dockerfile
@@ -1,18 +1,25 @@
-FROM python:3.10
+FROM python:3.10 as builder
 
-ENV TRANSFORMERS_CACHE=/cache
-ENV PYTHONDONTWRITEBYTECODE 1
-ENV PYTHONUNBUFFERED 1
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=true
+
+RUN python -m venv /opt/venv && \
+    /opt/venv/bin/pip install --pre torch  -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html && \
+    /opt/venv/bin/pip install transformers tqdm numpy scikit-learn scipy nltk sentencepiece flask Pillow && \
+    /opt/venv/bin/pip install --no-deps sentence-transformers
+
+FROM python:3.10-slim
+
+COPY --from=builder /opt/venv /opt/venv
+
+ENV TRANSFORMERS_CACHE=/cache \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PATH="/opt/venv/bin:$PATH"
 
 WORKDIR /usr/src/app
 
-RUN python -m venv /opt/venv
-ENV PATH="/opt/venv/bin:$PATH"
-
-RUN pip install --no-cache-dir torch==1.13.1+cpu -f https://download.pytorch.org/whl/torch_stable.html
-RUN pip install transformers tqdm numpy scikit-learn scipy nltk sentencepiece flask Pillow
-RUN pip install --no-deps sentence-transformers
-
 COPY . .
 
 CMD ["python", "src/main.py"]