From 43bd575002a99e3e1e91400ff46ae1f39f57b3bb Mon Sep 17 00:00:00 2001 From: Chirp Team Date: Thu, 27 Jul 2023 16:21:38 -0700 Subject: [PATCH] Internal PiperOrigin-RevId: 551671177 --- chirp/eval/eval_lib.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/chirp/eval/eval_lib.py b/chirp/eval/eval_lib.py index 667114c6..4afab992 100644 --- a/chirp/eval/eval_lib.py +++ b/chirp/eval/eval_lib.py @@ -464,7 +464,7 @@ def _prepare_eval_set( for class_name in eval_set_specification.class_names: choice_key, rng_key = jax.random.split(rng_key) - class_representative_mask = global_class_representative_mask & _df_eval( # pytype: disable=unsupported-operands # typed-pandas + class_representative_mask = global_class_representative_mask & _df_eval( embeddings_df, eval_set_specification.class_representative_classwise_mask_fn( class_name @@ -490,7 +490,7 @@ def _prepare_eval_set( ) search_corpus_mask = ( - global_search_corpus_mask # pytype: disable=unsupported-operands # typed-pandas + global_search_corpus_mask & _df_eval( embeddings_df, eval_set_specification.search_corpus_classwise_mask_fn(class_name),