FROM nroggendorff/train-llama:latest

ENV HOME=/home/user \
  HF_HOME=$HOME/.cache/huggingface \
  PATH=$HOME/.local/bin:$PATH
ENV DATASETS_CACHE=$HOME/.cache/huggingface/datasets \
  HF_DATASETS_CACHE=$HOME/.cache/huggingface/datasets \
  TRITON_CACHE_DIR=$HOME/.cache/.triton/

ENV TOKENIZERS_PARALLELISM=false \
  PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True