remove unused imports and fix docstyle

aws · Feb 29, 2024 · 4d58cfd · 4d58cfd
1 parent 5aeaa35
commit 4d58cfd
Show file tree

Hide file tree

Showing 5 changed files with 2 additions and 10 deletions.
diff --git a/src/sagemaker/accept_types.py b/src/sagemaker/accept_types.py
@@ -16,7 +16,6 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
-from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -73,7 +72,6 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
-    model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_SOURCE,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,

diff --git a/src/sagemaker/content_types.py b/src/sagemaker/content_types.py
@@ -16,7 +16,6 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
-from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -73,7 +72,6 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
-    model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_SOURCE,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,

diff --git a/src/sagemaker/deserializers.py b/src/sagemaker/deserializers.py
@@ -35,7 +35,6 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
-from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -96,7 +95,6 @@ def retrieve_default(
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
-    model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_SOURCE,
 ) -> BaseDeserializer:
     """Retrieves the default deserializer for the model matching the given arguments.
 

diff --git a/src/sagemaker/jumpstart/factory/model.py b/src/sagemaker/jumpstart/factory/model.py
@@ -204,8 +204,8 @@ def _add_instance_type_to_kwargs(
 
 
 def _add_image_uri_to_kwargs(kwargs: JumpStartModelInitKwargs) -> JumpStartModelInitKwargs:
-    """
-    Sets image uri based on default or override, returns full kwargs.
+    """Sets image uri based on default or override, returns full kwargs.
+
     Uses placeholder image uri for JumpStart proprietary models that uses ModelPackages
     """
 

diff --git a/src/sagemaker/serializers.py b/src/sagemaker/serializers.py
@@ -34,7 +34,6 @@
 
 from sagemaker.jumpstart import artifacts, utils as jumpstart_utils
 from sagemaker.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
-from sagemaker.jumpstart.enums import JumpStartModelType
 from sagemaker.session import Session
 
 
@@ -92,7 +91,6 @@ def retrieve_default(
     region: Optional[str] = None,
     model_id: Optional[str] = None,
     model_version: Optional[str] = None,
-    model_type: Optional[JumpStartModelType] = JumpStartModelType.OPEN_SOURCE,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,