Update: Pull latest tei container for sentence similiarity models

Author: samruds

src/sagemaker/serve/builder/transformers_builder.py

@@ -23,7 +23,7 @@
     _get_nb_instance,
 )
 from sagemaker.djl_inference.model import _get_model_config_properties_from_hf
-from sagemaker.huggingface import HuggingFaceModel
+from sagemaker.huggingface import HuggingFaceModel, get_huggingface_llm_image_uri
 from sagemaker.serve.model_server.multi_model_server.prepare import (
     _create_dir_structure,
 )
@@ -47,6 +47,7 @@ class Transformers(ABC):
     """Transformers build logic with ModelBuilder()"""
 
     def __init__(self):
+        self.model_metadata = None
         self.model = None
         self.serve_settings = None
         self.sagemaker_session = None
@@ -99,7 +100,26 @@ def _create_transformers_model(self) -> Type[Model]:
         if hf_model_md is None:
             raise ValueError("Could not fetch HF metadata")
 
-        if "pytorch" in hf_model_md.get("tags"):
+        model_task = None
+        if self.model_metadata:
+            model_task = self.model_metadata.get("HF_TASK")
+        else:
+            model_task = hf_model_md.get("pipeline_tag")
+        
+        if model_task == "sentence-similarity" and not self.image_uri:
+            self.image_uri = \
+                get_huggingface_llm_image_uri("huggingface-tei", session=self.sagemaker_session)
+
+            logger.info("Auto detected %s. Proceeding with the the deployment.", self.image_uri)
+
+            pysdk_model = HuggingFaceModel(
+                env=self.env_vars,
+                role=self.role_arn,
+                sagemaker_session=self.sagemaker_session,
+                image_uri=self.image_uri,
+                vpc_config=self.vpc_config,
+            )
+        elif "pytorch" in hf_model_md.get("tags"):
             self.pytorch_version = self._get_supported_version(
                 hf_config, base_hf_version, "pytorch"
             )

tests/unit/sagemaker/serve/builder/test_transformers_builder.py

@@ -110,7 +110,7 @@ def test_build_deploy_for_transformers_local_container_and_remote_container(
         return_value="ml.g5.24xlarge",
     )
     @patch("sagemaker.serve.builder.transformers_builder._capture_telemetry", side_effect=None)
-    def test_image_uri(
+    def test_image_uri_override(
         self,
         mock_get_nb_instance,
         mock_telemetry,
@@ -144,3 +144,52 @@ def test_image_uri(
 
         with self.assertRaises(ValueError) as _:
             model.deploy(mode=Mode.IN_PROCESS)
+
+    @patch(
+        "sagemaker.serve.builder.transformers_builder._get_nb_instance",
+        return_value="ml.g5.24xlarge",
+    )
+    @patch(
+        "sagemaker.huggingface.llm_utils.get_huggingface_model_metadata",
+        return_value="sentence-similarity",
+    )
+    @patch(
+        "from sagemaker.huggingface.get_huggingface_llm_image_uri",
+        return_value=MOCK_IMAGE_CONFIG
+    )
+    @patch("sagemaker.serve.builder.transformers_builder._capture_telemetry", side_effect=None)
+    def test_sentence_similarity_support(
+            self,
+            mock_get_nb_instance,
+            mock_task,
+            mock_image,
+            mock_telemetry,
+    ):
+        builder = ModelBuilder(
+            model=mock_model_id,
+            schema_builder=mock_schema_builder,
+            mode=Mode.LOCAL_CONTAINER,
+        )
+
+        builder._prepare_for_mode = MagicMock()
+        builder._prepare_for_mode.side_effect = None
+
+        model = builder.build()
+        builder.serve_settings.telemetry_opt_out = True
+
+        builder.modes[str(Mode.LOCAL_CONTAINER)] = MagicMock()
+        predictor = model.deploy(model_data_download_timeout=1800)
+
+        assert builder.image_uri == MOCK_IMAGE_CONFIG
+        assert builder.env_vars["MODEL_LOADING_TIMEOUT"] == "1800"
+        assert isinstance(predictor, TransformersLocalModePredictor)
+
+        assert builder.nb_instance_type == "ml.g5.24xlarge"
+
+        builder._original_deploy = MagicMock()
+        builder._prepare_for_mode.return_value = (None, {})
+        predictor = model.deploy(mode=Mode.SAGEMAKER_ENDPOINT, role="mock_role_arn")
+        assert "HF_MODEL_ID" in model.env
+
+        with self.assertRaises(ValueError) as _:
+            model.deploy(mode=Mode.IN_PROCESS)