Skip to content

Commit

Permalink
[0.27.0-dlc] pin datasets version in tensorrt llm (#2106)
Browse files Browse the repository at this point in the history
  • Loading branch information
sindhuvahinis authored Jun 26, 2024
1 parent 42963b7 commit cbd139a
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion serving/docker/tensorrt-llm.Dockerfile
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ ARG ammo_version=0.7.0
ARG janus_version=1.0.0
ARG pynvml_verison=11.5.0
ARG numpy_version=1.26.4
ARG datasets_version=2.17.1

EXPOSE 8080

Expand Down Expand Up @@ -73,7 +74,7 @@ RUN apt-get update && apt-get install -y g++ wget unzip openmpi-bin libopenmpi-d
# Install PyTorch
# Qwen needs transformers_stream_generator, tiktoken and einops
RUN pip install torch==${TORCH_VERSION} transformers==${transformers_version} accelerate==${accelerate_version} ${peft_wheel} sentencepiece \
mpi4py cuda-python==${cuda_python_version} onnx polygraphy pynvml==${pynvml_verison} datasets pydantic==${pydantic_version} scipy torchprofile bitsandbytes ninja \
mpi4py cuda-python==${cuda_python_version} onnx polygraphy pynvml==${pynvml_verison} datasets==${datasets_version} pydantic==${pydantic_version} scipy torchprofile bitsandbytes ninja \
transformers_stream_generator einops tiktoken jinja2 graphviz blobfile && \
pip3 cache purge

Expand Down

0 comments on commit cbd139a

Please sign in to comment.