fmeval.model_runners.bedrock_model_runner

  1"""
  2Module to manage model runners for Bedrock models.
  3"""
  4import json
  5import logging
  6from fmeval.util import require
  7from typing import Optional, Tuple, List, Union
  8from fmeval.constants import MIME_TYPE_JSON
  9from fmeval.model_runners.model_runner import ModelRunner
 10from fmeval.model_runners.util import get_bedrock_runtime_client
 11
 12logger = logging.getLogger(__name__)
 13
 14
 15class BedrockModelRunner(ModelRunner):
 16    """
 17    A class to manage the creation and deletion of Bedrock model runner when user provides
 18    a Bedrock model id.
 19    """
 20
 21    def __init__(
 22        self,
 23        model_id: str,
 24        content_template: str,
 25        output: Optional[str] = None,
 26        log_probability: Optional[str] = None,
 27        embedding: Optional[str] = None,
 28        content_type: str = MIME_TYPE_JSON,
 29        accept_type: str = MIME_TYPE_JSON,
 30    ):
 31        """
 32        :param model_id: Id of the Bedrock model to be used for model predictions
 33        :param content_template: String template to compose the model input from the prompt
 34        :param output: JMESPath expression of output in the model output
 35        :param log_probability: JMESPath expression of log probability in the model output
 36        :param embedding: JMESPath expression of embedding in the model output
 37        :param content_type: The content type of the request sent to the model for inference
 38        :param accept_type: The accept type of the request sent to the model for inference
 39        """
 40        super().__init__(content_template, output, log_probability, embedding, content_type, accept_type)
 41        self._model_id = model_id
 42        self._content_template = content_template
 43        self._output = output
 44        self._log_probability = log_probability
 45        self._content_type = content_type
 46        self._accept_type = accept_type
 47        self._embedding = embedding
 48
 49        require(
 50            output is not None or log_probability is not None or embedding is not None,
 51            "One of output jmespath expression, log probability or embedding jmespath expression must be provided",
 52        )
 53        require(self._accept_type == MIME_TYPE_JSON, f"Model accept type `{self._accept_type}` is not supported.")
 54        require(
 55            self._content_type == MIME_TYPE_JSON,
 56            f"Model content type `{self._content_type}` is not supported.",
 57        )
 58        self._bedrock_runtime_client = get_bedrock_runtime_client()
 59
 60    def predict(self, prompt: str) -> Union[Tuple[Optional[str], Optional[float]], List[float]]:
 61        """
 62        Invoke the Bedrock model and parse the model response.
 63        :param prompt: Input data for which you want the model to provide inference.
 64        """
 65        composed_data = self._composer.compose(prompt)
 66        body = json.dumps(composed_data)
 67        response = self._bedrock_runtime_client.invoke_model(
 68            body=body, modelId=self._model_id, accept=self._accept_type, contentType=self._content_type
 69        )
 70        model_output = json.loads(response.get("body").read())
 71
 72        embedding = (
 73            self._extractor.extract_embedding(data=model_output, num_records=1)
 74            if self._extractor.embedding_jmespath_expression
 75            else None
 76        )
 77        if embedding:
 78            return embedding
 79
 80        output = (
 81            self._extractor.extract_output(data=model_output, num_records=1)
 82            if self._extractor.output_jmespath_expression
 83            else None
 84        )
 85        log_probability = (
 86            self._extractor.extract_log_probability(data=model_output, num_records=1)
 87            if self._extractor.log_probability_jmespath_expression
 88            else None
 89        )
 90        return output, log_probability
 91
 92    def __reduce__(self):
 93        """
 94        Custom serializer method used by Ray when it serializes instances of this
 95        class in eval_algorithms.util.generate_model_predict_response_for_dataset.
 96        """
 97        serialized_data = (
 98            self._model_id,
 99            self._content_template,
100            self._output,
101            self._log_probability,
102            self._embedding,
103            self._content_type,
104            self._accept_type,
105        )
106        return self.__class__, serialized_data