From 04b2ab7f1db6d4c518bc10a1ad7739d359a528a4 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <6949769+tengomucho@users.noreply.github.com>
Date: Wed, 15 Jan 2025 10:12:35 +0100
Subject: [PATCH] feat(tpu): updated to optimum 0.2.3, added llama training
 example (#141)

* feat(tpu): updated to optimum 0.2.3, added llama training example

* fix: correct notebook filename for llama

* chore: remove training examples from tpu container

* chore: install optimum-tpu from pypi instead of git
---
 .../tpu/2.5.1/transformers/4.46.3/py310/Dockerfile         | 7 ++-----
 1 file changed, 2 insertions(+), 5 deletions(-)

diff --git a/containers/pytorch/training/tpu/2.5.1/transformers/4.46.3/py310/Dockerfile b/containers/pytorch/training/tpu/2.5.1/transformers/4.46.3/py310/Dockerfile
index dcfd60f9..e775e390 100644
--- a/containers/pytorch/training/tpu/2.5.1/transformers/4.46.3/py310/Dockerfile
+++ b/containers/pytorch/training/tpu/2.5.1/transformers/4.46.3/py310/Dockerfile
@@ -6,7 +6,7 @@ LABEL maintainer="Hugging Face"
 ARG DEBIAN_FRONTEND=noninteractive
 
 # Versions
-ARG OPTIMUM_TPU='0.2.2'
+ARG OPTIMUM_TPU='0.2.3'
 ARG TRANSFORMERS='4.46.3'
 ARG PEFT='0.13.2'
 ARG TRL='0.12.1'
@@ -30,10 +30,7 @@ RUN pip install --upgrade --no-cache-dir \
   jupyter notebook
 
 # Install Optimum TPU
-RUN pip install git+https://github.com/huggingface/optimum-tpu.git@v${OPTIMUM_TPU}
-# Add example
-ADD https://raw.githubusercontent.com/huggingface/optimum-tpu/v${OPTIMUM_TPU}/examples/language-modeling/gemma_tuning.ipynb \
-  /notebooks/gemma_tuning.ipynb
+RUN pip install --upgrade --no-cache-dir optimum-tpu==${OPTIMUM_TPU}
 
 # Install Google CLI single command
 RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main" \