def test_metric_fn_signature(self):
    add_task = functools.partial(self.add_task, source=self.function_source)

    def score_metric_fn(targets, scores):
        return {}

    def predict_metric_fn(targets, predictions):
        return {}
    valid_task = add_task('valid_metrics', metric_fns=[score_metric_fn, predict_metric_fn])
    self.assertSameElements([score_metric_fn, predict_metric_fn], valid_task.metric_fns)
    self.assertSameElements([score_metric_fn], valid_task.score_metric_fns)
    self.assertSameElements([predict_metric_fn], valid_task.predict_metric_fns)

    def extra_arg_metric_fn(targets, predictions, extra_param):
        return {}
    expected_error_message_prefix = "Metric functions must have positional arguments matching either ('targets', 'predictions') or ('targets', 'scores'). Got: "
    with self.assertRaisesWithLiteralMatch(ValueError, (expected_error_message_prefix + "('targets', 'predictions', 'extra_param')")):
        valid_task = add_task('extra_arg_metric', metric_fns=[extra_arg_metric_fn])

    def bad_order_metric_fn(predictions, targets):
        return {}
    with self.assertRaisesWithLiteralMatch(ValueError, (expected_error_message_prefix + "('predictions', 'targets')")):
        valid_task = add_task('bad_order_metric', metric_fns=[bad_order_metric_fn])

    def bad_default_metric_fn(targets, predictions=0):
        return {}
    with self.assertRaisesWithLiteralMatch(ValueError, (expected_error_message_prefix + "('targets',)")):
        valid_task = add_task('bad_default_metric', metric_fns=[bad_default_metric_fn])

    def ok_default_metric_fn(targets, predictions, extra_param=3):
        return {}
    valid_task_2 = add_task('valid_metrics_2', metric_fns=[ok_default_metric_fn])
    self.assertSameElements([ok_default_metric_fn], valid_task_2.metric_fns)
    self.assertEmpty(valid_task_2.score_metric_fns)
    self.assertSameElements([ok_default_metric_fn], valid_task_2.predict_metric_fns)

    def predict_metric_fn_with_types(targets: Sequence[Mapping[(str, Any)]], predictions: Sequence[Mapping[(str, Any)]]) -> Mapping[(str, metrics_lib.MetricValue)]:
        return {}
    valid_task_with_types = TaskRegistry.add('valid_metrics_with_types', source=self.function_source, output_features={'inputs': dataset_providers.Feature(test_utils.sentencepiece_vocab()), 'targets': dataset_providers.Feature(test_utils.sentencepiece_vocab())}, metric_fns=[predict_metric_fn_with_types])
    self.assertSameElements([predict_metric_fn_with_types], valid_task_with_types.metric_fns)
def test_metric_fn_signature(self):
    add_task = functools.partial(self.add_task, source=self.function_source)

    def score_metric_fn(targets, scores):
        return {}

    def predict_metric_fn(targets, predictions):
        return {}
    valid_task = add_task('valid_metrics', metric_fns=[score_metric_fn, predict_metric_fn])
    self.assertSameElements([score_metric_fn, predict_metric_fn], valid_task.metric_fns)
    self.assertSameElements([score_metric_fn], valid_task.score_metric_fns)
    self.assertSameElements([predict_metric_fn], valid_task.predict_metric_fns)

    def extra_arg_metric_fn(targets, predictions, extra_param):
        return {}
    expected_error_message_prefix = "Metric functions must have positional arguments matching either ('targets', 'predictions') or ('targets', 'scores'). Got: "
    with self.assertRaisesWithLiteralMatch(ValueError, (expected_error_message_prefix + "('targets', 'predictions', 'extra_param')")):
        valid_task = add_task('extra_arg_metric', metric_fns=[extra_arg_metric_fn])

    def bad_order_metric_fn(predictions, targets):
        return {}
    with self.assertRaisesWithLiteralMatch(ValueError, (expected_error_message_prefix + "('predictions', 'targets')")):
        valid_task = add_task('bad_order_metric', metric_fns=[bad_order_metric_fn])

    def bad_default_metric_fn(targets, predictions=0):
        return {}
    with self.assertRaisesWithLiteralMatch(ValueError, (expected_error_message_prefix + "('targets',)")):
        valid_task = add_task('bad_default_metric', metric_fns=[bad_default_metric_fn])

    def ok_default_metric_fn(targets, predictions, extra_param=3):
        return {}
    valid_task_2 = add_task('valid_metrics_2', metric_fns=[ok_default_metric_fn])
    self.assertSameElements([ok_default_metric_fn], valid_task_2.metric_fns)
    self.assertEmpty(valid_task_2.score_metric_fns)
    self.assertSameElements([ok_default_metric_fn], valid_task_2.predict_metric_fns)

    def predict_metric_fn_with_types(targets: Sequence[Mapping[(str, Any)]], predictions: Sequence[Mapping[(str, Any)]]) -> Mapping[(str, metrics_lib.MetricValue)]:
        return {}
    valid_task_with_types = TaskRegistry.add('valid_metrics_with_types', source=self.function_source, output_features={'inputs': dataset_providers.Feature(test_utils.sentencepiece_vocab()), 'targets': dataset_providers.Feature(test_utils.sentencepiece_vocab())}, metric_fns=[predict_metric_fn_with_types])
    self.assertSameElements([predict_metric_fn_with_types], valid_task_with_types.metric_fns)