Dockerfile.cuda-build

ARG TRANSLATOR_MODEL_NAME='winstxnhdw/nllb-200-distilled-1.3B-ct2-int8'
ARG LANGUAGE_DETECTOR_MODEL_NAME='facebook/fasttext-language-identification'


FROM python:3.12-slim AS python-builder

ENV POETRY_VIRTUALENVS_CREATE=false
ENV POETRY_HOME=/opt/poetry
ENV PATH=$POETRY_HOME/bin:$PATH

COPY pyproject.toml poetry.lock ./

RUN python -m venv $POETRY_HOME
RUN $POETRY_HOME/bin/pip install poetry
RUN poetry install --without=dev


FROM python:slim AS model-builder

ARG TRANSLATOR_MODEL_NAME
ARG LANGUAGE_DETECTOR_MODEL_NAME

RUN pip install huggingface_hub
RUN python -c \
    "from huggingface_hub import snapshot_download, hf_hub_download; \
    snapshot_download('$TRANSLATOR_MODEL_NAME'); \
    hf_hub_download('$LANGUAGE_DETECTOR_MODEL_NAME', 'model.bin')"


FROM python:3.12-slim

ARG TRANSLATOR_MODEL_NAME
ARG LANGUAGE_DETECTOR_MODEL_NAME
ENV HOME=/home/user
ENV PYTHONUNBUFFERED=1
ENV PYTHONDONTWRITEBYTECODE=1
ENV USE_CUDA=True
ENV LD_LIBRARY_PATH=/usr/local/lib/python3.12/site-packages/nvidia/cublas/lib

RUN useradd -m -u 1000 user

USER user

WORKDIR $HOME/app

COPY --chown=user --from=model-builder /root/.cache/huggingface $HOME/.cache/huggingface
COPY --chown=user --from=python-builder /usr/local/lib/python3.12/site-packages /usr/local/lib/python3.12/site-packages
COPY --chown=user . $HOME/app

CMD ["python", "main.py"]