feat: Allow EvalTask to take dataset as a Google sheet - simplifying rubric revision CUJ

PiperOrigin-RevId: 744088634

Author: vertex-sdk-bot

vertexai/preview/evaluation/eval_task.py

@@ -17,9 +17,9 @@
 """Evaluation Task class."""
 
 import logging
-import warnings
 from typing import Any, Callable, Dict, List, Literal, Optional, TYPE_CHECKING, Union
 import uuid
+import warnings
 
 from google.api_core import exceptions
 import vertexai
@@ -45,6 +45,7 @@
 
 if TYPE_CHECKING:
     import pandas as pd
+    from google.colab import sheets
 
 
 # pylint: disable=g-import-not-at-top
@@ -241,7 +242,7 @@ def custom_model_fn(input: str) -> str:
     def __init__(
         self,
         *,
-        dataset: Union["pd.DataFrame", str, Dict[str, Any]],
+        dataset: Union["pd.DataFrame", str, Dict[str, Any], "sheets.InteractiveSheet"],
         metrics: List[
             Union[
                 Literal[

vertexai/preview/evaluation/utils.py

@@ -20,20 +20,21 @@
 import io
 import json
 import os
+import re
+import sys
 import tempfile
 import threading
 import time
-import re
 from typing import (
     Any,
     Callable,
     Dict,
     List,
     Optional,
-    TYPE_CHECKING,
-    Union,
     Pattern,
     Tuple,
+    TYPE_CHECKING,
+    Union,
 )
 
 from google.cloud import bigquery
@@ -219,6 +220,12 @@ def load_dataset(
             'Pandas is not installed. Please install the SDK using "pip install'
             ' google-cloud-aiplatform[evaluation]"'
         )
+    if "google.colab" in sys.modules:
+        from google.colab import sheets
+
+        if isinstance(source, sheets.InteractiveSheet):
+            return source.as_df().copy()
+
     if isinstance(source, pd.DataFrame):
         return source.copy()
     elif isinstance(source, dict):