feat: view model evaluations when AustoSxS pipeline completes

PiperOrigin-RevId: 642441033
googleapis · Jun 12, 2024 · 17c59c4 · 17c59c4
1 parent 3481f40
commit 17c59c4
Showing 1 changed file with 21 additions and 0 deletions.
diff --git a/google/cloud/aiplatform/pipeline_jobs.py b/google/cloud/aiplatform/pipeline_jobs.py
@@ -373,6 +373,27 @@ def _run(
 
         self._block_until_complete()
 
+        # AutoSxS view model evaluations
+        for details in self.task_details:
+            if details.task_name == "model-evaluation-text-generation-pairwise":
+                model_a_eval = details.execution.metadata[
+                    "output:model_a_evaluation_path"
+                ]
+                model_b_eval = details.execution.metadata[
+                    "output:model_b_evaluation_path"
+                ]
+                if model_a_eval:
+                    _LOGGER.info("Model A")
+                    utils._ipython_utils.display_model_evaluation_button(
+                        aiplatform.ModelEvaluation(model_a_eval),
+                    )
+                if model_b_eval:
+                    _LOGGER.info("Model B")
+                    utils._ipython_utils.display_model_evaluation_button(
+                        aiplatform.ModelEvaluation(model_b_eval),
+                    )
+                break
+
     def submit(
         self,
         service_account: Optional[str] = None,