chore: remove deprecated generated_text_response_key from JumpStartSerializablePayload

Author: evakravi

src/sagemaker/jumpstart/payload_utils.py

@@ -115,80 +115,6 @@ def _construct_payload(
     return payload_to_use
 
 
-def _extract_generated_text_from_response(
-    response: dict,
-    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,
-    accept_type: Optional[str] = None,
-) -> str:
-    """Returns generated text extracted from full response payload.
-
-    Args:
-        response (dict): Dictionary-valued response from which to extract
-            generated text.
-        model_id (str): JumpStart model ID of the JumpStart model from which to extract
-            generated text.
-        model_version (str): Version of the JumpStart model for which to extract generated
-            text.
-        region (Optional[str]): Region for which to extract generated
-            text. (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).
-        accept_type (Optional[str]): The accept type to optionally specify for the response.
-            (Default: None).
-
-    Returns:
-        str: extracted generated text from the endpoint response payload.
-
-    Raises:
-        ValueError: If the model is invalid, the model does not support generated text extraction,
-        or if the response is malformed.
-    """
-
-    if not isinstance(response, dict):
-        raise ValueError(f"Response must be dictionary. Instead, got: {type(response)}")
-
-    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,
-    )
-    if payloads is None or len(payloads) == 0:
-        raise ValueError(f"Model ID '{model_id}' does not support generated text extraction.")
-
-    for payload in payloads.values():
-        if accept_type is None or payload.accept == accept_type:
-            generated_text_response_key: Optional[str] = payload.generated_text_response_key
-            if generated_text_response_key is None:
-                raise ValueError(
-                    f"Model ID '{model_id}' does not support generated text extraction."
-                )
-
-            generated_text_response_key_split = generated_text_response_key.split(".")
-            try:
-                return _extract_field_from_json(response, generated_text_response_key_split)
-            except KeyError:
-                raise ValueError(f"Response is malformed: {response}")
-
-    raise ValueError(f"Model ID '{model_id}' does not support generated text extraction.")
-
-
 class PayloadSerializer:
     """Utility class for serializing payloads associated with JumpStart models.
 

src/sagemaker/jumpstart/types.py

@@ -339,7 +339,6 @@ class JumpStartSerializablePayload(JumpStartDataHolderType):
         "content_type",
         "accept",
         "body",
-        "generated_text_response_key",
         "prompt_key",
     ]
 
@@ -371,7 +370,6 @@ def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None:
         self.content_type = json_obj["content_type"]
         self.body = json_obj["body"]
         accept = json_obj.get("accept")
-        self.generated_text_response_key = json_obj.get("generated_text_response_key")
         self.prompt_key = json_obj.get("prompt_key")
         if accept:
             self.accept = accept

tests/unit/sagemaker/jumpstart/constants.py

@@ -4129,7 +4129,6 @@
             "Dog": {
                 "content_type": "application/json",
                 "prompt_key": "hello.prompt",
-                "generated_text_response_key": "key1.key2.generated_text",
                 "body": {
                     "hello": {"prompt": "a dog"},
                     "seed": 43,

tests/unit/sagemaker/jumpstart/test_payload_utils.py

@@ -18,7 +18,6 @@
 
 from sagemaker.jumpstart.payload_utils import (
     PayloadSerializer,
-    _extract_generated_text_from_response,
     _construct_payload,
 )
 from sagemaker.jumpstart.types import JumpStartSerializablePayload
@@ -59,77 +58,6 @@ def test_construct_payload(self, patched_get_model_specs):
         )
 
 
-class TestResponseExtraction(TestCase):
-    @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
-    def test_extract_generated_text(self, patched_get_model_specs):
-        patched_get_model_specs.side_effect = get_special_model_spec
-
-        model_id = "response-keys"
-        region = "us-west-2"
-        generated_text = _extract_generated_text_from_response(
-            response={"key1": {"key2": {"generated_text": "top secret"}}},
-            model_id=model_id,
-            model_version="*",
-            region=region,
-        )
-
-        self.assertEqual(
-            _extract_generated_text_from_response(
-                response={"key1": {"key2": {"generated_text": "top secret"}}},
-                model_id=model_id,
-                model_version="*",
-                region=region,
-                accept_type="application/json",
-            ),
-            generated_text,
-        )
-
-        self.assertEqual(
-            generated_text,
-            "top secret",
-        )
-
-        with pytest.raises(ValueError):
-            _extract_generated_text_from_response(
-                response={"key1": {"key2": {"generated_texts": "top secret"}}},
-                model_id=model_id,
-                model_version="*",
-                region=region,
-            )
-
-        with pytest.raises(ValueError):
-            _extract_generated_text_from_response(
-                response={"key1": {"key2": {"generated_text": "top secret"}}},
-                model_id=model_id,
-                model_version="*",
-                region=region,
-                accept_type="blah/blah",
-            )
-
-        with pytest.raises(ValueError):
-            _extract_generated_text_from_response(
-                response={"key1": {"key2": {"generated_text": "top secret"}}},
-                model_id="env-var-variant-model",  # some model without the required metadata
-                model_version="*",
-                region=region,
-            )
-        with pytest.raises(ValueError):
-            _extract_generated_text_from_response(
-                response={"key1": {"generated_texts": "top secret"}},
-                model_id=model_id,
-                model_version="*",
-                region=region,
-            )
-
-        with pytest.raises(ValueError):
-            _extract_generated_text_from_response(
-                response="blah",
-                model_id=model_id,
-                model_version="*",
-                region=region,
-            )
-
-
 class TestPayloadSerializer(TestCase):
 
     payload_serializer = PayloadSerializer()

tests/unit/sagemaker/jumpstart/test_predictor.py

@@ -76,7 +76,7 @@ def test_jumpstart_serializable_payload_with_predictor(
         "JumpStartSerializablePayload: {'content_type': 'application/json', 'accept': 'application/json'"
         ", 'body': {'prompt': 'a dog', 'num_images_per_prompt': 2, 'num_inference_steps':"
         " 20, 'guidance_scale': 7.5, 'seed': 43, 'eta': 0.7, 'image':"
-        " '$s3_b64<inference-notebook-assets/inpainting_cow.jpg>'}, 'generated_text_response_key': None}"
+        " '$s3_b64<inference-notebook-assets/inpainting_cow.jpg>'}}"
     )
 
     js_predictor.predict(default_payload)