File tree Expand file tree Collapse file tree 2 files changed +9
-9
lines changed
src/sagemaker/serve/builder
tests/unit/sagemaker/serve/builder Expand file tree Collapse file tree 2 files changed +9
-9
lines changed Original file line number Diff line number Diff line change @@ -107,8 +107,9 @@ def _create_transformers_model(self) -> Type[Model]:
107107 model_task = hf_model_md .get ("pipeline_tag" )
108108
109109 if model_task == "sentence-similarity" and not self .image_uri :
110- self .image_uri = \
111- get_huggingface_llm_image_uri ("huggingface-tei" , session = self .sagemaker_session )
110+ self .image_uri = get_huggingface_llm_image_uri (
111+ "huggingface-tei" , session = self .sagemaker_session
112+ )
112113
113114 logger .info ("Auto detected %s. Proceeding with the the deployment." , self .image_uri )
114115
Original file line number Diff line number Diff line change @@ -154,16 +154,15 @@ def test_image_uri_override(
154154 return_value = "sentence-similarity" ,
155155 )
156156 @patch (
157- "from sagemaker.huggingface.get_huggingface_llm_image_uri" ,
158- return_value = MOCK_IMAGE_CONFIG
157+ "from sagemaker.huggingface.get_huggingface_llm_image_uri" , return_value = MOCK_IMAGE_CONFIG
159158 )
160159 @patch ("sagemaker.serve.builder.transformers_builder._capture_telemetry" , side_effect = None )
161160 def test_sentence_similarity_support (
162- self ,
163- mock_get_nb_instance ,
164- mock_task ,
165- mock_image ,
166- mock_telemetry ,
161+ self ,
162+ mock_get_nb_instance ,
163+ mock_task ,
164+ mock_image ,
165+ mock_telemetry ,
167166 ):
168167 builder = ModelBuilder (
169168 model = mock_model_id ,
You can’t perform that action at this time.
0 commit comments