feat: Copy output uploading features from GA Eval SDK to preview.

PiperOrigin-RevId: 740001047

Author: jsondai

vertexai/preview/evaluation/utils.py

@@ -18,6 +18,7 @@
 
 import functools
 import io
+import json
 import os
 import tempfile
 import threading
@@ -30,10 +31,15 @@
 from google.cloud.aiplatform import compat
 from google.cloud.aiplatform import initializer
 from google.cloud.aiplatform import utils
+from google.cloud.aiplatform.utils import _ipython_utils
+from vertexai.evaluation import _base as eval_base
 from google.cloud.aiplatform_v1beta1.services import (
     evaluation_service as gapic_evaluation_services,
 )
-
+from vertexai.evaluation.metrics import (
+    _base as metrics_base,
+    metric_prompt_template as metric_prompt_template_base,
+)
 
 if TYPE_CHECKING:
     import pandas as pd
@@ -255,6 +261,16 @@ def _read_gcs_file_contents(filepath: str) -> str:
     return blob.download_as_string().decode("utf-8")
 
 
+def _upload_file_to_gcs(upload_gcs_path: str, filename: str) -> None:
+    storage_client = storage.Client(
+        project=initializer.global_config.project,
+        credentials=initializer.global_config.credentials,
+    )
+    storage.Blob.from_string(
+        uri=upload_gcs_path, client=storage_client
+    ).upload_from_filename(filename)
+
+
 def _upload_pandas_df_to_gcs(
     df: "pd.DataFrame", upload_gcs_path: str, file_type: str
 ) -> None:
@@ -288,28 +304,94 @@ def _upload_pandas_df_to_gcs(
         ).upload_from_filename(filename=local_dataset_path)
 
 
+def _upload_evaluation_summary_to_gcs(
+    summary_metrics: Dict[str, float],
+    upload_gcs_path: str,
+    candidate_model_name: Optional[str] = None,
+    baseline_model_name: Optional[str] = None,
+    dataset_uri: Optional[str] = None,
+    metrics: Optional[List[Union[str, metrics_base._Metric]]] = None,
+) -> None:
+    """Uploads the evaluation summary to a GCS bucket."""
+    summary = {
+        "summary_metrics": summary_metrics,
+    }
+    if candidate_model_name:
+        summary["candidate_model_name"] = candidate_model_name
+    if baseline_model_name:
+        summary["baseline_model_name"] = baseline_model_name
+    if dataset_uri:
+        summary["dataset_uri"] = dataset_uri
+
+    if metrics:
+        metric_descriptions = {}
+        for metric in metrics:
+            if isinstance(metric, metrics_base._ModelBasedMetric) and isinstance(
+                metric._raw_metric_prompt_template,
+                metric_prompt_template_base._MetricPromptTemplate,
+            ):
+                metric_descriptions[metric.metric_name] = {
+                    "criteria": metric._raw_metric_prompt_template._criteria,
+                    "rating_rubric": metric._raw_metric_prompt_template._rating_rubric,
+                }
+        summary["metric_descriptions"] = metric_descriptions
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        local_summary_path = os.path.join(temp_dir, "summary_metrics.json")
+        json.dump(summary, open(local_summary_path, "w"))
+        _upload_file_to_gcs(upload_gcs_path, local_summary_path)
+
+
 def upload_evaluation_results(
-    dataset: "pd.DataFrame", destination_uri_prefix: str, file_name: str
+    eval_result: eval_base.EvalResult,
+    destination_uri_prefix: str,
+    file_name: Optional[str] = None,
+    candidate_model_name: Optional[str] = None,
+    baseline_model_name: Optional[str] = None,
+    dataset_uri: Optional[str] = None,
+    metrics: Optional[List[Union[str, metrics_base._Metric]]] = None,
 ) -> None:
     """Uploads eval results to GCS destination.
 
     Args:
-        dataset: Pandas dataframe to upload.
+        eval_result: Eval results to upload.
         destination_uri_prefix: GCS folder to store the data.
-        file_name: File name to store the data.
+        file_name: Optional. File name to store the metrics table.
+        candidate_model_name: Optional. Candidate model name.
+        baseline_model_name: Optional. Baseline model name.
+        dataset_uri: Optional. URI pointing to the dataset.
+        metrics: Optional. List of metrics used for evaluation.
     """
-
     if not destination_uri_prefix:
+        _ipython_utils.display_gen_ai_evaluation_results_button()
+        return
+    if eval_result.metrics_table is None:
         return
     if destination_uri_prefix.startswith(_GCS_PREFIX):
-        _, extension = os.path.splitext(file_name)
+        if not file_name:
+            file_name = f"eval_results_{utils.timestamped_unique_name()}.csv"
+        base_name, extension = os.path.splitext(file_name)
         file_type = extension.lower()[1:]
-        output_path = destination_uri_prefix + "/" + file_name
-        _upload_pandas_df_to_gcs(dataset, output_path, file_type)
+        output_folder = destination_uri_prefix + "/" + base_name
+        metrics_table_path = output_folder + "/" + file_name
+        _upload_pandas_df_to_gcs(
+            eval_result.metrics_table, metrics_table_path, file_type
+        )
+        _upload_evaluation_summary_to_gcs(
+            eval_result.summary_metrics,
+            output_folder + "/summary_metrics.json",
+            candidate_model_name,
+            baseline_model_name,
+            dataset_uri,
+            metrics,
+        )
+        _ipython_utils.display_gen_ai_evaluation_results_button(
+            metrics_table_path.split(_GCS_PREFIX)[1]
+        )
     else:
         raise ValueError(
             f"Unsupported destination URI: {destination_uri_prefix}."
-            " Please provide a valid GCS bucket URI prefix starting with"
+            f" Please provide a valid GCS bucket URI prefix starting with"
             f" {_GCS_PREFIX}."
         )