Skip to content

Commit 6280d81

Browse files
authored
change: cover tei with image_uris.retrieve API (#4689)
1 parent 9faa8be commit 6280d81

File tree

1 file changed

+2
-0
lines changed

1 file changed

+2
-0
lines changed

src/sagemaker/image_uris.py

+2
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@
3737
ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
3838
HUGGING_FACE_FRAMEWORK = "huggingface"
3939
HUGGING_FACE_LLM_FRAMEWORK = "huggingface-llm"
40+
HUGGING_FACE_TEI_FRAMEWORK = "huggingface-tei"
4041
HUGGING_FACE_LLM_NEURONX_FRAMEWORK = "huggingface-llm-neuronx"
4142
XGBOOST_FRAMEWORK = "xgboost"
4243
SKLEARN_FRAMEWORK = "sklearn"
@@ -477,6 +478,7 @@ def _validate_version_and_set_if_needed(version, config, framework):
477478
if version is None and framework in [
478479
DATA_WRANGLER_FRAMEWORK,
479480
HUGGING_FACE_LLM_FRAMEWORK,
481+
HUGGING_FACE_TEI_FRAMEWORK,
480482
HUGGING_FACE_LLM_NEURONX_FRAMEWORK,
481483
STABILITYAI_FRAMEWORK,
482484
]:

0 commit comments

Comments
 (0)