sagemaker-core 1.0.62__py3-none-any.whl → 2.1.1__py3-none-any.whl

sagemaker/core/jumpstart/payload_utils.py

@@ -0,0 +1,239 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""This module stores inference payload utilities for JumpStart models."""
+from __future__ import absolute_import
+import base64
+import json
+from typing import Any, Dict, List, Optional, Union
+import re
+import boto3
+
+from sagemaker.core.jumpstart.accessors import JumpStartS3PayloadAccessor
+from sagemaker.core.jumpstart.artifacts.payloads import _retrieve_example_payloads
+from sagemaker.core.jumpstart.constants import (
+    DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+)
+from sagemaker.core.jumpstart.enums import JumpStartModelType, MIMEType
+from sagemaker.core.jumpstart.types import JumpStartSerializablePayload
+from sagemaker.core.jumpstart.utils import (
+    get_jumpstart_content_bucket,
+    get_region_fallback,
+)
+from sagemaker.core.helper.session_helper import Session
+
+
+S3_BYTES_REGEX = r"^\$s3<(?P<s3_key>[a-zA-Z0-9-_/.]+)>$"
+S3_B64_STR_REGEX = r"\$s3_b64<(?P<s3_key>[a-zA-Z0-9-_/.]+)>"
+
+
+def _extract_field_from_json(
+    json_input: dict,
+    keys: List[str],
+) -> Any:
+    """Given a dictionary, returns value at specified keys.
+
+    Raises:
+        KeyError: If a key cannot be found in the json input.
+    """
+    curr_json = json_input
+    for idx, key in enumerate(keys):
+        if idx < len(keys) - 1:
+            curr_json = curr_json[key]
+            continue
+        return curr_json[key]
+
+
+def _construct_payload(
+    prompt: str,
+    model_id: str,
+    model_version: str,
+    region: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    alias: Optional[str] = None,
+) -> Optional[JumpStartSerializablePayload]:
+    """Returns example payload from prompt.
+
+    Args:
+        prompt (str): String-valued prompt to embed in payload.
+        model_id (str): JumpStart model ID of the JumpStart model for which to construct
+            the payload.
+        model_version (str): Version of the JumpStart model for which to retrieve the
+            payload.
+        region (Optional[str]): Region for which to retrieve the
+            payload. (Default: None).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated versions of model
+            specifications should be tolerated (exception not raised). If False, raises
+            an exception if the version of the model is deprecated. (Default: False).
+        sagemaker_session (sagemaker.session.Session): A SageMaker Session
+            object, used for SageMaker interactions. If not
+            specified, one is created using the default AWS configuration
+            chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        model_type (JumpStartModelType): The type of the model, can be open weights model or
+            proprietary model. (Default: JumpStartModelType.OPEN_WEIGHTS).
+    Returns:
+        Optional[JumpStartSerializablePayload]: serializable payload with prompt, or None if
+            this feature is unavailable for the specified model.
+    """
+    payloads: Optional[Dict[str, JumpStartSerializablePayload]] = _retrieve_example_payloads(
+        model_id=model_id,
+        model_version=model_version,
+        region=region,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        sagemaker_session=sagemaker_session,
+        model_type=model_type,
+    )
+    if payloads is None or len(payloads) == 0:
+        return None
+
+    payload_to_use: JumpStartSerializablePayload = (
+        payloads[alias] if alias else list(payloads.values())[0]
+    )
+
+    prompt_key: Optional[str] = payload_to_use.prompt_key
+    if prompt_key is None:
+        return None
+
+    payload_body = payload_to_use.body
+    prompt_key_split = prompt_key.split(".")
+    for idx, prompt_key in enumerate(prompt_key_split):
+        if idx < len(prompt_key_split) - 1:
+            payload_body = payload_body[prompt_key]
+        else:
+            payload_body[prompt_key] = prompt
+
+    return payload_to_use
+
+
+class PayloadSerializer:
+    """Utility class for serializing payloads associated with JumpStart models.
+
+    Many JumpStart models embed byte-streams into payloads corresponding to images, sounds,
+    and other content types which require downloading from S3.
+    """
+
+    def __init__(
+        self,
+        bucket: Optional[str] = None,
+        region: Optional[str] = None,
+        s3_client: Optional[boto3.client] = None,
+    ) -> None:
+        """Initializes PayloadSerializer object."""
+        self.bucket = bucket or get_jumpstart_content_bucket()
+        self.region = region or get_region_fallback(
+            s3_client=s3_client,
+        )
+        self.s3_client = s3_client
+
+    def get_bytes_payload_with_s3_references(
+        self,
+        payload_str: str,
+    ) -> bytes:
+        """Returns bytes object corresponding to referenced S3 object.
+
+        Raises:
+            ValueError: If the raw bytes payload is not formatted correctly.
+        """
+        s3_keys = re.compile(S3_BYTES_REGEX).findall(payload_str)
+        if len(s3_keys) != 1:
+            raise ValueError("Invalid bytes payload.")
+
+        s3_key = s3_keys[0]
+        serialized_s3_object = JumpStartS3PayloadAccessor.get_object_cached(
+            bucket=self.bucket, key=s3_key, region=self.region, s3_client=self.s3_client
+        )
+
+        return serialized_s3_object
+
+    def embed_s3_references_in_str_payload(
+        self,
+        payload: str,
+    ) -> str:
+        """Inserts serialized S3 content into string payload.
+
+        If no S3 content is embedded in payload, original string is returned.
+        """
+        return self._embed_s3_b64_references_in_str_payload(payload_body=payload)
+
+    def _embed_s3_b64_references_in_str_payload(
+        self,
+        payload_body: str,
+    ) -> str:
+        """Performs base 64 encoding of payloads embedded in a payload.
+
+        This is required so that byte-valued payloads can be transmitted efficiently
+        as a utf-8 encoded string.
+        """
+
+        s3_keys = re.compile(S3_B64_STR_REGEX).findall(payload_body)
+        for s3_key in s3_keys:
+            b64_encoded_string = base64.b64encode(
+                bytearray(
+                    JumpStartS3PayloadAccessor.get_object_cached(
+                        bucket=self.bucket, key=s3_key, region=self.region, s3_client=self.s3_client
+                    )
+                )
+            ).decode()
+            payload_body = payload_body.replace(f"$s3_b64<{s3_key}>", b64_encoded_string)
+        return payload_body
+
+    def embed_s3_references_in_json_payload(
+        self, payload_body: Union[list, dict, str, int, float]
+    ) -> Union[list, dict, str, int, float]:
+        """Finds all S3 references in payload and embeds serialized S3 data.
+
+        If no S3 references are found, the payload is returned un-modified.
+
+        Raises:
+            ValueError: If the payload has an unrecognized type.
+        """
+        if isinstance(payload_body, str):
+            return self.embed_s3_references_in_str_payload(payload_body)
+        if isinstance(payload_body, (float, int)):
+            return payload_body
+        if isinstance(payload_body, list):
+            return [self.embed_s3_references_in_json_payload(item) for item in payload_body]
+        if isinstance(payload_body, dict):
+            return {
+                key: self.embed_s3_references_in_json_payload(value)
+                for key, value in payload_body.items()
+            }
+        raise ValueError(f"Payload has unrecognized type: {type(payload_body)}")
+
+    def serialize(self, payload: JumpStartSerializablePayload) -> Union[str, bytes]:
+        """Returns payload string or bytes that can be inputted to inference endpoint.
+
+        Raises:
+            ValueError: If the payload has an unrecognized type.
+        """
+        content_type = MIMEType.from_suffixed_type(payload.content_type)
+        body = payload.body
+
+        if content_type in {MIMEType.JSON, MIMEType.LIST_TEXT, MIMEType.X_TEXT}:
+            body = self.embed_s3_references_in_json_payload(body)
+        else:
+            body = self.get_bytes_payload_with_s3_references(body)
+
+        if isinstance(body, dict):
+            body = json.dumps(body)
+        elif not isinstance(body, str) and not isinstance(body, bytes):
+            raise ValueError(f"Default payload '{body}' has unrecognized type: {type(body)}")
+
+        return body

sagemaker/core/jumpstart/region_config.json

@@ -0,0 +1,163 @@
+{
+    "af-south-1": {
+        "content_bucket": "jumpstart-cache-prod-af-south-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-af-south-1"
+    },
+    "ap-east-1": {
+        "content_bucket": "jumpstart-cache-prod-ap-east-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-east-1"
+    },
+    "ap-northeast-1": {
+        "content_bucket": "jumpstart-cache-prod-ap-northeast-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-northeast-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-ap-northeast-1"
+    },
+    "ap-northeast-2": {
+        "content_bucket": "jumpstart-cache-prod-ap-northeast-2",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-northeast-2",
+        "neo_content_bucket": "sagemaker-sd-models-prod-ap-northeast-2"
+    },
+    "ap-northeast-3": {
+        "content_bucket": "jumpstart-cache-prod-ap-northeast-3",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-northeast-3",
+        "neo_content_bucket": "sagemaker-sd-models-prod-ap-northeast-3"
+    },
+    "ap-south-1": {
+        "content_bucket": "jumpstart-cache-prod-ap-south-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-south-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-ap-south-1"
+    },
+    "ap-south-2": {
+        "content_bucket": "jumpstart-cache-prod-ap-south-2",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-south-2"
+    },
+    "ap-southeast-1": {
+        "content_bucket": "jumpstart-cache-prod-ap-southeast-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-ap-southeast-1"
+    },
+    "ap-southeast-2": {
+        "content_bucket": "jumpstart-cache-prod-ap-southeast-2",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-2",
+        "neo_content_bucket": "sagemaker-sd-models-prod-ap-southeast-2"
+    },
+    "ap-southeast-3": {
+        "content_bucket": "jumpstart-cache-prod-ap-southeast-3",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-3"
+    },
+    "ap-southeast-4": {
+        "content_bucket": "jumpstart-cache-prod-ap-southeast-4",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-4"
+    },
+    "ap-southeast-5": {
+        "content_bucket": "jumpstart-cache-prod-ap-southeast-5",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-5"
+    },
+    "ap-southeast-7": {
+        "content_bucket": "jumpstart-cache-prod-ap-southeast-7",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ap-southeast-7"
+    },
+    "ca-central-1": {
+        "content_bucket": "jumpstart-cache-prod-ca-central-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ca-central-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-ca-central-1"
+    },
+    "ca-west-1": {
+        "content_bucket": "jumpstart-cache-prod-ca-west-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-ca-west-1"
+    },
+    "cn-north-1": {
+        "content_bucket": "jumpstart-cache-prod-cn-north-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-cn-north-1"
+    },
+    "cn-northwest-1": {
+        "content_bucket": "jumpstart-cache-prod-cn-northwest-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-cn-northwest-1"
+    },
+    "eu-central-1": {
+        "content_bucket": "jumpstart-cache-prod-eu-central-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-eu-central-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-eu-central-1"
+    },
+    "eu-central-2": {
+        "content_bucket": "jumpstart-cache-prod-eu-central-2",
+        "gated_content_bucket": "jumpstart-private-cache-prod-eu-central-2"
+    },
+    "eu-north-1": {
+        "content_bucket": "jumpstart-cache-prod-eu-north-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-eu-north-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-eu-north-1"
+    },
+    "eu-south-1": {
+        "content_bucket": "jumpstart-cache-prod-eu-south-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-eu-south-1"
+    },
+    "eu-south-2": {
+        "content_bucket": "jumpstart-cache-prod-eu-south-2",
+        "gated_content_bucket": "jumpstart-private-cache-prod-eu-south-2"
+    },
+    "eu-west-1": {
+        "content_bucket": "jumpstart-cache-prod-eu-west-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-eu-west-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-eu-west-1"
+    },
+    "eu-west-2": {
+        "content_bucket": "jumpstart-cache-prod-eu-west-2",
+        "gated_content_bucket": "jumpstart-private-cache-prod-eu-west-2",
+        "neo_content_bucket": "sagemaker-sd-models-prod-eu-west-2"
+    },
+    "eu-west-3": {
+        "content_bucket": "jumpstart-cache-prod-eu-west-3",
+        "gated_content_bucket": "jumpstart-private-cache-prod-eu-west-3",
+        "neo_content_bucket": "sagemaker-sd-models-prod-eu-west-3"
+    },
+    "il-central-1": {
+        "content_bucket": "jumpstart-cache-prod-il-central-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-il-central-1"
+    },
+    "me-central-1": {
+        "content_bucket": "jumpstart-cache-prod-me-central-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-me-central-1"
+    },
+    "me-south-1": {
+        "content_bucket": "jumpstart-cache-prod-me-south-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-me-south-1"
+    },
+    "mx-central-1": {
+        "content_bucket": "jumpstart-cache-prod-mx-central-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-mx-central-1"
+    },
+    "sa-east-1": {
+        "content_bucket": "jumpstart-cache-prod-sa-east-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-sa-east-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-sa-east-1"
+    },
+    "us-east-1": {
+        "content_bucket": "jumpstart-cache-prod-us-east-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-east-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-us-east-1"
+    },
+    "us-east-2": {
+        "content_bucket": "jumpstart-cache-prod-us-east-2",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-east-2",
+        "neo_content_bucket": "sagemaker-sd-models-prod-us-east-2"
+    },
+    "us-gov-east-1": {
+        "content_bucket": "jumpstart-cache-prod-us-gov-east-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-gov-east-1"
+    },
+    "us-gov-west-1": {
+        "content_bucket": "jumpstart-cache-prod-us-gov-west-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-gov-west-1"
+    },
+    "us-west-1": {
+        "content_bucket": "jumpstart-cache-prod-us-west-1",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-west-1",
+        "neo_content_bucket": "sagemaker-sd-models-prod-us-west-1"
+    },
+    "us-west-2": {
+        "content_bucket": "jumpstart-cache-prod-us-west-2",
+        "gated_content_bucket": "jumpstart-private-cache-prod-us-west-2",
+        "neo_content_bucket": "sagemaker-sd-models-prod-us-west-2"
+    }
+}

sagemaker/core/jumpstart/search.py

@@ -0,0 +1,171 @@
+import re
+import logging
+from typing import List, Iterator, Optional
+from sagemaker.core.helper.session_helper import Session
+from sagemaker.core.resources import HubContent
+
+logger = logging.getLogger(__name__)
+
+
+class _Filter:
+    """
+    A filter that evaluates logical expressions against a list of keyword strings.
+
+    Supports logical operators (AND, OR, NOT), parentheses for grouping, and wildcard patterns
+    (e.g., `text-*`, `*ai`, `@task:foo`).
+
+    Example:
+        filt = _Filter("(@framework:huggingface OR text-*) AND NOT deprecated")
+        filt.match(["@framework:huggingface", "text-generation"])  # Returns True
+    """
+
+    def __init__(self, expression: str) -> None:
+        """
+        Initialize the filter with a string expression.
+
+        Args:
+            expression (str): A logical expression to evaluate against keywords.
+                Supports AND, OR, NOT, parentheses, and wildcard patterns (*).
+        """
+        self.expression: str = expression
+
+    def match(self, keywords: List[str]) -> bool:
+        """
+        Evaluate the filter expression against a list of keywords.
+
+        Args:
+            keywords (List[str]): A list of keyword strings to test.
+
+        Returns:
+            bool: True if the expression evaluates to True for the given keywords, else False.
+        """
+        expr: str = self._convert_expression(self.expression)
+        try:
+            return eval(expr, {"__builtins__": {}}, {"keywords": keywords, "any": any})
+        except Exception:
+            return False
+
+    def _convert_expression(self, expr: str) -> str:
+        """
+        Convert the logical filter expression into a Python-evaluable string.
+
+        Args:
+            expr (str): The raw expression to convert.
+
+        Returns:
+            str: A Python expression string using 'any' and logical operators.
+        """
+        tokens: List[str] = re.findall(
+            r"\bAND\b|\bOR\b|\bNOT\b|[^\s()]+|\(|\)", expr, flags=re.IGNORECASE
+        )
+
+        def wildcard_condition(pattern: str) -> str:
+            pattern = pattern.strip('"').strip("'")
+            stripped = pattern.strip("*")
+
+            if pattern.startswith("*") and pattern.endswith("*"):
+                return f"{repr(stripped)} in k"
+            elif pattern.startswith("*"):
+                return f"k.endswith({repr(stripped)})"
+            elif pattern.endswith("*"):
+                return f"k.startswith({repr(stripped)})"
+            else:
+                return f"k == {repr(pattern)}"
+
+        def convert_token(token: str) -> str:
+            upper = token.upper()
+            if upper == "AND":
+                return "and"
+            elif upper == "OR":
+                return "or"
+            elif upper == "NOT":
+                return "not"
+            elif token in ("(", ")"):
+                return token
+            else:
+                return f"any({wildcard_condition(token)} for k in keywords)"
+
+        converted_tokens = [convert_token(tok) for tok in tokens]
+        return " ".join(converted_tokens)
+
+
+def _list_all_hub_models(hub_name: str, sm_client: Session) -> Iterator[HubContent]:
+    """
+    Retrieve all model entries from the specified hub and yield them one by one.
+
+    This function paginates through the SageMaker Hub API to retrieve all published models of type "Model"
+    and yields them as `HubContent` objects.
+
+    Args:
+        hub_name (str): The name of the hub to query.
+        sm_client (Session): The SageMaker session.
+
+    Yields:
+        HubContent: A `HubContent` object representing a single model entry from the hub.
+    """
+    next_token = None
+
+    while True:
+        # Prepare the request parameters
+        params = {"HubName": hub_name, "HubContentType": "Model", "MaxResults": 100}
+
+        # Add NextToken if it exists
+        if next_token:
+            params["NextToken"] = next_token
+
+        # Make the API call
+        response = sm_client.list_hub_contents(**params)
+
+        # Yield each content summary
+        for content in response["HubContentSummaries"]:
+            yield HubContent(
+                hub_name=hub_name,
+                hub_content_arn=content["HubContentArn"],
+                hub_content_type="Model",
+                hub_content_name=content["HubContentName"],
+                hub_content_version=content["HubContentVersion"],
+                hub_content_description=content.get("HubContentDescription", ""),
+                hub_content_search_keywords=content.get("HubContentSearchKeywords", []),
+            )
+
+        # Check if there are more results
+        next_token = response.get("NextToken", None)
+        if not next_token or len(response["HubContentSummaries"]) == 0:
+            break  # Exit the loop if there are no more pages
+
+
+def search_public_hub_models(
+    query: str,
+    hub_name: Optional[str] = "SageMakerPublicHub",
+    sagemaker_session: Optional[Session] = None,
+) -> List[HubContent]:
+    """
+    Search and filter models from hub using a keyword expression.
+
+    Args:
+        query (str): A logical expression used to filter models by keywords.
+            Example: "@task:text-generation AND NOT @framework:legacy"
+        hub_name (Optional[str]): The name of the hub to query. Defaults to "SageMakerPublicHub".
+        sagemaker_session (Optional[Session]): An optional SageMaker `Session` object. If not provided,
+            a default session will be created and a warning will be logged.
+
+    Returns:
+        List[HubContent]: A list of filtered `HubContent` model objects that match the query.
+    """
+    if sagemaker_session is None:
+        sagemaker_session = Session()
+        logger.warning("SageMaker session not provided. Using default Session.")
+    sm_client = sagemaker_session.sagemaker_client
+
+    models = _list_all_hub_models(hub_name, sm_client)
+    filt = _Filter(query)
+    results: List[HubContent] = []
+
+    for model in models:
+        keywords = model.hub_content_search_keywords
+        normalized_keywords = [kw.replace(" ", "-") for kw in keywords]
+
+        if filt.match(normalized_keywords):
+            results.append(model)
+
+    return results

sagemaker/core/jumpstart/serializers.py

@@ -0,0 +1,81 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""JumpStart serializers module - provides retrieve_default function for backward compatibility."""
+from __future__ import absolute_import
+
+from typing import Optional
+
+from sagemaker.core.serializers import BaseSerializer
+from sagemaker.core.jumpstart import artifacts, utils as jumpstart_utils
+from sagemaker.core.jumpstart.constants import DEFAULT_JUMPSTART_SAGEMAKER_SESSION
+from sagemaker.core.jumpstart.enums import JumpStartModelType
+from sagemaker.core.helper.session_helper import Session
+
+
+def retrieve_default(
+    region: Optional[str] = None,
+    model_id: Optional[str] = None,
+    model_version: Optional[str] = None,
+    hub_arn: Optional[str] = None,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    model_type: JumpStartModelType = JumpStartModelType.OPEN_WEIGHTS,
+    sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
+    config_name: Optional[str] = None,
+) -> BaseSerializer:
+    """Retrieves the default serializer for the model matching the given arguments.
+
+    Args:
+        region (str): The AWS Region for which to retrieve the default serializer.
+            Defaults to ``None``.
+        model_id (str): The model ID of the model for which to
+            retrieve the default serializer. (Default: None).
+        model_version (str): The version of the model for which to retrieve the
+            default serializer. (Default: None).
+        hub_arn (str): The arn of the SageMaker Hub for which to retrieve
+            model details from. (Default: None).
+        tolerate_vulnerable_model (bool): True if vulnerable versions of model
+            specifications should be tolerated (exception not raised). If False, raises an
+            exception if the script used by this version of the model has dependencies with known
+            security vulnerabilities. (Default: False).
+        tolerate_deprecated_model (bool): True if deprecated models should be tolerated
+            (exception not raised). False if these models should raise an exception.
+            (Default: False).
+        model_type (JumpStartModelType): The model type. (Default: JumpStartModelType.OPEN_WEIGHTS).
+        sagemaker_session (sagemaker.session.Session): A SageMaker Session
+            object, used for SageMaker interactions. If not
+            specified, one is created using the default AWS configuration
+            chain. (Default: sagemaker.jumpstart.constants.DEFAULT_JUMPSTART_SAGEMAKER_SESSION).
+        config_name (Optional[str]): Name of the JumpStart Model config to apply. (Default: None).
+    Returns:
+        BaseSerializer: The default serializer to use for the model.
+
+    Raises:
+        ValueError: If the combination of arguments specified is not supported.
+    """
+    if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
+        raise ValueError(
+            "Must specify JumpStart `model_id` and `model_version` when retrieving serializers."
+        )
+
+    return artifacts._retrieve_default_serializer(
+        model_id=model_id,
+        model_version=model_version,
+        hub_arn=hub_arn,
+        region=region,
+        tolerate_deprecated_model=tolerate_deprecated_model,
+        tolerate_vulnerable_model=tolerate_vulnerable_model,
+        sagemaker_session=sagemaker_session,
+        model_type=model_type,
+        config_name=config_name,
+    )