def predict()

in src/fmeval/model_runners/sm_model_runner.py [0:0]

• 25 lines of code
• 5 McCabe index (conditional complexity)

    def predict(self, prompt: str) -> Union[Tuple[Optional[str], Optional[float]], List[float]]:
        """
        Invoke the SageMaker endpoint and parse the model response.
        :param prompt: Input data for which you want the model to provide inference.
        """
        composed_data = self._composer.compose(prompt)
        model_output = self._predictor.predict(
            data=composed_data,
            custom_attributes=self._custom_attributes,
            component_name=self._component_name,
        )

        embedding = (
            self._extractor.extract_embedding(data=model_output, num_records=1)
            if self._extractor.embedding_jmespath_expression
            else None
        )
        if embedding:
            return embedding

        output = (
            self._extractor.extract_output(data=model_output, num_records=1)
            if self._extractor.output_jmespath_expression
            else None
        )
        log_probability = (
            self._extractor.extract_log_probability(data=model_output, num_records=1)
            if self._extractor.log_probability_jmespath_expression
            else None
        )
        return output, log_probability