From 89050654fe1395231dfbb9b7c6f9384bd4c26db6 Mon Sep 17 00:00:00 2001
From: Patrick Schork <patrick@numer.ai>
Date: Wed, 8 Nov 2023 13:23:58 -0800
Subject: [PATCH] Adds treelite package Upgrades tensorflow & keras to 2.14.0
 Fixes predictions validation assertion

---
 predict.py        |  2 +-
 py3.10/Dockerfile |  7 +++++--
 py3.11/Dockerfile |  7 +++++--
 py3.9/Dockerfile  |  3 +++
 requirements.txt  | 13 +++++++------
 5 files changed, 21 insertions(+), 11 deletions(-)

diff --git a/predict.py b/predict.py
index cbd70d0..ac2c71c 100644
--- a/predict.py
+++ b/predict.py
@@ -171,7 +171,7 @@ def main(args):
             elif predictions.isna().any().any():
                 logging.error("Pickle function returned at least 1 NaN prediction")
                 exit_with_help(1)
-            elif not (predictions.between(0, 1).all().all()):
+            elif not (predictions.iloc[:, 0].between(0, 1).all().all()):
                 logging.error(
                     "Pickle function returned invalid predictions. Ensure values are between 0 and 1."
                 )
diff --git a/py3.10/Dockerfile b/py3.10/Dockerfile
index 164d4b4..89e21ae 100644
--- a/py3.10/Dockerfile
+++ b/py3.10/Dockerfile
@@ -2,10 +2,13 @@ FROM python:3.10-slim
 
 RUN apt-get update -yqq \
     && apt-get install -yqq --no-install-recommends \
-    build-essential \
-    python3-dev
+        build-essential \
+        cmake \
+        git \
+        python3-dev
 
 COPY ./requirements.txt .
+RUN pip install --upgrade pip
 RUN pip install -r requirements.txt
 
 COPY ./predict.py .
diff --git a/py3.11/Dockerfile b/py3.11/Dockerfile
index 06a3c67..598d54a 100644
--- a/py3.11/Dockerfile
+++ b/py3.11/Dockerfile
@@ -2,10 +2,13 @@ FROM python:3.11-slim
 
 RUN apt-get update -yqq \
     && apt-get install -yqq --no-install-recommends \
-    build-essential \
-    python3-dev
+        build-essential \
+        cmake \
+        git \
+        python3-dev
 
 COPY ./requirements.txt .
+RUN pip install --upgrade pip
 RUN pip install -r requirements.txt
 
 COPY ./predict.py .
diff --git a/py3.9/Dockerfile b/py3.9/Dockerfile
index 5bccd0a..88404aa 100644
--- a/py3.9/Dockerfile
+++ b/py3.9/Dockerfile
@@ -3,9 +3,12 @@ FROM python:3.9-slim
 RUN apt-get update -yqq \
     && apt-get install -yqq --no-install-recommends \
         build-essential \
+        cmake \
+        git \
         python3-dev
 
 COPY ./requirements.txt .
+RUN pip install --upgrade pip
 RUN pip install -r requirements.txt
 
 COPY ./predict.py .
diff --git a/requirements.txt b/requirements.txt
index c3160fa..c6a89d0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -29,7 +29,7 @@ idna==3.4
 jax==0.4.11
 Jinja2==3.1.2
 joblib==1.2.0
-keras==2.12.0
+keras==2.14.0
 kiwisolver==1.4.4
 libclang==16.0.0
 lightgbm==4.0.0
@@ -68,16 +68,17 @@ scipy==1.10.1
 six==1.16.0
 sympy==1.12
 tenacity==8.2.2
-tensorboard==2.12.3
-tensorboard-data-server==0.7.0
-tensorflow==2.12.0
-tensorflow-estimator==2.12.0
-tensorflow-io-gcs-filesystem==0.32.0
+tensorboard==2.14.1
+tensorboard-data-server==0.7.2
+tensorflow==2.14.0
+tensorflow-estimator==2.14.0
+tensorflow-io-gcs-filesystem==0.34.0
 termcolor==2.3.0
 threadpoolctl==3.1.0
 torch==2.0.1
 torchmetrics==0.11.4
 tqdm==4.65.0
+treelite==3.9.1
 typing_extensions==4.6.3
 tzdata==2023.3
 urllib3==1.26.15