File tree 2 files changed +9
-9
lines changed
src/sagemaker/serve/builder
tests/unit/sagemaker/serve/builder
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]:
107
107
model_task = hf_model_md .get ("pipeline_tag" )
108
108
109
109
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
+ )
112
113
113
114
logger .info ("Auto detected %s. Proceeding with the the deployment." , self .image_uri )
114
115
Original file line number Diff line number Diff line change @@ -154,16 +154,15 @@ def test_image_uri_override(
154
154
return_value = "sentence-similarity" ,
155
155
)
156
156
@patch (
157
- "from sagemaker.huggingface.get_huggingface_llm_image_uri" ,
158
- return_value = MOCK_IMAGE_CONFIG
157
+ "sagemaker.huggingface.get_huggingface_llm_image_uri" , return_value = MOCK_IMAGE_CONFIG
159
158
)
160
159
@patch ("sagemaker.serve.builder.transformers_builder._capture_telemetry" , side_effect = None )
161
160
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 ,
167
166
):
168
167
builder = ModelBuilder (
169
168
model = mock_model_id ,
You can’t perform that action at this time.
0 commit comments