ARG SERVING_BUILD_IMAGE=tensorflow/tensorflow:2.16.1-gpu

FROM ${SERVING_BUILD_IMAGE}

WORKDIR /workspace

RUN apt-get update -y && apt-get install -y \

cmake

COPY requirements.txt requirements.txt

RUN pip install --upgrade --no-cache-dir pip \

&& pip install --no-cache-dir -r requirements.txt

COPY --from=apache/beam_python3.11_sdk:2.54.0 /opt/apache/beam /opt/apache/beam

COPY gemma_2b gemma_2B

COPY custom_model_gemma.py custom_model_gemma.py

ENV KERAS_BACKEND="tensorflow"

ENTRYPOINT ["/opt/apache/beam/boot"]