wehub-resource-sync
1084 lines
37 KiB
Python
1084 lines
37 KiB
Python
# Copyright 2024 Google LLC
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# https://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
# pylint: disable=too-many-lines
|
|
|
|
"""Utility functions and classes for the VAPO notebook."""
|
|
|
|
import csv
|
|
import io
|
|
import json
|
|
import random
|
|
import re
|
|
import string
|
|
import subprocess
|
|
from collections.abc import Callable
|
|
from typing import Any, Union
|
|
|
|
import ipywidgets as widgets
|
|
import jinja2
|
|
import jinja2.meta
|
|
import pandas as pd
|
|
import plotly.graph_objects as go
|
|
from IPython.core.display import DisplayHandle
|
|
from IPython.display import HTML, display
|
|
from google.cloud import aiplatform, storage
|
|
from jsonschema import ValidationError, validate
|
|
from tenacity import retry, wait_random_exponential
|
|
from tensorflow.io import gfile
|
|
from vertexai import generative_models
|
|
from vertexai.evaluation import EvalTask
|
|
from vertexai.generative_models import (
|
|
Content,
|
|
GenerationConfig,
|
|
GenerativeModel,
|
|
Part,
|
|
SafetySetting,
|
|
Tool,
|
|
ToolConfig,
|
|
)
|
|
|
|
|
|
def is_target_required_metric(eval_metric: str) -> bool:
|
|
"""Check if the metric requires the target label."""
|
|
return eval_metric in [
|
|
"bleu",
|
|
"exact_match",
|
|
"question_answering_correctness",
|
|
"rouge_1",
|
|
"rouge_2",
|
|
"rouge_l",
|
|
"rouge_l_sum",
|
|
"tool_call_valid",
|
|
"tool_name_match",
|
|
"tool_parameter_key_match",
|
|
"tool_parameter_kv_match",
|
|
]
|
|
|
|
|
|
def is_run_target_required(eval_metric_types: list[str], source_model: str) -> bool:
|
|
"""Check if the run requires the target label."""
|
|
if source_model:
|
|
return False
|
|
|
|
label_required = False
|
|
for metric in eval_metric_types:
|
|
label_required = label_required or is_target_required_metric(metric)
|
|
return label_required
|
|
|
|
|
|
_TARGET_KEY = "target"
|
|
|
|
|
|
def load_file_from_gcs(dataset: str) -> str:
|
|
"""Loads the file from GCS and returns it as a string."""
|
|
if dataset.startswith("gs://"):
|
|
with gfile.GFile(dataset, "r") as f:
|
|
return f.read()
|
|
else:
|
|
raise ValueError(
|
|
"Unsupported file location. Only GCS paths starting with 'gs://' are"
|
|
" supported."
|
|
)
|
|
|
|
|
|
def parse_jsonl(data_str: str) -> list[dict[str, str]]:
|
|
"""Parses the content of a JSONL file and returns a list of dictionaries."""
|
|
data = []
|
|
lines = data_str.splitlines()
|
|
for line in lines:
|
|
if line:
|
|
try:
|
|
data.append(json.loads(line))
|
|
except json.JSONDecodeError as e:
|
|
raise ValueError(
|
|
f"Error decoding JSON on line: {line}. Error: {e}"
|
|
) from e
|
|
return data
|
|
|
|
|
|
def parse_and_validate_csv(data_str: str) -> list[dict[str, str]]:
|
|
"""Parses and validates the content of a CSV file and returns a list of dictionaries."""
|
|
data = []
|
|
csv_reader = csv.reader(io.StringIO(data_str))
|
|
|
|
# Extract and validate headers
|
|
try:
|
|
headers = next(csv_reader)
|
|
if not headers:
|
|
raise ValueError("The CSV file has an empty or invalid header row.")
|
|
except StopIteration as e:
|
|
raise ValueError("The CSV file is empty.") from e
|
|
|
|
# Validate and process rows
|
|
for row_number, row in enumerate(csv_reader, start=2):
|
|
if len(row) != len(headers):
|
|
raise ValueError(
|
|
f"Row {row_number} has an inconsistent number of fields. "
|
|
f"Expected {len(headers)} fields but found {len(row)}."
|
|
)
|
|
# Create dictionary for each row using headers as keys
|
|
item = dict(zip(headers, row))
|
|
data.append(item)
|
|
|
|
return data
|
|
|
|
|
|
def load_dataset(dataset: str) -> list[dict[str, str]]:
|
|
"""Loads and parses the dataset based on its file type ('.jsonl' or '.csv')."""
|
|
# Load the file from GCS
|
|
data_str = load_file_from_gcs(dataset)
|
|
|
|
# Parse based on file type
|
|
if dataset.endswith(".jsonl"):
|
|
return parse_jsonl(data_str)
|
|
|
|
if dataset.endswith(".csv"):
|
|
return parse_and_validate_csv(data_str)
|
|
|
|
raise ValueError(
|
|
"Unsupported file type. Please provide a file with '.jsonl' or '.csv'"
|
|
" extension."
|
|
)
|
|
|
|
|
|
def validate_prompt_and_data(
|
|
template: str,
|
|
dataset_path: str,
|
|
placeholder_to_content: str,
|
|
label_enforced: bool,
|
|
) -> None:
|
|
"""Validates the prompt template and the dataset."""
|
|
data = load_dataset(dataset_path)
|
|
placeholder_to_content_json = json.loads(placeholder_to_content)
|
|
template = re.sub(r"(?<!{){(?!{)(\s*\w+\s*)(?<!})}(?!})", r"{{\1}}", template)
|
|
env = jinja2.Environment()
|
|
try:
|
|
parsed_content = env.parse(template)
|
|
except jinja2.exceptions.TemplateSyntaxError as e:
|
|
raise ValueError(f"Invalid template: {template}") from e
|
|
|
|
template_variables = jinja2.meta.find_undeclared_variables(parsed_content)
|
|
extra_keys = set()
|
|
for ex in data:
|
|
ex.update(placeholder_to_content_json)
|
|
missing_keys = [key for key in template_variables if key not in ex]
|
|
extra_keys.update([key for key in ex if key not in template_variables])
|
|
if label_enforced:
|
|
if _TARGET_KEY not in ex:
|
|
raise ValueError(
|
|
f"The example {ex} doesn't have a key corresponding to the target"
|
|
f" var: {_TARGET_KEY}"
|
|
)
|
|
if not ex[_TARGET_KEY]:
|
|
raise ValueError(f"The following example has an empty target: {ex}")
|
|
if missing_keys:
|
|
raise ValueError(
|
|
f"The example {ex} doesn't have a key corresponding to following"
|
|
f" template vars: {missing_keys}"
|
|
)
|
|
if extra_keys:
|
|
raise Warning(
|
|
"Warning: extra keys in the examples not used in the prompt template"
|
|
f" template {extra_keys}"
|
|
)
|
|
|
|
|
|
def run_custom_job(
|
|
display_name: str,
|
|
container_uri: str,
|
|
container_args: dict[str, str],
|
|
) -> str:
|
|
"""A sample to create custom jobs."""
|
|
worker_pool_specs = [
|
|
{
|
|
"replica_count": 1,
|
|
"container_spec": {
|
|
"image_uri": container_uri,
|
|
"args": [f"--{k}={v}" for k, v in container_args.items()],
|
|
},
|
|
"machine_spec": {
|
|
"machine_type": "n1-standard-4",
|
|
},
|
|
}
|
|
]
|
|
|
|
custom_job = aiplatform.CustomJob(
|
|
display_name=display_name,
|
|
worker_pool_specs=worker_pool_specs,
|
|
)
|
|
custom_job.submit()
|
|
return custom_job
|
|
|
|
|
|
def run_apd(config: dict[str, str], bucket_uri: str, display_name: str) -> str:
|
|
"""A function to the vertex prompt optimizer."""
|
|
print(f"\n\nJob display name: {display_name}")
|
|
version = "preview_v1_0"
|
|
container_uri = "us-docker.pkg.dev/vertex-ai-restricted/builtin-algorithm/apd"
|
|
config_path = f"{bucket_uri}/{display_name}/input_config.json"
|
|
|
|
with gfile.GFile(config_path, "w") as f:
|
|
json.dump(config, f)
|
|
|
|
aiplatform.init(
|
|
project=config["project"],
|
|
location=config["target_model_location"],
|
|
staging_bucket=f"{bucket_uri}/{display_name}",
|
|
)
|
|
|
|
return run_custom_job(
|
|
display_name=display_name,
|
|
container_uri=f"{container_uri}:{version}",
|
|
container_args={"config": config_path},
|
|
)
|
|
|
|
|
|
def update_best_display(
|
|
df: pd.DataFrame,
|
|
textarea: widgets.Textarea,
|
|
best_score_label: widgets.Label,
|
|
eval_metric: str,
|
|
) -> None:
|
|
"""Update the best prompt display."""
|
|
df["score"] = df[f"metrics.{eval_metric}/mean"]
|
|
|
|
best_template = df.loc[df["score"].argmax(), "prompt"]
|
|
best_score = df.loc[df["score"].argmax(), "score"]
|
|
original_score = df.loc[0, "score"]
|
|
|
|
def placeholder_llm() -> str:
|
|
return "{{llm()}}"
|
|
|
|
env = jinja2.Environment(loader=jinja2.BaseLoader())
|
|
env.globals["llm"] = placeholder_llm
|
|
|
|
best_template = best_template.replace("store('answer', llm())", "llm()")
|
|
textarea.value = best_template
|
|
improvement = best_score - original_score
|
|
no_improvement_str = "\nNo better template is found yet." if not improvement else ""
|
|
best_score_label.value = (
|
|
f"Score: {best_score} Improvement: {improvement: .3f} {no_improvement_str}"
|
|
)
|
|
|
|
|
|
def generate_dataframe(filename: str) -> pd.DataFrame:
|
|
"""Generates a pandas dataframe from a json file."""
|
|
if not gfile.exists(filename):
|
|
return pd.DataFrame()
|
|
|
|
with gfile.GFile(filename, "r") as f:
|
|
try:
|
|
data = json.load(f)
|
|
except json.JSONDecodeError:
|
|
return pd.DataFrame()
|
|
return pd.json_normalize(data)
|
|
|
|
|
|
def left_aligned_df_html(df: pd.DataFrame) -> HTML:
|
|
"""Displays a Pandas DataFrame in Colab with left-aligned values."""
|
|
# Convert to HTML table, but keep the HTML in a variable
|
|
html_table = df.to_html(index=False, classes="left-aligned")
|
|
|
|
# Add CSS styling to left-align table data cells and override default styles
|
|
styled_html = f"""
|
|
<style>
|
|
.left-aligned td, .left-aligned th {{ text-align: left !important; }}
|
|
</style>
|
|
{html_table}
|
|
"""
|
|
|
|
# Display the styled HTML table
|
|
return HTML(styled_html)
|
|
|
|
|
|
def extract_top_level_function_name(source_code: str) -> str | None:
|
|
"""Extract the top level function name from the source code."""
|
|
match = re.search(r"^def\s+([a-zA-Z_]\w*)\s*\(", source_code, re.MULTILINE)
|
|
if match:
|
|
return match.group(1)
|
|
return None
|
|
|
|
|
|
class ProgressForm:
|
|
"""A class to display the progress of the optimization job."""
|
|
|
|
# pylint: disable=too-many-instance-attributes
|
|
|
|
def __init__(self, params: dict[str, str]) -> None:
|
|
"""Initialize the progress form."""
|
|
self.instruction_progress_bar = None
|
|
self.instruction_display = None
|
|
self.instruction_best = None
|
|
self.instruction_score = None
|
|
self.demo_progress_bar = None
|
|
self.demo_display = None
|
|
self.demo_best = None
|
|
self.demo_score = None
|
|
|
|
self.job_state_display = display(
|
|
HTML("<span>Job State: Not Started!</span>"), display_id=True
|
|
)
|
|
self.status_display = display(HTML(""), display_id=True)
|
|
|
|
if params["optimization_mode"] in ["instruction", "instruction_and_demo"]:
|
|
(
|
|
self.instruction_progress_bar,
|
|
self.instruction_display,
|
|
self.instruction_best,
|
|
self.instruction_score,
|
|
) = self.create_progress_ui("Instruction", params["num_steps"])
|
|
|
|
if params["optimization_mode"] in ["demonstration", "instruction_and_demo"]:
|
|
(
|
|
self.demo_progress_bar,
|
|
self.demo_display,
|
|
self.demo_best,
|
|
self.demo_score,
|
|
) = self.create_progress_ui(
|
|
"Demonstration", params["num_demo_set_candidates"]
|
|
)
|
|
|
|
if len(params["eval_metrics_types"]) == 1:
|
|
self.eval_metric = params["eval_metrics_types"][0]
|
|
else:
|
|
self.eval_metric = "composite_metric"
|
|
|
|
self.output_path = params["output_path"]
|
|
self.instruction_df = None
|
|
self.demo_df = None
|
|
|
|
# pylint: disable=too-many-positional-arguments,too-many-arguments
|
|
def update_progress(
|
|
self,
|
|
progress_bar: widgets.IntProgress | None,
|
|
templates_file: str,
|
|
df: pd.DataFrame | None,
|
|
df_display: DisplayHandle,
|
|
best_textarea: widgets.Textarea,
|
|
best_score: widgets.Label,
|
|
eval_metric: str,
|
|
) -> pd.DataFrame:
|
|
"""Update the progress of the optimization job."""
|
|
|
|
def get_last_step(df: pd.DataFrame) -> int:
|
|
if df.empty:
|
|
return -1
|
|
return int(df["step"].max())
|
|
|
|
if progress_bar is None or df is None:
|
|
return pd.DataFrame()
|
|
|
|
new_df = generate_dataframe(templates_file)
|
|
|
|
last_step = get_last_step(df)
|
|
new_last_step = get_last_step(new_df)
|
|
if new_last_step > last_step:
|
|
df_display.update(left_aligned_df_html(new_df))
|
|
update_best_display(new_df, best_textarea, best_score, eval_metric)
|
|
progress_bar.value = progress_bar.value + new_last_step - last_step
|
|
|
|
return new_df
|
|
|
|
def create_progress_ui(
|
|
self, opt_mode: str, num_opt_steps: str
|
|
) -> tuple[widgets.IntProgress, DisplayHandle, widgets.Textarea, widgets.Label]:
|
|
"""Create the progress UI for a specific optimization mode."""
|
|
print(f"\n\n{opt_mode} Optimization")
|
|
progress_bar = widgets.IntProgress(
|
|
value=0, min=0, max=int(num_opt_steps), step=1, description="Progress"
|
|
)
|
|
display(progress_bar)
|
|
print("\nGenerated Templates:")
|
|
templates_display = display("No template is evaluated yet!", display_id=True)
|
|
|
|
print("\nBest Template so far:")
|
|
best_textarea = widgets.Textarea(
|
|
value="NA",
|
|
disabled=False,
|
|
layout=widgets.Layout(width="80%", height="150px"),
|
|
)
|
|
display(best_textarea)
|
|
|
|
best_score = widgets.Label(value="Score: NA Improvement: NA")
|
|
display(best_score)
|
|
|
|
return progress_bar, templates_display, best_textarea, best_score
|
|
|
|
def monitor_progress(self, job: aiplatform.CustomJob) -> bool:
|
|
"""Monitor the progress of the optimization job."""
|
|
self.job_state_display.update(HTML(f"<span>Job State: {job.state.name}</span>"))
|
|
|
|
# Initial display of the templates.
|
|
instruction_templates_file = f"{self.output_path}/instruction/templates.json"
|
|
demo_templates_file = f"{self.output_path}/demonstration/templates.json"
|
|
|
|
if not job.done():
|
|
self.instruction_df = self.update_progress(
|
|
self.instruction_progress_bar,
|
|
instruction_templates_file,
|
|
self.instruction_df,
|
|
self.instruction_display,
|
|
self.instruction_best,
|
|
self.instruction_score,
|
|
self.eval_metric,
|
|
)
|
|
self.demo_df = self.update_progress(
|
|
self.demo_progress_bar,
|
|
demo_templates_file,
|
|
self.demo_df,
|
|
self.demo_display,
|
|
self.demo_best,
|
|
self.demo_score,
|
|
self.eval_metric,
|
|
)
|
|
return True
|
|
|
|
if job.state.name != "JOB_STATE_SUCCEEDED":
|
|
errors = [f"Error: Job failed with error {job.error}."]
|
|
for err_file in [
|
|
f"{self.output_path}/instruction/error.json",
|
|
f"{self.output_path}/demonstration/error.json",
|
|
f"{self.output_path}/error.json",
|
|
]:
|
|
if gfile.exists(err_file):
|
|
with gfile.GFile(err_file, "r") as f:
|
|
error_json = json.load(f)
|
|
errors.append(f"Detailed error: {error_json['Error']}")
|
|
errors.append(
|
|
f"Please feel free to send {err_file} to the VAPO team to help"
|
|
" resolving the issue."
|
|
)
|
|
|
|
errors.append(
|
|
"All the templates found before failure can be found under"
|
|
f" {self.output_path}"
|
|
)
|
|
errors.append(
|
|
"Please consider rerunning to make sure the failure is intransient."
|
|
)
|
|
err = "\n".join(errors)
|
|
err = err.replace("\n", "<br>")
|
|
self.status_display.update(HTML(f'<span style="color: red;">{err}</span>'))
|
|
else:
|
|
self.status_display.update(
|
|
HTML(
|
|
'<span style="color: green;">Job succeeded!</span> <span>All the'
|
|
f" artifacts can be found under {self.output_path}</span>"
|
|
)
|
|
)
|
|
return False
|
|
|
|
|
|
def display_dataframe(df: pd.DataFrame) -> None:
|
|
"""Display a pandas dataframe in Colab."""
|
|
|
|
# Function to wrap text in a scrollable div
|
|
def wrap_in_scrollable_div(text: str) -> str:
|
|
return f'<div class="scrollable">{text}</div>'
|
|
|
|
# Apply the function to every cell using the format method
|
|
styled_html = df.style.format(wrap_in_scrollable_div).to_html(index=False)
|
|
|
|
# Display the HTML in the notebook
|
|
display(HTML(styled_html))
|
|
|
|
|
|
def split_gcs_path(gcs_path: str) -> tuple[str, str]:
|
|
"""Splits a full GCS path into bucket name and prefix."""
|
|
if gcs_path.startswith("gs://"):
|
|
path_without_scheme = gcs_path[5:] # Remove the 'gs://' part
|
|
parts = path_without_scheme.split("/", 1)
|
|
bucket_name = parts[0]
|
|
prefix = parts[1] if len(parts) > 1 else ""
|
|
return bucket_name, prefix
|
|
|
|
raise ValueError("Invalid GCS path. Must start with 'gs://'")
|
|
|
|
|
|
def list_gcs_objects(full_gcs_path: str) -> list[str]:
|
|
"""Lists all the objects in the given GCS path."""
|
|
bucket_name, prefix = split_gcs_path(full_gcs_path)
|
|
storage_client = storage.Client()
|
|
bucket = storage_client.bucket(bucket_name)
|
|
blobs = bucket.list_blobs(
|
|
prefix=prefix
|
|
) # List all objects that start with the prefix
|
|
|
|
return [blob.name for blob in blobs]
|
|
|
|
|
|
def find_directories_with_files(
|
|
full_gcs_path: str, required_files: list[str]
|
|
) -> list[str]:
|
|
"""Finds directories containing specific files under the given full GCS path."""
|
|
bucket_name, prefix = split_gcs_path(full_gcs_path)
|
|
all_paths = list_gcs_objects(f"gs://{bucket_name}/{prefix}")
|
|
directories = set()
|
|
|
|
# Create a dictionary to track files found in each directory
|
|
file_presence: dict[str, set[str]] = {}
|
|
for path in all_paths:
|
|
# Get the directory part of the path
|
|
directory = "/".join(path.split("/")[:-1])
|
|
filename = path.split("/")[-1] # Get the filename part of the path
|
|
if directory:
|
|
if directory not in file_presence:
|
|
file_presence[directory] = set()
|
|
file_presence[directory].add(filename)
|
|
|
|
# Check which directories have all required files
|
|
for directory, files in file_presence.items():
|
|
if all(file in files for file in required_files):
|
|
directories.add(f"gs://{bucket_name}/{directory}")
|
|
|
|
return list(directories)
|
|
|
|
|
|
def extract_metric_name(metric_string: str) -> str:
|
|
"""Extract the metric name from a string."""
|
|
# Use a regular expression to find the metric name
|
|
match = re.search(r"\.(\w+)/", metric_string)
|
|
# Return the matched group if found
|
|
return match.group(1) if match else metric_string
|
|
|
|
|
|
def read_file_from_gcs(filename: str) -> str:
|
|
"""Read a file from GCS."""
|
|
with gfile.GFile(filename, "r") as f:
|
|
return f.read()
|
|
|
|
|
|
def process_results(df: pd.DataFrame) -> pd.DataFrame:
|
|
"""Process the results removing columns that could be confusing."""
|
|
columns_to_drop = []
|
|
# Dropping columns that could be confusing.
|
|
for col in df.columns:
|
|
if "confidence" in col:
|
|
columns_to_drop.append(col)
|
|
if "raw_eval_resp" in col:
|
|
columns_to_drop.append(col)
|
|
if col == "instruction":
|
|
columns_to_drop.append(col)
|
|
if col == "context":
|
|
columns_to_drop.append(col)
|
|
return df.drop(columns=columns_to_drop)
|
|
|
|
|
|
class ResultsUI:
|
|
"""A UI to display the results of a VAPO run."""
|
|
|
|
def __init__(self, path: str) -> None:
|
|
"""Initialize the UI."""
|
|
required_files = ["eval_results.json", "templates.json"]
|
|
runs = find_directories_with_files(path, required_files)
|
|
|
|
self.run_label = widgets.Label("Select Run:")
|
|
|
|
if not runs:
|
|
self.run_dropdown = widgets.Dropdown(
|
|
options=[],
|
|
value=None,
|
|
layout=widgets.Layout(width="200px"),
|
|
disabled=True,
|
|
)
|
|
self.dropdown_description = widgets.Label("Select Template:")
|
|
self.template_dropdown = widgets.Dropdown(
|
|
options=[],
|
|
value=None,
|
|
layout=widgets.Layout(width="400px"),
|
|
disabled=True,
|
|
)
|
|
self.results_output = widgets.Output(
|
|
layout=widgets.Layout(
|
|
height="600px", overflow="auto", margin="20px 0px 0px 0px"
|
|
)
|
|
)
|
|
self.templates = []
|
|
self.eval_results = []
|
|
return
|
|
|
|
self.run_dropdown = widgets.Dropdown(
|
|
options=runs, value=runs[0], layout=widgets.Layout(width="200px")
|
|
)
|
|
self.run_dropdown.observe(self.display_run_handler, names="value")
|
|
|
|
# Create a label widget for the description
|
|
self.dropdown_description = widgets.Label("Select Template:")
|
|
self.template_dropdown = widgets.Dropdown(
|
|
options=[],
|
|
value=None,
|
|
layout=widgets.Layout(width="400px"),
|
|
disabled=True,
|
|
)
|
|
self.template_dropdown.observe(self.display_template_handler, names="value")
|
|
self.results_output = widgets.Output(
|
|
layout=widgets.Layout(
|
|
height="600px", overflow="auto", margin="20px 0px 0px 0px"
|
|
)
|
|
)
|
|
self.display_run(runs[0])
|
|
|
|
def display_template_handler(self, change: dict[str, str | None]) -> None:
|
|
"""Display the template and the corresponding evaluation results."""
|
|
if change["new"] is None:
|
|
return
|
|
|
|
df_index = int(change["new"].split(" ")[1])
|
|
self.display_eval_results(df_index)
|
|
|
|
def display_run_handler(self, change: dict[str, str | None]) -> None:
|
|
"""Display the run and the corresponding templates."""
|
|
if change["new"] is None:
|
|
return
|
|
|
|
path = change["new"]
|
|
self.display_run(path)
|
|
|
|
def display_run(self, path: str) -> None:
|
|
"""Display the results of a VAPO run."""
|
|
self.run_dropdown.disabled = True
|
|
filename = f"{path}/eval_results.json"
|
|
eval_results = json.loads(read_file_from_gcs(filename))
|
|
|
|
filename = f"{path}/templates.json"
|
|
templates = json.loads(read_file_from_gcs(filename))
|
|
|
|
if len(templates) == len(eval_results):
|
|
offset = 0
|
|
elif len(templates) == len(eval_results) + 1:
|
|
# In some setups it is possible to have 1 more template than results.
|
|
offset = 1
|
|
else:
|
|
raise ValueError(
|
|
"Number of templates doesn't match number of eval results"
|
|
f" {len(templates)} vs {len(eval_results)}"
|
|
)
|
|
self.templates = [
|
|
pd.json_normalize(template) for template in templates[offset:]
|
|
]
|
|
metric_columns = [col for col in self.templates[0].columns if "metric" in col]
|
|
|
|
self.eval_results = [
|
|
process_results(pd.read_json(io.StringIO(result["metrics_table"])))
|
|
for result in eval_results
|
|
]
|
|
options = []
|
|
for i, template in enumerate(self.templates):
|
|
metrics = []
|
|
for col in metric_columns:
|
|
value = template[col].tolist()[0]
|
|
short_col = extract_metric_name(col)
|
|
metrics.append(f"{short_col}: {value}")
|
|
metrics_str = " ".join(metrics)
|
|
options.append(f"Template {i} {metrics_str}")
|
|
|
|
self.template_dropdown.disabled = False
|
|
self.template_dropdown.options = options
|
|
self.run_dropdown.disabled = False
|
|
|
|
def display_eval_results(self, index: int) -> None:
|
|
"""Display the evaluation results for a specific template."""
|
|
with self.results_output:
|
|
self.results_output.clear_output(wait=True) # Clear previous output
|
|
display_dataframe(self.templates[index])
|
|
print()
|
|
display_dataframe(self.eval_results[index])
|
|
|
|
def get_container(self) -> widgets.Output:
|
|
"""Get the container widget for the results UI."""
|
|
return widgets.VBox(
|
|
[
|
|
self.run_label,
|
|
self.run_dropdown,
|
|
self.dropdown_description,
|
|
self.template_dropdown,
|
|
self.results_output,
|
|
]
|
|
)
|
|
|
|
|
|
def get_id(length: int = 8) -> str:
|
|
"""Generate a uuid of a specified length (default=8)."""
|
|
return "".join(random.choices(string.ascii_lowercase + string.digits, k=length))
|
|
|
|
|
|
def get_auth_token() -> str:
|
|
"""A function to collect the authorization token"""
|
|
result = subprocess.run(
|
|
["gcloud", "auth", "print-identity-token", "-q"],
|
|
capture_output=True,
|
|
text=True,
|
|
check=True,
|
|
)
|
|
return result.stdout.strip()
|
|
|
|
|
|
def init_new_model(
|
|
model_name: str,
|
|
generation_config: GenerationConfig | None = None,
|
|
safety_settings: list[SafetySetting] | None = None,
|
|
**kwargs: Any,
|
|
) -> GenerativeModel:
|
|
"""Initialize a new model with configurable generation and safety settings."""
|
|
if generation_config is None:
|
|
generation_config = GenerationConfig(
|
|
candidate_count=1, max_output_tokens=2048, temperature=0
|
|
)
|
|
if safety_settings is None:
|
|
safety_settings = [
|
|
generative_models.SafetySetting(
|
|
category=generative_models.HarmCategory.HARM_CATEGORY_HATE_SPEECH,
|
|
method=generative_models.SafetySetting.HarmBlockMethod.SEVERITY,
|
|
threshold=generative_models.HarmBlockThreshold.BLOCK_NONE,
|
|
),
|
|
generative_models.SafetySetting(
|
|
category=generative_models.HarmCategory.HARM_CATEGORY_DANGEROUS_CONTENT,
|
|
method=generative_models.SafetySetting.HarmBlockMethod.SEVERITY,
|
|
threshold=generative_models.HarmBlockThreshold.BLOCK_NONE,
|
|
),
|
|
generative_models.SafetySetting(
|
|
category=generative_models.HarmCategory.HARM_CATEGORY_SEXUALLY_EXPLICIT,
|
|
method=generative_models.SafetySetting.HarmBlockMethod.SEVERITY,
|
|
threshold=generative_models.HarmBlockThreshold.BLOCK_NONE,
|
|
),
|
|
generative_models.SafetySetting(
|
|
category=generative_models.HarmCategory.HARM_CATEGORY_HARASSMENT,
|
|
method=generative_models.SafetySetting.HarmBlockMethod.SEVERITY,
|
|
threshold=generative_models.HarmBlockThreshold.BLOCK_NONE,
|
|
),
|
|
]
|
|
|
|
model = GenerativeModel(
|
|
model_name=model_name,
|
|
generation_config=generation_config,
|
|
safety_settings=safety_settings,
|
|
**kwargs,
|
|
)
|
|
return model
|
|
|
|
|
|
@retry(wait=wait_random_exponential(multiplier=1, max=120))
|
|
async def async_generate(
|
|
prompt: str,
|
|
model: GenerativeModel,
|
|
function_handler: dict[str, Callable] | None = None,
|
|
tools: Tool | None = None,
|
|
tool_config: ToolConfig | None = None,
|
|
**kwargs: Any,
|
|
) -> Union[str, None]:
|
|
"""Generates a response from the model, optionally handling function calls."""
|
|
user_prompt_content = Content(role="user", parts=[Part.from_text(prompt)])
|
|
|
|
try:
|
|
# Initial generation - potentially calling a function.
|
|
response = await model.generate_content_async(
|
|
prompt,
|
|
tools=[tools] if tools else None, # Only provide tools if they exist
|
|
tool_config=tool_config if tool_config else None, # Same for tool_config
|
|
**kwargs,
|
|
)
|
|
|
|
# Handle function calls if applicable
|
|
if (
|
|
function_handler
|
|
and response
|
|
and response.candidates
|
|
and response.candidates[0].content.parts[0].function_call
|
|
):
|
|
while response.candidates[0].content.parts[0].function_call:
|
|
function_call = response.candidates[0].content.parts[0].function_call
|
|
function_name = function_call.name
|
|
|
|
if function_name in function_handler:
|
|
function_args = function_call.args # No need for manual conversion
|
|
api_response = function_handler[function_name](function_args)
|
|
|
|
response = await model.generate_content_async(
|
|
[
|
|
user_prompt_content,
|
|
response.candidates[0].content,
|
|
Content(
|
|
parts=[
|
|
Part.from_function_response(
|
|
name=function_name,
|
|
response={"content": api_response},
|
|
)
|
|
]
|
|
),
|
|
],
|
|
tools=[tools] if tools else None, # Conditional tool passing
|
|
tool_config=tool_config if tool_config else None,
|
|
)
|
|
else:
|
|
break # Exit loop if function not found
|
|
|
|
# Extract and return text if generation was successful
|
|
if response and response.candidates and response.candidates[0].content.parts:
|
|
return (
|
|
response.candidates[0].content.parts[0].text
|
|
) # More robust text extraction
|
|
return None
|
|
|
|
except Exception as e: # pylint: disable=broad-except
|
|
print(f"Error calling the model: {e}") # Include the actual error message
|
|
return "Could not call the model. Please try it again in a few minutes."
|
|
|
|
|
|
# pylint: disable=too-many-positional-arguments,too-many-arguments
|
|
def evaluate_task(
|
|
df: pd.DataFrame,
|
|
prompt_col: str,
|
|
reference_col: str,
|
|
response_col: str,
|
|
experiment_name: str,
|
|
eval_metrics: list[str],
|
|
eval_sample_n: int,
|
|
) -> dict[str, float]:
|
|
"""Evaluate task using Agent Platform Evaluation."""
|
|
# Generate a unique id for the experiment run
|
|
idx = get_id()
|
|
|
|
# Rename the columns to match the expected format
|
|
eval_dataset = df[[prompt_col, reference_col, response_col]].rename(
|
|
columns={
|
|
prompt_col: "prompt",
|
|
reference_col: "reference",
|
|
response_col: "response",
|
|
}
|
|
)
|
|
|
|
# Drop rows with missing values
|
|
eval_dataset = eval_dataset.dropna()
|
|
|
|
# Sample a subset of the dataset
|
|
eval_dataset = eval_dataset.sample(n=eval_sample_n, random_state=8).reset_index(
|
|
drop=True
|
|
)
|
|
|
|
# Create an EvalTask object
|
|
eval_task = EvalTask(
|
|
dataset=eval_dataset,
|
|
metrics=eval_metrics,
|
|
experiment=experiment_name,
|
|
)
|
|
|
|
# Evaluate the task
|
|
result = eval_task.evaluate(experiment_run_name=f"{experiment_name}-{idx}")
|
|
|
|
# Return the summary metrics
|
|
return result.summary_metrics
|
|
|
|
|
|
def print_df_rows(
|
|
df: pd.DataFrame, columns: list[str] | None = None, n: int = 3
|
|
) -> None:
|
|
"""Print a subset of rows from a DataFrame."""
|
|
# Apply column filtering if specified
|
|
if columns:
|
|
df = df[columns]
|
|
|
|
# Style definitions for improved readability
|
|
base_style = (
|
|
"font-family: monospace; font-size: 14px; white-space: pre-wrap; width:"
|
|
" auto; overflow-x: auto;"
|
|
)
|
|
header_style = base_style + "font-weight: bold;"
|
|
|
|
# Iterate through the specified number of rows
|
|
for _, row in df.head(n).iterrows():
|
|
# Display each column name as a bold header
|
|
for column in df.columns:
|
|
display(
|
|
HTML(
|
|
"<span"
|
|
f" style='{header_style}'>{column.replace('_', ' ').title()}:"
|
|
" </span>"
|
|
)
|
|
)
|
|
display(
|
|
HTML(f"<span style='{base_style}'>{row[column]}</span><br>")
|
|
) # Display value and line break
|
|
display(HTML("<hr>")) # Add separator between rows for clarity
|
|
|
|
|
|
def plot_eval_metrics(
|
|
eval_results: list[tuple[str, dict[str, float]]],
|
|
metrics: list[str] | None = None,
|
|
) -> None:
|
|
"""Plot a bar plot for the evaluation results."""
|
|
# Create data for the bar plot
|
|
data = []
|
|
for eval_result in eval_results:
|
|
title, summary_metrics = eval_result
|
|
if metrics:
|
|
summary_metrics = {
|
|
k: v
|
|
for k, v in summary_metrics.items()
|
|
if any(selected_metric in k for selected_metric in metrics)
|
|
}
|
|
|
|
summary_metrics = {k: v for k, v in summary_metrics.items() if "mean" in k}
|
|
data.append(
|
|
go.Bar(
|
|
x=list(summary_metrics.keys()),
|
|
y=list(summary_metrics.values()),
|
|
name=title,
|
|
)
|
|
)
|
|
|
|
# Update the figure with the data
|
|
fig = go.Figure(data=data)
|
|
|
|
# Add the title
|
|
fig.update_layout(
|
|
title=go.layout.Title(text="Evaluation Metrics", x=0.5),
|
|
xaxis_title="Metric Name",
|
|
yaxis_title="Mean Value",
|
|
)
|
|
|
|
# Change the bar mode
|
|
fig.update_layout(barmode="group")
|
|
|
|
# Show the plot
|
|
fig.show()
|
|
|
|
|
|
def create_target_column(row: dict[str, Any]) -> str:
|
|
"""Creates a JSON string representing tool calls from input row."""
|
|
tool_calls = (
|
|
[{"name": row["tool_names"], "arguments": row["tool_arguments"]}]
|
|
if row.get("tool_names")
|
|
else []
|
|
)
|
|
|
|
return json.dumps({"content": "", "tool_calls": tool_calls})
|
|
|
|
|
|
def tool_config_to_dict(tool_config: ToolConfig | None) -> dict[str, Any] | None:
|
|
"""Converts a ToolConfig object to a dictionary."""
|
|
if tool_config is None:
|
|
return None
|
|
|
|
# pylint: disable=protected-access
|
|
config = tool_config._gapic_tool_config.function_calling_config
|
|
return {
|
|
"function_calling_config": {
|
|
"mode": config.mode.name,
|
|
"allowed_function_names": list(config.allowed_function_names),
|
|
}
|
|
}
|
|
|
|
|
|
def replace_type_key(data: dict[str, Any]) -> dict[str, Any]:
|
|
"""Recursively replaces "type_" with "type" in a dictionary or list."""
|
|
|
|
def _recursive_replace(item: Any) -> Any:
|
|
if isinstance(item, dict):
|
|
return {
|
|
("type" if k == "type_" else k): _recursive_replace(v)
|
|
for k, v in item.items()
|
|
}
|
|
if isinstance(item, list):
|
|
return [_recursive_replace(elem) for elem in item]
|
|
return item
|
|
|
|
new_data = {}
|
|
for key, value in data.items():
|
|
if key == "function_declarations" and isinstance(value, list):
|
|
new_data[key] = [_recursive_replace(tool) for tool in value]
|
|
else:
|
|
new_data[key] = value
|
|
|
|
return new_data
|
|
|
|
|
|
def validate_tools(spec: str) -> None:
|
|
"""Validates the tools specification."""
|
|
# Define the JSON schema for validation
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"tools": {
|
|
"type": "array",
|
|
"minItems": 1, # Ensures that 'tools' is not an empty array
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"function_declarations": {
|
|
"type": "array",
|
|
# Ensures this is not an empty array
|
|
"minItems": 1,
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"name": {"type": "string"},
|
|
"description": {"type": "string"},
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"type": {"type": "string"},
|
|
"properties": {"type": "object"},
|
|
"required": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
},
|
|
},
|
|
"required": ["type", "properties"],
|
|
},
|
|
},
|
|
"required": ["name", "description", "parameters"],
|
|
},
|
|
}
|
|
},
|
|
"required": ["function_declarations"],
|
|
},
|
|
}
|
|
},
|
|
"required": ["tools"],
|
|
}
|
|
|
|
json_spec = json.loads(spec)
|
|
try:
|
|
# Validate the JSON specification against the schema
|
|
validate(instance=json_spec, schema=schema)
|
|
except ValidationError as e:
|
|
raise ValueError(f"Invalid Tools specification: {e}") from e
|
|
|
|
|
|
def validate_tool_config(tool_config: str) -> None:
|
|
"""Validates the format of the tool_config."""
|
|
schema = {
|
|
"type": "object",
|
|
"properties": {
|
|
"function_calling_config": {
|
|
"type": "object",
|
|
"properties": {
|
|
"mode": {"type": "string", "enum": ["AUTO", "ANY", "NONE"]},
|
|
"allowed_function_names": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
},
|
|
},
|
|
"required": ["mode"],
|
|
}
|
|
},
|
|
"required": ["function_calling_config"],
|
|
}
|
|
|
|
try:
|
|
validate(instance=json.loads(tool_config), schema=schema)
|
|
except ValidationError as e:
|
|
raise ValueError(f"Invalid tool_config: {tool_config}") from e
|