aws · danielezhu · Jul 11, 2024 · Jul 10, 2024 · Jul 11, 2024 · Jul 11, 2024
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -37,7 +37,7 @@ bert-score = "^0.3.13"
 scikit-learn = "^1.3.1"
 jiwer = "^3.0.3"
 transformers = "^4.36.0"
-sagemaker = "^2.219.0"
+sagemaker = "^2.225.0"
 testbook = "^0.4.2"
 ipykernel = "^6.26.0"
 mypy-boto3-bedrock = "^1.33.2"

diff --git a/src/fmeval/constants.py b/src/fmeval/constants.py
@@ -123,6 +123,7 @@ class DatasetColumns(Enum):
 # Jumpstart
 JUMPSTART_MODEL_ID = "jumpstart_model_id"
 JUMPSTART_MODEL_VERSION = "jumpstart_model_version"
+JUMPSTART_MODEL_TYPE = "jumpstart_model_type"
 MODEL_ID = "model_id"
 SPEC_KEY = "spec_key"
 DEFAULT_PAYLOADS = "default_payloads"

diff --git a/src/fmeval/model_runners/extractors/__init__.py b/src/fmeval/model_runners/extractors/__init__.py
@@ -1,6 +1,14 @@
 from typing import Optional
 
-from fmeval.constants import MIME_TYPE_JSON, JUMPSTART_MODEL_ID, JUMPSTART_MODEL_VERSION, IS_EMBEDDING_MODEL
+from sagemaker.jumpstart.enums import JumpStartModelType
+
+from fmeval.constants import (
+    MIME_TYPE_JSON,
+    JUMPSTART_MODEL_ID,
+    JUMPSTART_MODEL_VERSION,
+    JUMPSTART_MODEL_TYPE,
+    IS_EMBEDDING_MODEL,
+)
 from fmeval.exceptions import EvalAlgorithmClientError
 from fmeval.model_runners.extractors.json_extractor import JsonExtractor
 from fmeval.model_runners.extractors.jumpstart_extractor import JumpStartExtractor
@@ -25,6 +33,9 @@ def create_extractor(
         extractor = JumpStartExtractor(
             jumpstart_model_id=kwargs[JUMPSTART_MODEL_ID],
             jumpstart_model_version=kwargs[JUMPSTART_MODEL_VERSION] if JUMPSTART_MODEL_VERSION in kwargs else "*",
+            jumpstart_model_type=kwargs[JUMPSTART_MODEL_TYPE]
+            if JUMPSTART_MODEL_TYPE in kwargs
+            else JumpStartModelType.OPEN_WEIGHTS,
             is_embedding_model=kwargs[IS_EMBEDDING_MODEL] if IS_EMBEDDING_MODEL in kwargs else False,
         )
     else:  # pragma: no cover

diff --git a/src/fmeval/model_runners/extractors/jumpstart_extractor.py b/src/fmeval/model_runners/extractors/jumpstart_extractor.py
@@ -1,4 +1,5 @@
 import json
+import logging
 import os
 from typing import Union, List, Dict, Optional
 from urllib import request
@@ -7,26 +8,30 @@
 from functional import seq
 from jmespath.exceptions import JMESPathError
 from sagemaker import Session
+from sagemaker.jumpstart.enums import JumpStartScriptScope
+from sagemaker.jumpstart.utils import verify_model_region_and_return_specs
 
 from fmeval import util
 from fmeval.constants import (
     MODEL_ID,
-    SPEC_KEY,
     GENERATED_TEXT_JMESPATH_EXPRESSION,
     SDK_MANIFEST_FILE,
     PROPRIETARY_SDK_MANIFEST_FILE,
-    DEFAULT_PAYLOADS,
     JUMPSTART_BUCKET_BASE_URL_FORMAT,
     JUMPSTART_BUCKET_BASE_URL_FORMAT_ENV_VAR,
     INPUT_LOG_PROBS_JMESPATH_EXPRESSION,
     EMBEDDING_JMESPATH_EXPRESSION,
+    SPEC_KEY,
+    DEFAULT_PAYLOADS,
 )
 from fmeval.exceptions import EvalAlgorithmClientError, EvalAlgorithmInternalError
 from fmeval.model_runners.extractors.extractor import Extractor
 
 # The expected model response location for Jumpstart that do produce the log probabilities
 from fmeval.model_runners.util import get_sagemaker_session
 
+logger = logging.getLogger(__name__)
+
 
 class JumpStartExtractor(Extractor):
     """
@@ -37,6 +42,7 @@ def __init__(
         self,
         jumpstart_model_id: str,
         jumpstart_model_version: str,
+        jumpstart_model_type: str,
         is_embedding_model: Optional[bool] = False,
         sagemaker_session: Optional[Session] = None,
     ):
@@ -51,6 +57,7 @@ def __init__(
         """
         self._model_id = jumpstart_model_id
         self._model_version = jumpstart_model_version
+        self._model_type = jumpstart_model_type
         self._sagemaker_session = sagemaker_session if sagemaker_session else get_sagemaker_session()
         self._is_embedding_model = is_embedding_model
 
@@ -62,22 +69,41 @@ def __init__(
             lambda x: x.get(MODEL_ID, None) == jumpstart_model_id
         )
         util.require(model_manifest, f"Model {jumpstart_model_id} is not a valid JumpStart Model")
-        model_spec_key = self.get_jumpstart_sdk_spec(
+
+        model_spec = self.get_jumpstart_sdk_spec(
             model_manifest.get(SPEC_KEY, None),
             self._sagemaker_session.boto_region_name,
         )
-        util.require(
-            DEFAULT_PAYLOADS in model_spec_key, f"JumpStart Model: {jumpstart_model_id} is not supported at this time"
-        )
+        if DEFAULT_PAYLOADS not in model_spec:
+            # Model spec contains alt configs, which should
+            # be obtained through JumpStart util function.
+            logger.info(
+                "default_payloads not found as a top-level attribute of model spec"
+                "Searching for default_payloads in inference configs instead."
+            )
+            model_spec = verify_model_region_and_return_specs(
+                region=self._sagemaker_session.boto_region_name,
+                model_id=self._model_id,
+                version=self._model_version,
+                model_type=self._model_type,
+                scope=JumpStartScriptScope.INFERENCE,
+                sagemaker_session=self._sagemaker_session,
+            )
+            configs = model_spec.inference_configs  # type: ignore[attr-defined]
+            util.require(configs, f"JumpStart Model: {jumpstart_model_id} is not supported at this time")
+            default_payloads = configs.get_top_config_from_ranking().resolved_metadata_config[DEFAULT_PAYLOADS]
+        else:
+            # Continue to extract default payloads by manually parsing the spec json object.
+            # TODO: update this code when the `default_payloads` attribute of JumpStartModelSpecs
+            #  returns the full data, including fields like generated_text.
+            default_payloads = model_spec[DEFAULT_PAYLOADS]
+
+        util.require(default_payloads, f"JumpStart Model: {jumpstart_model_id} is not supported at this time")
 
-        output_jmespath_expressions = None
-        input_log_probs_jmespath_expressions = None
         try:
-            output_jmespath_expressions = jmespath.compile(GENERATED_TEXT_JMESPATH_EXPRESSION).search(
-                model_spec_key[DEFAULT_PAYLOADS]
-            )
+            output_jmespath_expressions = jmespath.compile(GENERATED_TEXT_JMESPATH_EXPRESSION).search(default_payloads)
             input_log_probs_jmespath_expressions = jmespath.compile(INPUT_LOG_PROBS_JMESPATH_EXPRESSION).search(
-                model_spec_key[DEFAULT_PAYLOADS]
+                default_payloads
             )
         except (TypeError, JMESPathError) as e:
             raise EvalAlgorithmInternalError(

diff --git a/src/fmeval/model_runners/sm_jumpstart_model_runner.py b/src/fmeval/model_runners/sm_jumpstart_model_runner.py
@@ -54,7 +54,17 @@ def __init__(
         :param component_name: Name of the Amazon SageMaker inference component corresponding
                             the predictor
         """
+        sagemaker_session = get_sagemaker_session()
+        util.require(
+            is_endpoint_in_service(sagemaker_session, endpoint_name),
+            f"Endpoint {endpoint_name} is not in service",
+        )
+        # Default model type is always OPEN_WEIGHTS. See https://tinyurl.com/yc58s6wj
+        jumpstart_model_type = JumpStartModelType.OPEN_WEIGHTS
+        if is_proprietary_js_model(sagemaker_session.boto_region_name, model_id):
+            jumpstart_model_type = JumpStartModelType.PROPRIETARY
         is_text_embedding_model = is_text_embedding_js_model(model_id)
+
         super().__init__(
             content_template=content_template,
             output=output,
@@ -64,6 +74,7 @@ def __init__(
             accept_type=MIME_TYPE_JSON,
             jumpstart_model_id=model_id,
             jumpstart_model_version=model_version,
+            jumpstart_model_type=jumpstart_model_type,
             is_embedding_model=is_text_embedding_model,
         )
         self._endpoint_name = endpoint_name
@@ -77,17 +88,6 @@ def __init__(
         self._component_name = component_name
         self._is_embedding_model = is_text_embedding_model
 
-        sagemaker_session = get_sagemaker_session()
-        util.require(
-            is_endpoint_in_service(sagemaker_session, self._endpoint_name),
-            f"Endpoint {self._endpoint_name} is not in service",
-        )
-
-        # Default model type is always OPEN_WEIGHTS. See https://tinyurl.com/yc58s6wj
-        jumpstart_model_type = JumpStartModelType.OPEN_WEIGHTS
-        if is_proprietary_js_model(sagemaker_session.boto_region_name, self._model_id):
-            jumpstart_model_type = JumpStartModelType.PROPRIETARY
-
         predictor = sagemaker.predictor.retrieve_default(
             endpoint_name=self._endpoint_name,
             model_id=self._model_id,

diff --git a/test/integration/test_create_extractor.py b/test/integration/test_create_extractor.py
@@ -0,0 +1,42 @@
+import pytest
+
+from fmeval.constants import MIME_TYPE_JSON
+from fmeval.exceptions import EvalAlgorithmClientError
+from fmeval.model_runners.extractors import create_extractor, JumpStartExtractor
+
+
+class TestCreateExtractor:
+    """
+    These tests are under integration tests instead of unit tests because
+    credentials are required to call the JumpStart util function
+    verify_model_region_and_return_specs.
+
+    See test/unit/model_runners/extractors/test_create_extractor.py
+    for corresponding unit tests.
+    """
+
+    def test_create_extractor_jumpstart(self):
+        """
+        GIVEN a model whose default payloads are not found at the top level of
+            the model spec, but instead nested under the inference_configs attribute.
+        WHEN create_extractor is called with this model id.
+        THEN a JumpStartExtractor is successfully created for this model.
+        """
+        # default payloads found in inference_component_configs
+        jumpstart_model_id = "huggingface-llm-mistral-7b"
+        assert isinstance(
+            create_extractor(model_accept_type=MIME_TYPE_JSON, jumpstart_model_id=jumpstart_model_id),
+            JumpStartExtractor,
+        )
+
+    def test_create_extractor_jumpstart_no_default_payloads(self):
+        """
+        GIVEN a model whose spec does not contain default payloads data anywhere.
+        WHEN a create_extractor is called with this model id.
+        THEN the correct exception is raised.
+        """
+        with pytest.raises(
+            EvalAlgorithmClientError,
+            match="JumpStart Model: xgboost-regression-snowflake is not supported at this time",
+        ):
+            create_extractor(model_accept_type=MIME_TYPE_JSON, jumpstart_model_id="xgboost-regression-snowflake")
diff --git a/test/unit/model_runners/extractors/test_create_extractor.py b/test/unit/model_runners/extractors/test_create_extractor.py
@@ -9,9 +9,17 @@ def test_create_extractor():
     assert isinstance(create_extractor(model_accept_type=MIME_TYPE_JSON, output_location="output"), JsonExtractor)
 
 
-def test_create_extractor_jumpstart():
+@pytest.mark.parametrize(
+    "jumpstart_model_id", ["huggingface-llm-falcon-7b-bf16"]  # default payloads found top level of model spec
+)
+def test_create_extractor_jumpstart(jumpstart_model_id):
+    """
+    Note: the test case for a model whose default payloads are found in inference_configs
+        (instead of as a top-level attribute of the model spec) is an integration test,
+        since unit tests don't run with the credentials required.
+    """
     assert isinstance(
-        create_extractor(model_accept_type=MIME_TYPE_JSON, jumpstart_model_id="huggingface-llm-falcon-7b-bf16"),
+        create_extractor(model_accept_type=MIME_TYPE_JSON, jumpstart_model_id=jumpstart_model_id),
         JumpStartExtractor,
     )