Skip to content

Commit 43ce1ba

Browse files
committed
Fix formatting
1 parent 2e00238 commit 43ce1ba

File tree

2 files changed

+9
-9
lines changed

2 files changed

+9
-9
lines changed

src/sagemaker/serve/builder/transformers_builder.py

+3-2
Original file line numberDiff line numberDiff 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

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

+6-7
Original file line numberDiff line numberDiff 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+
"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,

0 commit comments

Comments
 (0)