Fix formatting

aws · May 15, 2024 · 513e604 · 513e604
1 parent 2e00238
commit 513e604
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 9 deletions.
diff --git a/src/sagemaker/serve/builder/transformers_builder.py b/src/sagemaker/serve/builder/transformers_builder.py
@@ -107,8 +107,9 @@ def _create_transformers_model(self) -> Type[Model]:
             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)
+            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)
 

diff --git a/tests/unit/sagemaker/serve/builder/test_transformers_builder.py b/tests/unit/sagemaker/serve/builder/test_transformers_builder.py
@@ -154,16 +154,15 @@ def test_image_uri_override(
         return_value="sentence-similarity",
     )
     @patch(
-        "from sagemaker.huggingface.get_huggingface_llm_image_uri",
-        return_value=MOCK_IMAGE_CONFIG
+        "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,
+        self,
+        mock_get_nb_instance,
+        mock_task,
+        mock_image,
+        mock_telemetry,
     ):
         builder = ModelBuilder(
             model=mock_model_id,