From f45250cc64126e3c3bc5cf631f999891bb56f447 Mon Sep 17 00:00:00 2001
From: pagezyhf <165770107+pagezyhf@users.noreply.github.com>
Date: Fri, 23 Aug 2024 09:53:56 +0200
Subject: [PATCH] update container URIs (#680)

* update container URIs

* Update docs/source/containers.mdx

Co-authored-by: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com>

* TGI 0.0.24

---------

Co-authored-by: Jingya HUANG <44135271+JingyaHuang@users.noreply.github.com>
---
 docs/source/containers.mdx | 19 +++++++++++++++----
 1 file changed, 15 insertions(+), 4 deletions(-)

diff --git a/docs/source/containers.mdx b/docs/source/containers.mdx
index c05c59b3d..a03c31861 100644
--- a/docs/source/containers.mdx
+++ b/docs/source/containers.mdx
@@ -15,15 +15,26 @@ specific language governing permissions and limitations under the License.
 We provide pre-built Optimum Neuron containers for Amazon SageMaker. These containers come with all of the Hugging Face libraries and dependencies pre-installed, so you can start using them right away.
 We have containers for training and inference, and optimized text generation containers with TGI. The table is up to date and only includes the latest versions of each container. You can find older versions in the [Deep Learning Container Release Notes](https://github.com/aws/deep-learning-containers/releases?q=hf-neuronx&expanded=true)
 
-We recommend using the `sagemaker` Python SDK to retrieve the image URI for the container you want to use.
+We recommend using the `sagemaker` Python SDK to retrieve the image URI for the container you want to use. Here is a code snippet to retrieve the latest Text Generation Inference container Image URI:
+```python
+from sagemaker.huggingface import get_huggingface_llm_image_uri
+ 
+# retrieve the llm image uri
+llm_image = get_huggingface_llm_image_uri(
+  "huggingface-neuronx"
+)
+
+print(f"llm image uri: {llm_image}")
+
+```
 
 ## Available Optimum Neuron Containers
 
 | Type                       | Optimum Version | Image URI                                   |
 |-----------------------------|-----------------|---------------------------------------------|
-| Training  | 0.0.21           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training-neuronx:1.13.1-transformers4.36.2-neuronx-py310-sdk2.18.0-ubuntu20.04`   |
-| Inference      | 0.0.22           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference-neuronx:2.1.2-transformers4.36.2-neuronx-py310-sdk2.18.0-ubuntu20.04`      |
-| Text Generation Inference        | 0.0.22           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.2-optimum0.0.22-neuronx-py310-ubuntu22.04`        |
+| Training  | 0.0.24           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training-neuronx:2.1.2-transformers4.41.1-neuronx-py310-sdk2.19.1-ubuntu20.04`   |
+| Inference      | 0.0.24           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference-neuronx:2.1.2-transformers4.41.1-neuronx-py310-sdk2.19.1-ubuntu20.04`      |
+| Text Generation Inference        | 0.0.24           | `763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.2-optimum0.0.24-neuronx-py310-ubuntu22.04`        |
 
 
 Please replace `763104351884` with the correct [AWS account ID](https://github.com/aws/sagemaker-python-sdk/blob/master/src/sagemaker/image_uri_config/huggingface-neuronx.json) and `region` with the AWS region you are working in.