Allow None values in ModelCosineSimilarity evaluator.

PiperOrigin-RevId: 645424050
tensorflow · Jun 21, 2024 · 807bd62 · 807bd62
1 parent 84a9c2b
commit 807bd62
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/tensorflow_model_analysis/metrics/model_cosine_similarity.py b/tensorflow_model_analysis/metrics/model_cosine_similarity.py
@@ -148,7 +148,7 @@ def add_input(
             model_name=self._baseline_model_name,
             output_name=self._output_name,
             flatten=False,
-            allow_none=False,
+            allow_none=True,
         )
     )
 
@@ -159,7 +159,7 @@ def add_input(
             model_name=self._model_name,
             output_name=self._output_name,
             flatten=False,
-            allow_none=False,
+            allow_none=True,
         )
     )
     accumulator.merge(

diff --git a/tensorflow_model_analysis/metrics/model_cosine_similarity_test.py b/tensorflow_model_analysis/metrics/model_cosine_similarity_test.py
@@ -101,7 +101,10 @@ def test_cosine_similarity(
     examples = []
     for baseline_prediction, candidate_prediction in prediction_pairs:
       examples.append({
-          constants.LABELS_KEY: [0],
+          constants.LABELS_KEY: {
+              baseline_model_name: None,
+              candidate_model_name: None,
+          },
           constants.PREDICTIONS_KEY: {
               baseline_model_name: baseline_prediction,
               candidate_model_name: candidate_prediction,