def extract_log_probability()

in src/fmeval/model_runners/extractors/json_extractor.py [0:0]

• 17 lines of code
• 4 McCabe index (conditional complexity)

    def extract_log_probability(self, data: Union[List, Dict], num_records: int) -> Union[List[float], float]:
        """
        Extract log probability from model response.

        :param data: Model response. The log_probability_jmespath_expression is used to extract the log probabilities
                     of the input tokens. Each record in the extracted probabilities will be a float or list of floats.
                     Examples for the extracted probabilities:
                        - data: 0.1, num_records: 1, num tokens: 1 (or probabilities already summed up)
                        - data: [0.1], num_records: 1, num tokens: 1 (or probabilities already summed up)
                        - data: [0.1, 0.2], num_records: 1, num tokens: 2
        :param num_records: number of inference records in the model output
        :return: float or list of float where each float is sum of log probabilities.
        """
        assert num_records == 1, "JSON extractor does not support batch requests"
        util.require(
            self.log_probability_jmespath_expression,
            "Extractor cannot extract log_probability as log_probability_jmespath_expression is not provided",
        )
        log_probs = self.log_probability_jmespath.search(data)
        util.require(
            log_probs is not None, f"JMESpath {self.log_probability_jmespath_expression} could not find any data"
        )
        if isinstance(log_probs, float):
            return log_probs
        util.require(
            isinstance(log_probs, List) and all(isinstance(value, float) for value in log_probs),
            f"Extractor found: {log_probs} which does not match expected {float} or list of {float}",
        )
        return sum(log_probs)