![\"Google](\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\"){kind=logo}
Open in Colab\n", + " \n", + "
![\"Google](\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\")
Open in Colab Enterprise\n", + " \n", + "
![\"Vertex](\"https://www.gstatic.com/images/branding/gcpiconscolors/vertexai/v1/32px.svg\")
Open in Vertex AI Workbench\n", + " \n", + "
![\"GitHub](\"https://www.svgrepo.com/download/217753/github.svg\")
View on GitHub\n", + " \n", + "
Case {st.session_state.current_index + 1} of {len(st.session_state.human_rated_dict['user_input'])}
", + unsafe_allow_html=True, + ) + + st.markdown("---") + st.subheader("Launch Eval") + + st.subheader("Metrics Selection") + + col1, col2 = st.columns(2) + + with col1: + st.write("**Model-Based**") + metric_names = MetricPromptTemplateExamples.list_example_metric_names() + selected_model_based_metrics = st.multiselect( + "Select from model-based metrics", + metric_names, + key="selected_model_based_metrics", + label_visibility="collapsed", + ) + if selected_model_based_metrics: + num_metrics = len(selected_model_based_metrics) + if st.session_state.metric_preview_index >= num_metrics: + st.session_state.metric_preview_index = 0 + + current_metric_name = selected_model_based_metrics[ + st.session_state.metric_preview_index + ] + + st.markdown( + f"**Previewing Template: {current_metric_name} ({st.session_state.metric_preview_index + 1}/{num_metrics})**" + ) + + try: + metric_object = get_metric_object_by_name(current_metric_name) + if isinstance( + metric_object, + PointwiseMetricPromptTemplate | PairwiseMetricPromptTemplate, + ): + st.text_area( + "Template Preview", + metric_object.metric_prompt_template, + height=200, + ) + except Exception as e: + st.error( + f"Could not retrieve template for {current_metric_name}: {e}" + ) + + if num_metrics > 1: + prev_col, next_col = st.columns(2) + with prev_col: + if st.button( + "Previous Template", + disabled=st.session_state.metric_preview_index <= 0, + ): + st.session_state.metric_preview_index -= 1 + st.rerun() + with next_col: + if st.button( + "Next Template", + disabled=st.session_state.metric_preview_index + >= num_metrics - 1, + ): + st.session_state.metric_preview_index += 1 + st.rerun() + + with col2: + st.write("**Computation-Based Pointwise**") + computation_based_pointwise = [ + "bleu", + "rouge_1", + "rouge_2", + "rouge_l", + "rouge_l_sum", + "exact_match", + ] + st.multiselect( + "Select from computation-based pointwise metrics", + computation_based_pointwise, + key="selected_cbp", + label_visibility="collapsed", + ) + + st.selectbox( + "Select Evaluation Model", + options=[ + "gemini-2.0-flash-lite", + "gemini-2.5-flash", + "gemini-2.5-pro", + ], + key="selected_evaluation_model", + ) + + st.button("Launch Eval", key="launch_eval_button") + + if st.session_state.launch_eval_button: + selected_mbp_names = st.session_state.get( + "selected_model_based_metrics", [] + ) + selected_cbp_metrics = st.session_state.get("selected_cbp", []) + + all_metrics = selected_mbp_names + selected_cbp_metrics + + if not all_metrics: + st.warning("Please select at least one evaluation metric.") + return + + evaluation_data_list = [] + for idx, include_item in enumerate(st.session_state.include_in_evaluations): + if include_item: + user_input_values = st.session_state.human_rated_dict["user_input"][ + idx + ] + prompt_template = ( + st.session_state.local_prompt.prompt_to_run.prompt_data + ) + logger.info(f"Prompt template: {prompt_template}") + system_instruction = ( + st.session_state.local_prompt.prompt_to_run.system_instruction + ) + prediction = str( + st.session_state.human_rated_dict["assistant_response"][idx] + ) + + # Process reference value like in the old code + reference_val = st.session_state.human_rated_dict["ground_truth"][ + idx + ] + final_reference_str = "" + if isinstance(reference_val, int | float | bool): + final_reference_str = json.dumps({"value": reference_val}) + elif isinstance(reference_val, str): + try: + parsed_json = json.loads(reference_val) + if isinstance(parsed_json, int | float | bool): + final_reference_str = json.dumps({"value": parsed_json}) + else: + final_reference_str = reference_val + except json.JSONDecodeError: + final_reference_str = reference_val + elif isinstance(reference_val, dict | list): + final_reference_str = json.dumps(reference_val) + else: + final_reference_str = str(reference_val) + + instruction = ( + prompt_template.format(**user_input_values) + if isinstance(user_input_values, dict) + else prompt_template + ) + context = system_instruction if system_instruction else "" + prompt_str = ( + json.dumps(user_input_values) + if isinstance(user_input_values, dict) + else str(user_input_values) + ) + + eval_item = { + "context": context, + "instruction": instruction, + "prompt": prompt_str, + "prediction": prediction, + "reference": final_reference_str, + } + evaluation_data_list.append(eval_item) + + if not evaluation_data_list: + st.warning( + "No items were selected for evaluation. Please check the 'Include in Evaluation' checkboxes." + ) + return + + df_dataset = pd.DataFrame(evaluation_data_list) + st.session_state.df_dataset_eval = df_dataset + st.session_state.all_metrics_eval = all_metrics + logger.info(f"Evaluation DataFrame columns: {df_dataset.columns.tolist()}") + logger.info(f"Evaluation DataFrame head:\n{df_dataset.head()}") + + task = EvalTask( + dataset=df_dataset, + metrics=all_metrics, + experiment=os.getenv("EXPERIMENT_NAME"), + ) + st.session_state.eval_result = task.evaluate( + response_column_name="prediction", + baseline_model_response_column_name="reference", + ) + + st.markdown("---") + st.subheader("View Eval") + + st.button("View Evaluation Results", key="eval_results_button") + + if st.session_state.eval_result and st.session_state.eval_results_button: + print(st.session_state.eval_result.metrics_table) + st.dataframe(st.session_state.eval_result.metrics_table) + if st.session_state.eval_result: + st.markdown("---") + st.subheader("Summary Scores") + + mean_scores = {} + if ( + st.session_state.eval_result + and hasattr(st.session_state.eval_result, "metrics_table") + and not st.session_state.eval_result.metrics_table.empty + ): + for col in st.session_state.eval_result.metrics_table.columns: + if col.endswith("/score"): + scores = pd.to_numeric( + st.session_state.eval_result.metrics_table[col], + errors="coerce", + ) + if not scores.dropna().empty: + mean_scores[col] = scores.dropna().mean() + if mean_scores: + for metric, score in mean_scores.items(): + st.metric(label=f"Mean {metric}", value=f"{score:.2f}") + else: + st.metric( + label="Mean Automated Score", + value="N/A", + ) + + st.markdown("---") + st.subheader("Save to Prompt Records") + save_to_records = st.checkbox( + "I want to save the results of this evaluation to the prompt records.", + key="save_to_records_checkbox", + ) + if st.button("Save to Prompt Records", key="save_to_records_button"): + if save_to_records: + prompt_name = st.session_state.selected_prompt + prompt_version = st.session_state.selected_version + data_file = st.session_state.input_data_uri + + if "df_dataset_eval" in st.session_state: + data = st.session_state.df_dataset_eval.to_dict( + orient="records" + ) + else: + st.error( + "Evaluation data not found in session state. Please re-run evaluation." + ) + return + + if "all_metrics_eval" in st.session_state: + metrics = st.session_state.all_metrics_eval + else: + st.error( + "Metrics not found in session state. Please re-run evaluation." + ) + return + + scores_df = st.session_state.eval_result.metrics_table + scores = scores_df.to_dict(orient="records") + + mean_scores_to_save = {} + for col in scores_df.columns: + if col.endswith("/score"): + s = pd.to_numeric(scores_df[col], errors="coerce") + if not s.dropna().empty: + mean_scores_to_save[col] = s.dropna().mean() + + record_data = { + "prompt_name": prompt_name, + "prompt_version": prompt_version, + "data_file": data_file, + "metrics": metrics, + "mean_scores": mean_scores_to_save, + "scores": scores, + "evaluation_data": data, + "timestamp": datetime.datetime.now().isoformat(), + } + + try: + filename = f"record_{prompt_name}_v{prompt_version}_{datetime.datetime.now().strftime('%Y%m%d%H%M%S')}.json" + bucket = st.session_state.storage_client.bucket( + os.getenv("BUCKET") + ) + blob = bucket.blob(f"records/{filename}") + + json_data = json.dumps(record_data, indent=4) + blob.upload_from_string( + json_data, content_type="application/json" + ) + + gcs_path = f"gs://{os.getenv('BUCKET')}/records/{filename}" + + st.success( + f"Successfully saved to prompt records at: {gcs_path}" + ) + st.json(json_data) + except Exception as e: + st.error(f"Failed to save to GCS: {e}") + else: + st.warning( + "Please check the box to confirm you want to save to prompt records." + ) + + +if __name__ == "__main__": + main() diff --git a/tools/llmevalkit/pages/4_Prompt_Optimization.py b/tools/llmevalkit/pages/4_Prompt_Optimization.py new file mode 100644 index 00000000000..702a1ea3d19 --- /dev/null +++ b/tools/llmevalkit/pages/4_Prompt_Optimization.py @@ -0,0 +1,418 @@ +# Copyright 2025 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 +# +# http://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. + + +"""Streamlit page for running Vertex AI Prompt Optimization. + +This script provides a user interface for: +- Loading existing prompts from Vertex AI Prompt Registry. +- Loading datasets from a Google Cloud Storage bucket. +- Generating baseline responses and evaluating them against a ground truth. +- Configuring and launching a Vertex AI CustomJob for prompt optimization. +- Displaying baseline evaluation results. + +File Source: +https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/prompts/prompt_optimizer/vapo_lib.py +""" + +import json +import logging +import os +from argparse import Namespace + +import pandas as pd +import streamlit as st +from dotenv import load_dotenv +from etils import epath +from google.cloud import aiplatform, storage +from src import vapo_lib +from src.gcp_prompt import GcpPrompt as gcp_prompt +from vertexai.preview import prompts + +load_dotenv("src/.env") + +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + +TARGET_MODELS = ["gemini-2.0-flash-001", "gemini-2.0-flash-lite-001"] + + +def initialize_session_state() -> None: + """Initializes the session state variables.""" + if "op_id" not in st.session_state: + st.session_state.op_id = vapo_lib.get_id() + + if "local_prompt" not in st.session_state: + st.session_state.local_prompt = gcp_prompt() + + if "storage_client" not in st.session_state: + st.session_state["storage_client"] = storage.Client() + + if "data_uris" not in st.session_state: + st.session_state["data_uris"] = refresh_bucket() + + if "dataset" not in st.session_state: + st.session_state["dataset"] = None + + if "cached_data_files" not in st.session_state: + st.session_state.cached_data_files = {} + if "last_selected_dataset_for_cache" not in st.session_state: + st.session_state.last_selected_dataset_for_cache = None + + +def refresh_bucket() -> list[str]: + """Refreshes the list of available dataset URIs from the GCS bucket. + + This function lists all blobs in the configured GCS bucket, filters for + CSV and JSONL files located within the 'datasets/' prefix, and constructs a list + of their full gs:// URI paths. + + Returns: + A list of strings, where each string is a GCS URI to a dataset file. + """ + logger.info("Bucket: %s", os.getenv("BUCKET")) + bucket = st.session_state.storage_client.bucket(os.getenv("BUCKET")) + blobs = bucket.list_blobs() + data_uris = [] + for i in blobs: + if i.name.split("/")[0] == "datasets" and ( + i.name.endswith(".csv") or i.name.endswith(".jsonl") + ): + data_uris.append(f"gs://{i.bucket.name}/{i.name}") + logger.info("Data URIs: %s", data_uris) + return data_uris + + +def prompt_selection() -> None: + """Handles the prompt selection and loading.""" + st.selectbox( + "Select Existing Prompt", + options=st.session_state.local_prompt.existing_prompts.keys(), + placeholder="Select Prompt...", + key="selected_prompt", + ) + if st.session_state.selected_prompt: + logger.info("Prompt Meta: %s", st.session_state.local_prompt.prompt_meta) + st.session_state.local_prompt.prompt_meta["name"] = ( + st.session_state.selected_prompt + ) + versions = [ + i.version_id + for i in prompts.list_versions( + st.session_state.local_prompt.existing_prompts[ + st.session_state.selected_prompt + ] + ) + ] + + st.selectbox( + "Select Version", + options=versions, + placeholder="Select Version...", + key="selected_version", + ) + + st.button("Load Prompt", key="load_existing_prompt_button") + if st.session_state.load_existing_prompt_button: + logger.info( + "Selected Prompt ID: %s", + st.session_state.local_prompt.existing_prompts[ + st.session_state.selected_prompt + ], + ) + logger.info("Version: %s", st.session_state.selected_version) + st.session_state.local_prompt.load_prompt( + st.session_state.local_prompt.existing_prompts[ + st.session_state.selected_prompt + ], + st.session_state.selected_prompt, + st.session_state.selected_version, + ) + logger.info("Local Prompt Meta: %s", st.session_state.local_prompt.prompt_meta) + logger.info( + "Local Prompt Meta Dict Keys: %s", + st.session_state.local_prompt.prompt_meta.keys(), + ) + + st.session_state.prompt_name = ( + st.session_state.local_prompt.prompt_to_run.prompt_name + ) + st.session_state.prompt_data = ( + st.session_state.local_prompt.prompt_to_run.prompt_data + ) + st.session_state.model_name = ( + st.session_state.local_prompt.prompt_to_run.model_name.split("/")[-1] + ) + st.session_state.system_instructions = ( + st.session_state.local_prompt.prompt_to_run.system_instruction + ) + st.session_state.response_schema = json.dumps( + st.session_state.local_prompt.prompt_meta.get("response_schema", {}) + ) + st.session_state.generation_config = json.dumps( + st.session_state.local_prompt.prompt_meta.get("generation_config", {}) + ) + st.session_state.meta_tags = st.session_state.local_prompt.prompt_meta[ + "meta_tags" + ] + + +def dataset_selection() -> None: + """Handles the dataset selection and loading.""" + data_sets = list({i.split("/")[4] for i in st.session_state.data_uris}) + logger.info("Data Sets: %s", data_sets) + st.selectbox( + "Select an Existing Dataset", options=[None, *data_sets], key="selected_dataset" + ) + + files_to_display_in_selectbox = [] + if st.session_state.selected_dataset: + if ( + st.session_state.selected_dataset + != st.session_state.last_selected_dataset_for_cache + or st.session_state.selected_dataset + not in st.session_state.cached_data_files + ): + logger.info( + "Cache miss or dataset changed for files. Fetching for: %s", + st.session_state.selected_dataset, + ) + bucket = st.session_state.storage_client.bucket(os.getenv("BUCKET")) + prefix = f"datasets/{st.session_state.selected_dataset}/" + blobs_iterator = bucket.list_blobs(prefix=prefix) + + current_dataset_files = [] + for blob in blobs_iterator: + if ( + blob.name.endswith(".csv") or blob.name.endswith(".jsonl") + ) and not blob.name.endswith("/"): + filename = blob.name[len(prefix) :] + if filename: + current_dataset_files.append(filename) + + st.session_state.cached_data_files[st.session_state.selected_dataset] = ( + sorted(set(current_dataset_files)) + ) + st.session_state.last_selected_dataset_for_cache = ( + st.session_state.selected_dataset + ) + logger.info( + "Cached files for %s: %s", + st.session_state.selected_dataset, + st.session_state.cached_data_files[st.session_state.selected_dataset], + ) + + if "selected_file_from_dataset" in st.session_state: + st.session_state.selected_file_from_dataset = None + logger.info("Reset selected_file_from_dataset due to dataset change.") + + files_to_display_in_selectbox = st.session_state.cached_data_files.get( + st.session_state.selected_dataset, [] + ) + + st.selectbox( + "Select a file from this dataset:", + options=[None, *files_to_display_in_selectbox], + key="selected_file_from_dataset", + ) + + st.button("Load Dataset", key="load_existing_dataset_button") + if st.session_state.load_existing_dataset_button: + gcs_uri = f"gs://{os.getenv('BUCKET')}/datasets/{st.session_state.selected_dataset}/{st.session_state.selected_file_from_dataset}" + logger.info("Loading file: %s", gcs_uri) + if st.session_state.selected_file_from_dataset.endswith(".jsonl"): + st.session_state.dataset = pd.read_json(gcs_uri, lines=True) + else: + st.session_state.dataset = pd.read_csv(gcs_uri) + + if st.session_state.dataset is not None: + st.dataframe(st.session_state.dataset) + + +def get_optimization_args( + input_optimization_data_file_uri, output_optimization_run_uri +): + """Gets the arguments for the optimization job.""" + response_schema_str = st.session_state.local_prompt.prompt_meta.get( + "response_schema", "{}" + ) + try: + response_schema = ( + json.loads(response_schema_str) + if isinstance(response_schema_str, str) + else response_schema_str + ) + except json.JSONDecodeError: + response_schema = {} + + if response_schema and response_schema != {}: + response_mime_type = "application/json" + response_schema_arg = response_schema + else: + response_mime_type = "text/plain" + response_schema_arg = "" + + has_multimodal = False + if ( + st.session_state.dataset is not None + and "image" in st.session_state.dataset.columns + ): + has_multimodal = True + + return Namespace( + system_instruction=st.session_state.local_prompt.prompt_to_run.system_instruction, + prompt_template=( + f"{st.session_state.local_prompt.prompt_to_run.prompt_data}" + "\n\tAnswer: {target}" + ), + target_model="gemini-2.0-flash-001", + optimization_mode="instruction", + eval_metrics_types=[ + "question_answering_correctness", + ], + eval_metrics_weights=[ + 1.0, + ], + aggregation_type="weighted_sum", + input_data_path=input_optimization_data_file_uri, + output_path=f"gs://{output_optimization_run_uri}", + project=os.getenv("PROJECT_ID"), + num_steps=10, + num_demo_set_candidates=10, + demo_set_size=3, + target_model_location="us-central1", + source_model="", + source_model_location="", + target_model_qps=1, + optimizer_model_qps=1, + eval_qps=1, + source_model_qps="", + response_mime_type=response_mime_type, + response_schema=response_schema_arg, + language="English", + placeholder_to_content=json.loads("{}"), + data_limit=10, + translation_source_field_name="", + has_multimodal_inputs=has_multimodal, + ) + + +def start_optimization() -> None: + """Starts the optimization job.""" + st.divider() + + st.subheader("Run Optimization") + st.button("Start Optimization", key="start_optimization_button") + + if st.session_state.start_optimization_button: + workspace_uri = ( + epath.Path(os.getenv("BUCKET")) / "optimization" / st.session_state.op_id + ) + logger.info("Workspace URI: %s", workspace_uri) + + input_data_uri = epath.Path(workspace_uri) / "data" + logger.info("Input Data URI: %s", input_data_uri) + + workspace_uri.mkdir(parents=True, exist_ok=True) + input_data_uri.mkdir(parents=True, exist_ok=True) + + output_optimization_data_uri = epath.Path(workspace_uri) / "optimization_jobs" + logger.info("Output Data URI: %s", output_optimization_data_uri) + + prompt_optimization_job = ( + f"{st.session_state.selected_prompt}-" + f"{st.session_state.selected_version}-" + f"{st.session_state.selected_dataset}-" + f"{st.session_state.op_id}" + ) + output_optimization_run_uri = str( + output_optimization_data_uri / prompt_optimization_job + ) + input_optimization_data_file_uri = ( + f"gs://{input_data_uri}/{prompt_optimization_job}.jsonl" + ) + logger.info("Input Optimization Data URI: %s", input_optimization_data_file_uri) + if st.session_state.dataset is not None: + st.session_state.dataset.to_json( + str(input_optimization_data_file_uri), orient="records", lines=True + ) + else: + st.error("Please load a dataset first.") + return + + args = get_optimization_args( + input_optimization_data_file_uri, output_optimization_run_uri + ) + + with st.expander("Prompt Otimization Config"): + st.json(vars(args)) + + args = vars(args) + + config_file_uri = "gs://" + str(workspace_uri / "config" / "config.json") + + with epath.Path(config_file_uri).open("w") as config_file: + json.dump(args, config_file) + config_file.close() + st.success(f"Successfully wrote config file to {config_file_uri}") + + worker_pool_specs = [ + { + "machine_spec": { + "machine_type": "n1-standard-4", + }, + "replica_count": 1, + "container_spec": { + "image_uri": os.getenv("APD_CONTAINER_URI"), + "args": ["--config=" + config_file_uri], + }, + } + ] + + custom_job = aiplatform.CustomJob( + display_name=prompt_optimization_job, + worker_pool_specs=worker_pool_specs, + staging_bucket=str(workspace_uri), + ) + + custom_job.run(service_account=os.getenv("APD_SERVICE_ACCOUNT"), sync=False) + + st.success("Successfully Started Job!!") + + +def main() -> None: + """Streamlit page for Prompt Optimization.""" + st.set_page_config( + layout="wide", page_title="Prompt Optimization", page_icon="assets/favicon.ico" + ) + + initialize_session_state() + + st.header("Prompt Optimization") + + st.selectbox( + "Select Target Model for Optimization:", + options=TARGET_MODELS, + key="target_model_optimization", + ) + + prompt_selection() + dataset_selection() + start_optimization() + + +if __name__ == "__main__": + main() diff --git a/tools/llmevalkit/pages/5_Prompt_Optimization_Results.py b/tools/llmevalkit/pages/5_Prompt_Optimization_Results.py new file mode 100644 index 00000000000..d803bb9ce26 --- /dev/null +++ b/tools/llmevalkit/pages/5_Prompt_Optimization_Results.py @@ -0,0 +1,478 @@ +## Copyright 2025 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 + +import json +import logging +import os + +import pandas as pd +import streamlit as st +from dotenv import load_dotenv +from google.cloud import storage +from src import vapo_lib + +# Load environment variables +load_dotenv("src/.env") + +# Configure logging to the console +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + +# --- Constants --- +BASE_OPTIMIZATION_PREFIX = "optimization/" +OPTIMIZATION_JOBS_SUBDIR = "optimization_jobs/" + +from google.cloud import aiplatform + + +def list_custom_training_jobs(project_id: str, location: str): + """Lists all custom training jobs and their statuses in a given project and location. + + Args: + project_id: The Google Cloud project ID. + location: The region for the Vertex AI jobs, e.g., "us-central1". + + Returns: + A list of dictionaries, where each dictionary contains details of a custom job. + """ + # Initialize the Vertex AI client + # The API endpoint is determined by the location + client_options = {"api_endpoint": f"{location}-aiplatform.googleapis.com"} + client = aiplatform.gapic.JobServiceClient(client_options=client_options) + + # The parent resource path format + parent = f"projects/{project_id}/locations/{location}" + + # Make the API request to list custom jobs + response = client.list_custom_jobs(parent=parent) + + # Process the response and format the output + jobs_list = [] + print(f"Fetching jobs from project '{project_id}' in '{location}'...") + for job in response: + job_info = { + "display_name": job.display_name, + "name": job.name, + "status": job.state.name, # .name gets the string representation of the enum + } + jobs_list.append(job_info) + + print(f"Found {len(jobs_list)} jobs.") + return jobs_list + + +# --- Example Usage --- +if __name__ == "__main__": + # Replace with your project ID and desired location + PROJECT_ID = os.getenv("PROJECT_ID") + LOCATION = os.getenv("LOCATION") + + # Ensure you have authenticated with Google Cloud CLI: + # gcloud auth application-default login + + # And have the necessary permissions (e.g., "Vertex AI User" role) + + try: + all_jobs = list_custom_training_jobs(project_id=PROJECT_ID, location=LOCATION) + + # Print the results + if all_jobs: + print("\n--- Job Statuses ---") + for job in all_jobs: + print(f" - Name: {job['display_name']:<40} Status: {job['status']}") + print("--------------------\n") + else: + print("No custom jobs found.") + + except Exception as e: + print( + "\nAn error occurred. Please ensure your project ID and location are correct," + ) + print(f"and that you have authenticated correctly. Error: {e}") + + +def safe_json_loads(s): + """Safely loads a JSON string, returning the original value on failure.""" + if not isinstance(s, str): + return s + try: + return json.loads(s) + except (json.JSONDecodeError, TypeError): + return s + + +@st.cache_data(ttl=300) +def list_gcs_directories( + bucket_name: str, prefix: str, _storage_client: storage.Client +) -> list[str]: + """Lists 'directories' in GCS under a given prefix. + A 'directory' is inferred from the common prefixes of objects. + Caches the result for 5 minutes to improve performance. + """ + if not bucket_name: + st.warning("BUCKET environment variable is not set.") + return [] + if not _storage_client: + st.warning("Storage client is not initialized.") + return [] + + bucket = _storage_client.bucket(bucket_name) + retrieved_prefixes = set() + try: + for page in bucket.list_blobs(prefix=prefix, delimiter="/").pages: + retrieved_prefixes.update(page.prefixes) + + # The retrieved prefixes are the "subdirectories". + # e.g., for prefix 'optimization/', a retrieved prefix might be 'optimization/op_id/'. + # We want to extract just 'op_id'. + dir_names = [] + for p in retrieved_prefixes: + name = p.replace(prefix, "").strip("/") + if name: + dir_names.append(name) + return sorted(set(dir_names)) + except Exception as e: + st.error( + f"Error listing GCS directories under gs://{bucket_name}/{prefix}: {e}" + ) + logger.error( + f"Error listing GCS directories under gs://{bucket_name}/{prefix}: {e}", + exc_info=True, + ) + return [] + + +def _display_interactive_results(results_ui: vapo_lib.ResultsUI) -> None: + """Processes results from a VAPO run and displays them in an interactive + Streamlit UI with tabs for each prompt version. + """ + try: + if ( + not hasattr(results_ui, "templates") + or not results_ui.templates + or not hasattr(results_ui, "eval_results") + ): + logger.info( + "ResultsUI object does not have 'templates' or 'eval_results', or templates list is empty. Falling back." + ) + else: + processed_results_for_tabs = [] + for i, template_summary_df in enumerate(results_ui.templates): + if ( + not isinstance(template_summary_df, pd.DataFrame) + or template_summary_df.empty + ): + logger.warning( + f"Template summary data at index {i} is not a non-empty DataFrame. Skipping." + ) + continue + + # Get the detailed results to perform the custom calculation + detailed_eval_df = pd.DataFrame() + if i < len(results_ui.eval_results) and isinstance( + results_ui.eval_results[i], pd.DataFrame + ): + detailed_eval_df = results_ui.eval_results[i] + + # Add a custom exact_match calculation. This is more robust than simple + # string comparison as it handles differences in JSON key order and whitespace. + if ( + not detailed_eval_df.empty + and "ground_truth" in detailed_eval_df.columns + and "reference" in detailed_eval_df.columns + ): + # Parse the JSON strings into Python objects before comparing. + parsed_ground_truths = detailed_eval_df["ground_truth"].apply( + safe_json_loads + ) + parsed_references = detailed_eval_df["reference"].apply( + safe_json_loads + ) + + # Create a boolean series for the comparison + is_match = parsed_ground_truths.eq(parsed_references) + + # Map boolean to 'yes'/'no' for display in the detailed table + detailed_eval_df["calculated_exact_match"] = is_match.map( + {True: "yes", False: "no"} + ) + + # Calculate the mean from the boolean series for the summary metric + new_exact_match_mean = is_match.mean() + template_summary_df["metrics.calculated_exact_match/mean"] = ( + new_exact_match_mean + ) + + prompt_text = "Prompt text not found in template data." + if "prompt" in template_summary_df.columns: + prompt_text = template_summary_df["prompt"].iloc[0] + else: + logger.warning( + f"Column 'prompt' not found in template_summary_df at index {i}." + ) + + # Determine the primary score and build the tab name. + primary_score_label = "Score" + primary_score_value = "N/A" + if "metrics.calculated_exact_match/mean" in template_summary_df.columns: + primary_score_label = "Calculated Exact Match" + primary_score_value = template_summary_df[ + "metrics.calculated_exact_match/mean" + ].iloc[0] + else: + # Fallback to the first available metric + mean_metric_columns = [ + col + for col in template_summary_df.columns + if col.startswith("metrics.") and "/mean" in col + ] + if mean_metric_columns: + first_metric_col = mean_metric_columns[0] + primary_score_label = ( + first_metric_col.replace("metrics.", "") + .replace("/mean", "") + .replace("_", " ") + .title() + ) + primary_score_value = template_summary_df[ + first_metric_col + ].iloc[0] + + # Build the tab name with all available metrics for a quick overview. + tab_name_metrics_parts = [] + mean_metric_columns = [ + col + for col in template_summary_df.columns + if col.startswith("metrics.") and "/mean" in col + ] + for metric_col in mean_metric_columns: + metric_name_short = metric_col.replace("metrics.", "").replace( + "/mean", "" + ) + metric_val = template_summary_df[metric_col].iloc[0] + if metric_name_short == "calculated_exact_match" and isinstance( + metric_val, float + ): + tab_name_metrics_parts.append( + f"{metric_name_short}: {metric_val:.1%}" + ) + else: + tab_name_metrics_parts.append( + f"{metric_name_short}: {metric_val:.3f}" + if isinstance(metric_val, float) + else f"{metric_name_short}: {metric_val}" + ) + + tab_name = f"Template {i}" + if tab_name_metrics_parts: + tab_name += f" ({', '.join(tab_name_metrics_parts)})" + + current_summary_df_display = template_summary_df.copy() + if "prompt" in current_summary_df_display.columns: + current_summary_df_display = current_summary_df_display.drop( + columns=["prompt"] + ) + + processed_results_for_tabs.append( + { + "name": tab_name, + "template_text": prompt_text, + "primary_score_label": primary_score_label, + "primary_score_value": primary_score_value, + "summary_metrics_df": current_summary_df_display, + "detailed_eval_df": detailed_eval_df, + } + ) + + if ( + processed_results_for_tabs + ): # If we successfully processed data, show the new UI + st.write("### Interactive Prompt Versions") + tab_titles = [res["name"] for res in processed_results_for_tabs] + tabs = st.tabs(tab_titles) + + for i, tab_content in enumerate(tabs): + with tab_content: + result_data = processed_results_for_tabs[i] + + st.subheader("Prompt Template") + # Sanitize tab name for key + clean_key_name = "".join( + filter(str.isalnum, result_data["name"]) + ) + st.text_area( + "Template", + value=result_data["template_text"], + height=200, + disabled=True, + key=f"template_view_{clean_key_name}_{i}", + ) + + st.subheader("Primary Score") + score_val = result_data["primary_score_value"] + score_label = result_data["primary_score_label"] + if score_label == "Calculated Exact Match" and isinstance( + score_val, float + ): + st.metric(label=score_label, value=f"{score_val:.2%}") + else: + st.metric( + label=score_label, + value=f"{score_val:.4f}" + if isinstance(score_val, float) + else str(score_val), + ) + + if not result_data["summary_metrics_df"].empty: + st.subheader("Summary Metrics (from templates.json)") + st.dataframe(result_data["summary_metrics_df"]) + + if not result_data["detailed_eval_df"].empty: + st.subheader( + "Detailed Evaluation Results (from eval_results.json)" + ) + st.dataframe(result_data["detailed_eval_df"]) + else: + st.caption( + "No detailed evaluation results available for this template." + ) + else: + st.warning("No valid results could be processed for display.") + + except Exception as e: + st.error(f"An error occurred while trying to display results: {e}") + logger.error(f"Error in results display section: {e}", exc_info=True) + st.markdown( + "For now, you can access the results directly at the GCS path shown above." + ) + + +def main() -> None: + """Renders the Streamlit page for viewing Prompt Optimization Results.""" + st.set_page_config( + layout="wide", + page_title="Prompt Optimization Results", + page_icon="assets/favicon.ico", + ) + st.header("Prompt Optimization Results Browser") + + if "storage_client" not in st.session_state: + try: + st.session_state["storage_client"] = storage.Client() + logger.info("Storage client initialized.") + except Exception as e: + st.error(f"Failed to initialize Google Cloud Storage client: {e}") + logger.error( + f"Failed to initialize Google Cloud Storage client: {e}", exc_info=True + ) + st.session_state["storage_client"] = None + return + + bucket_name = os.getenv("BUCKET") + if not bucket_name: + st.error("BUCKET environment variable is not set. Please configure it in .env.") + return + + # --- Step 1: Select Operation ID --- + op_ids = list_gcs_directories( + bucket_name, BASE_OPTIMIZATION_PREFIX, st.session_state.storage_client + ) + if not op_ids: + st.info( + f"No optimization operation IDs found under gs://{bucket_name}/{BASE_OPTIMIZATION_PREFIX}" + ) + return + + if "op_id" in st.session_state and st.session_state.op_id: + st.caption( + f"Hint: The last optimization run you initiated had the ID: `{st.session_state.op_id}`." + ) + + selected_op_id = st.selectbox( + "Select an Operation ID:", options=[None, *op_ids], key="selected_op_id_results" + ) + if not selected_op_id: + st.write("Please select an Operation ID to see its optimization job runs.") + return + + st.divider() + + # --- Step 2: Select Experiment Run --- + st.subheader(f"Optimization Job Runs for Operation ID: {selected_op_id}") + optimization_jobs_prefix = ( + f"{BASE_OPTIMIZATION_PREFIX}{selected_op_id}/{OPTIMIZATION_JOBS_SUBDIR}" + ) + experiment_runs = list_gcs_directories( + bucket_name, optimization_jobs_prefix, st.session_state.storage_client + ) + + if not experiment_runs: + st.info( + f"No completed optimization job runs found under gs://{bucket_name}/{optimization_jobs_prefix}" + ) + return + + selected_run = st.selectbox( + "Select an Optimization Job Run:", + options=[None, *experiment_runs], + key="selected_experiment_run", + ) + if not selected_run: + st.write("Please select an optimization job run to view its results.") + return + + st.divider() + + # --- Step 3: Check Job Status and Display Results --- + st.subheader(f"Results for: {selected_run}") + + project_id = os.getenv("PROJECT_ID") + location = os.getenv("LOCATION") + + if not project_id or not location: + st.error("PROJECT_ID or REGION environment variables are not set.") + return + + try: + jobs = list_custom_training_jobs(project_id=project_id, location=location) + job_status = "Not Found" + for job in jobs: + if job["display_name"] == selected_run: + job_status = job["status"] + break + + st.info(f"Status for job '{selected_run}': **{job_status}**") + + if job_status == "JOB_STATE_FAILED": + st.error( + "This optimization job has failed. Please check the logs in the Vertex AI console for more details." + ) + return + if job_status not in ["JOB_STATE_SUCCEEDED", "JOB_STATE_CANCELLED"]: + st.warning( + f"Job is currently in status: {job_status}. Results may be incomplete." + ) + + except Exception as e: + st.error(f"Could not retrieve job status. Error: {e}") + logger.error( + f"Failed to retrieve job status for {selected_run}: {e}", exc_info=True + ) + + run_uri = f"gs://{bucket_name}/{optimization_jobs_prefix}{selected_run}" + st.info(f"Loading results from: {run_uri}") + results_ui = vapo_lib.ResultsUI(run_uri) + _display_interactive_results(results_ui) + + +if __name__ == "__main__": + main() diff --git a/tools/llmevalkit/pages/6_Prompt_Records.py b/tools/llmevalkit/pages/6_Prompt_Records.py new file mode 100644 index 00000000000..d1bc6e2734d --- /dev/null +++ b/tools/llmevalkit/pages/6_Prompt_Records.py @@ -0,0 +1,132 @@ +# Copyright 2025 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 +# +# http://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. + +import json +import logging +import os + +import pandas as pd +import streamlit as st +from dotenv import load_dotenv +from google.cloud import storage + +load_dotenv("src/.env") + +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + + +def load_records_from_gcs(bucket_name: str, prefix: str) -> pd.DataFrame: + """Loads all JSON record files from a GCS prefix and returns a DataFrame.""" + try: + storage_client = storage.Client() + bucket = storage_client.bucket(bucket_name) + blobs = bucket.list_blobs(prefix=prefix) + + all_records = [] + for blob in blobs: + if blob.name.endswith(".json"): + logger.info(f"Loading record from {blob.name}") + try: + record_data = json.loads(blob.download_as_string()) + if isinstance(record_data, list): + all_records.extend(record_data) + else: + all_records.append(record_data) + except json.JSONDecodeError: + logger.warning(f"Could not decode JSON from {blob.name}") + except Exception as e: + logger.exception(f"Failed to process blob {blob.name}: {e}") + + if not all_records: + st.warning(f"No JSON records found at gs://{bucket_name}/{prefix}") + return pd.DataFrame() + + return pd.json_normalize(all_records) + + except Exception as e: + st.error(f"Failed to load or parse records from GCS: {e}") + logger.error("Error loading records: %s", e, exc_info=True) + return pd.DataFrame() + + +def main() -> None: + """Renders the Prompt Records Leaderboard page.""" + st.set_page_config( + layout="wide", + page_title="Prompt Records Leaderboard", + page_icon="assets/favicon.ico", + ) + st.header("Prompt Records Leaderboard") + st.markdown( + "This page allows you to view and compare the evaluation results of different prompt versions." + ) + + records_prefix = "records/" + + if "leaderboard_df" not in st.session_state: + st.session_state.leaderboard_df = pd.DataFrame() + + if st.button("Load/Refresh Leaderboard"): + with st.spinner("Loading records from GCS..."): + st.session_state.leaderboard_df = load_records_from_gcs( + os.getenv("BUCKET"), records_prefix + ) + if not st.session_state.leaderboard_df.empty: + st.success("Leaderboard loaded successfully.") + else: + st.info("Leaderboard is empty or could not be loaded.") + + if st.session_state.leaderboard_df.empty: + st.info("Click the button above to load the leaderboard data.") + return + + st.divider() + + prompt_names = st.session_state.leaderboard_df["prompt_name"].unique().tolist() + selected_prompt = st.selectbox( + "Select a Prompt to Compare Versions", options=[None, *prompt_names] + ) + + if selected_prompt: + st.subheader(f"Comparison for: {selected_prompt}") + + prompt_df = st.session_state.leaderboard_df[ + st.session_state.leaderboard_df["prompt_name"] == selected_prompt + ].copy() + + if prompt_df.empty: + st.info("No records found for the selected prompt.") + return + + # Extract scores into separate columns for easier analysis + scores_df = pd.json_normalize(prompt_df["scores"]) + scores_df.columns = [f"score.{col}" for col in scores_df.columns] + + comparison_df = pd.concat([prompt_df.reset_index(drop=True), scores_df], axis=1) + + # Clean up the view + display_columns = [ + col + for col in comparison_df.columns + if col not in ["scores", "evaluation_data"] + and not (col.startswith("score.") and col[6:].isdigit()) + ] + st.dataframe(comparison_df[display_columns]) + + +if __name__ == "__main__": + main() diff --git a/tools/llmevalkit/prompt-management-tutorial.ipynb b/tools/llmevalkit/prompt-management-tutorial.ipynb new file mode 100644 index 00000000000..60d9d54d068 --- /dev/null +++ b/tools/llmevalkit/prompt-management-tutorial.ipynb @@ -0,0 +1,580 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "id": "81450b47de75" + }, + "outputs": [], + "source": [ + "# Copyright 2025 Google LLC\n", + "#\n", + "# Licensed under the Apache License, Version 2.0 (the \"License\");\n", + "# you may not use this file except in compliance with the License.\n", + "# You may obtain a copy of the License at\n", + "#\n", + "# http://www.apache.org/licenses/LICENSE-2.0\n", + "#\n", + "# Unless required by applicable law or agreed to in writing, software\n", + "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", + "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", + "# See the License for the specific language governing permissions and\n", + "# limitations under the License." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "a204d0ab284d" + }, + "source": [ + "# Tutorial for Running Prompt Management and Evaluation\n", + "\n", + "| \n",
+ " \n",
+ " Open in Colab\n", + " \n", + " | \n",
+ " \n",
+ " \n",
+ " Open in Colab Enterprise\n", + " \n", + " | \n",
+ " \n",
+ " \n",
+ " Open in Vertex AI Workbench\n", + " \n", + " | \n",
+ " \n",
+ " \n",
+ " View on GitHub\n", + " \n", + " | \n",
+ "
![\"LinkedIn](\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\"){kind=icon}
\n",
+ "\n",
+ "\n",
+ "\n",
+ " ![\"Bluesky](\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\"){kind=logo}
\n",
+ "\n",
+ "\n",
+ "\n",
+ " ![\"X](\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\"){kind=icon}
\n",
+ "\n",
+ "\n",
+ "\n",
+ " ![\"Reddit](\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\"){kind=logo}
\n",
+ "\n",
+ "\n",
+ "\n",
+ " ![\"Facebook](\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\"){kind=logo}
\n",
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "8aee03ecd776"
+ },
+ "source": [
+ "| Author(s) |\n",
+ "| --- |\n",
+ "| [Mike Santoro](https://github.com/Michael-Santoro) |"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "7cb0282fcd41"
+ },
+ "source": [
+ "## 1. Overview\n",
+ "\n",
+ "This tutorial provides a comprehensive guide to prompt engineering, covering the entire lifecycle from creation to evaluation and optimization. It's broken down into the following sections:\n",
+ "\n",
+ "1. **Prompt Management:** This section focuses on the core tasks of creating, editing, and managing prompts. You can: \n",
+ " - **Create new prompts:** Define the prompt's name, text, the model it's designed for, and any system instructions. \n",
+ " - **Load and edit existing prompts:** Browse a library of saved prompts, load a specific version, and make modifications.\n",
+ " - **Test prompts:** Before saving, you can provide sample input and generate a response to see how the prompt performs.\n",
+ " - **Versioning:** Each time you save a change to a prompt, a new version is created, allowing you to track its evolution and compare different iterations.\n",
+ "\n",
+ "2. **Dataset Creation:** A crucial part of prompt engineering is having good data to test and evaluate your prompts. This section allows you to:\n",
+ "\n",
+ " - **Create new datasets:** A dataset is essentially a folder in Google Cloud Storage where you can group related files.\n",
+ " - **Upload data:** You can upload files in CSV, JSON, or JSONL format to your datasets. This data will be used for evaluating your prompts.\n",
+ "\n",
+ "3. **Evaluation:** Once you have a prompt and a dataset, you need to see how well the prompt performs. The evaluation section helps you with this by:\n",
+ "\n",
+ " - **Running evaluations:** You can select a prompt and a dataset and run an evaluation. This will generate responses from the model for each item in your dataset.\n",
+ " - **Human-in-the-loop rating:** For a more nuanced evaluation, you can manually review the model's responses and rate them.\n",
+ " - **Automated metrics:** The tutorial also supports automated evaluation metrics to get a quantitative measure of your prompt's performance.\n",
+ "\n",
+ "4. **Prompt Optimization:** This section helps you automatically improve your prompts. It uses Vertex AI's prompt optimization capabilities to:\n",
+ "\n",
+ " - **Configure and launch optimization jobs:** You can set up and run a job that will take your prompt and a dataset and try to find a better-performing version of the prompt.\n",
+ "\n",
+ "5. **Prompt Optimization Results:** After an optimization job has run, this section allows you to:\n",
+ "\n",
+ " - **View the results:** You can see the different prompt versions that the optimizer came up with and how they performed.\n",
+ " - **Compare versions:** The results are presented in a way that makes it easy to compare the different optimized prompts and choose the best one.\n",
+ "\n",
+ "6. **Prompt Records:** This is a leaderboard that shows you the evaluation results of all your different prompt versions. It helps you to:\n",
+ "\n",
+ " - **Track performance over time:** See how your prompts have improved with each new version.\n",
+ " - **Compare different prompts:** You can compare the performance of different prompts for the same task.\n",
+ "\n",
+ "In summary, this tutorial provides a complete and integrated environment for all your prompt engineering needs, from initial creation to sophisticated optimization and evaluation.\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "bbd5f8c2144a"
+ },
+ "source": [
+ "## 2. Before you start"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "6b2e004206a7"
+ },
+ "source": [
+ "### Clone the GitHub Repo"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "bcafdf59c8b9"
+ },
+ "outputs": [],
+ "source": [
+ "! git clone https://github.com/GoogleCloudPlatform/generative-ai.git"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "36e497494e38"
+ },
+ "outputs": [],
+ "source": [
+ "! gsutil cp gs://github-repo/prompts/prompt_optimizer/mathvista_dataset/mathvista_input.jsonl mathvista_input.jsonl"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "7445a03507f7"
+ },
+ "source": [
+ "### Install Python Dependencies"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "44fee8ab2679"
+ },
+ "outputs": [],
+ "source": [
+ "% pip install -r requirements.txt"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "870814a62e87"
+ },
+ "source": [
+ "### Authenticate your notebook environment (Colab only)\n",
+ "\n",
+ "Authenticate your environment on Google Colab."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "9b6bfee6ba31"
+ },
+ "outputs": [],
+ "source": [
+ "import sys\n",
+ "\n",
+ "if \"google.colab\" in sys.modules:\n",
+ " from google.colab import auth\n",
+ "\n",
+ " auth.authenticate_user()"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "f5cee1b8b3ab"
+ },
+ "source": [
+ "### Alternative Authenticate"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "4632dbaa3b73"
+ },
+ "outputs": [],
+ "source": [
+ "# fmt: off\n",
+ "PROJECT_ID = \"[your-project-id]\" # @param {type: \"string\", placeholder: \"[your-project-id]\", isTemplate: true}\n",
+ "LOCATION = \"[your-project-region]\" # @param{type: \"string\", placeholder: \"[your-project-region]\", isTemplate: true}\n",
+ "# fmt: on\n",
+ "\n",
+ "! gcloud auth application-default login\n",
+ "! gcloud config set project {PROJECT_ID}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "390a62d0e8de"
+ },
+ "source": [
+ "### Set Google Cloud project information\n",
+ "\n",
+ "**TO-DO: Check these APIs**\n",
+ "To get started using Vertex AI, you must have an existing Google Cloud project and [enable the following APIs](https://console.cloud.google.com/flows/enableapi?apiid=cloudresourcemanager.googleapis.com,aiplatform.googleapis.com,cloudfunctions.googleapis.com,run.googleapis.com).\n",
+ "\n",
+ "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment)."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "ad1be5b6c02d"
+ },
+ "outputs": [],
+ "source": [
+ "! cp src/.env.example src/.env"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "94dd1d71983f"
+ },
+ "source": [
+ "### Copy sample.env and Modify\n",
+ "\n",
+ "- BUCKET_NAME - Pick an existing bucket or make a new one below\n",
+ "- PROJECT_ID\n",
+ "- SERVICE_ACCOUNT - Created Below\n",
+ "\n",
+ "\n",
+ "#### Create a New Bucket (Not Required if using existing)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "e27ab4bb8a87"
+ },
+ "outputs": [],
+ "source": [
+ "# fmt: off\n",
+ "BUCKET_NAME = \"[your-bucket-name]\" # @param {type: \"string\", placeholder: \"[your-bucket-name]\", isTemplate: true}\n",
+ "# fmt: on\n",
+ "\n",
+ "BUCKET_URI = f\"gs://{BUCKET_NAME}\"\n",
+ "\n",
+ "\n",
+ "! gsutil mb -l {LOCATION} {BUCKET_URI}"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "33cc66a9351d"
+ },
+ "source": [
+ "#### Create a Service Account"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "492567ee96db"
+ },
+ "outputs": [],
+ "source": [
+ "PROJECT_NUMBER = !gcloud projects describe {PROJECT_ID} --format=\"get(projectNumber)\"[0]\n",
+ "PROJECT_NUMBER = PROJECT_NUMBER[0]\n",
+ "SERVICE_ACCOUNT = f\"{PROJECT_NUMBER}-compute@developer.gserviceaccount.com\"\n",
+ "\n",
+ "for role in ['aiplatform.user', 'storage.objectAdmin']:\n",
+ "\n",
+ " ! gcloud projects add-iam-policy-binding {PROJECT_ID} \\\n",
+ " --member=serviceAccount:{SERVICE_ACCOUNT} \\\n",
+ " --role=roles/{role} --condition=None"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "4bc0ae8e0c5f"
+ },
+ "source": [
+ "## 3. Run the App"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "21fed89471c8"
+ },
+ "outputs": [],
+ "source": [
+ "! cd generative-ai/llmevalkit && streamlit run index.py & npx localtunnel --port 8501"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "90a521d73796"
+ },
+ "source": [
+ "Click the link and use just the external ip as the password.\n",
+ "\n",
+ "📝 **Note:** You can run `wget -q -O - https://loca.lt/mytunnelpassword` to get the external ip (i.e 35.194.128.20)\n",
+ "\n",
+ "📝 **Note:** If you are having issues displaying the app, clear your cache."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "1e31b6cc7645"
+ },
+ "source": [
+ ""
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "9ca8d05072cd"
+ },
+ "source": [
+ "## 4. Work with the App"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "c7c11886dc17"
+ },
+ "source": [
+ "### 1. Prompt Management\n",
+ "\n",
+ "In the Prompt Name field enter:\n",
+ "\n",
+ "```\n",
+ "math_prompt_test\n",
+ "```\n",
+ "\n",
+ "In the Prompt Data field enter:\n",
+ "\n",
+ "```\n",
+ "Problem: {{query}}\n",
+ "Image: {{image}} @@@image/jpeg\n",
+ "Answer: {{target}}\n",
+ "```\n",
+ "\n",
+ "In the Model Name field enter:\n",
+ "```\n",
+ "gemini-2.0-flash-001\n",
+ "```\n",
+ "\n",
+ "In the System Instructions field enter:\n",
+ "```\n",
+ "Solve the problem given the image.\n",
+ "```\n",
+ "\n",
+ "Click `Save`\n",
+ "\n",
+ "Copy this text for testing:\n",
+ "\n",
+ "```\n",
+ "{\"query\": \"Hint: Please answer the question and provide the correct option letter, e.g., A, B, C, D, at the end.\\nQuestion: As shown in the figure, CD is the diameter of \\u2299O, chord DE \\u2225 OA, if the degree of \\u2220D is 50.0, then the degree of \\u2220C is ()\\nChoices:\\n(A) 25\\u00b0\\n(B) 30\\u00b0\\n(C) 40\\u00b0\\n(D) 50\\u00b0\", \"image\": \"gs://github-repo/prompts/prompt_optimizer/mathvista_dataset/images/643.jpg\", \"target\": \"25\\u00b0\"}\n",
+ "```\n",
+ "\n",
+ "🖱️ Click `Generate`.\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "306b19d374b6"
+ },
+ "source": [
+ "### 2. Dataset Creation\n",
+ "\n",
+ "Download a copy of the dataset. Then upload this file in the application.\n",
+ "\n",
+ "**Dataset Name:** `mathvista`\n",
+ "\n",
+ "You can preview the dataset at the bottom of the page."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {
+ "id": "814fb0866bd5"
+ },
+ "outputs": [],
+ "source": [
+ "! gsutil cp gs://github-repo/prompts/prompt_optimizer/mathvista_dataset/mathvista_input.jsonl ."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "0b06fa624a7a"
+ },
+ "source": [
+ "### 3. Evaluation\n",
+ "\n",
+ "We will now run an evaluation, prior to doing any tweaking to get a baseline.\n",
+ "\n",
+ "- **Existing Dataset:** 'mathvista'\n",
+ "- **Dataset File:** 'mathvista_input.jsonl'\n",
+ "- **Number of Samples:** '100'\n",
+ "- **Ground Truth Column Name:** 'target'\n",
+ "- **Existing Prompt:** 'math_prompt_test'\n",
+ "- **Version:** '1'\n",
+ "\n",
+ "Click Load Prompt, and Upload and Get Response... ⏰ Wait!!\n",
+ "\n",
+ "Review the responses.\n",
+ "\n",
+ "- **Model-Based:** 'question-answering-quality'\n",
+ "\n",
+ "Launch the Eval... ⏰ Wait!!\n",
+ "\n",
+ "View the Evaluation Results, and save to prompt records. This will save this initial version to the prompt records for the baseline.\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "c76e0e5065fa"
+ },
+ "source": [
+ "### 4. Prompt Optimization\n",
+ "\n",
+ "🔧 Set-Up Prompt Optimization.\n",
+ "\n",
+ "- **Target Model:** 'gemini-2.0-flash-001'\n",
+ "- **Existing Prompt:** 'math_prompt_test'\n",
+ "- **Version:** '1'\n",
+ "\n",
+ "🖱️ Click Load Prompt.\n",
+ "\n",
+ "- **Select Existing Dataset:** 'mathvista'\n",
+ "- **Select the File:** 'mathvista_input.jsonl'\n",
+ "\n",
+ "🖱️ Click Load Dataset.\n",
+ "\n",
+ "Preview the dataset.\n",
+ "\n",
+ "🖱️ Click Start Optimization.\n",
+ "\n",
+ "**Note:** If Interested in viewing the progress, Navigate to https://console.cloud.google.com/vertex-ai/training/custom-jobs\n",
+ "\n",
+ "⏰ Wait!! This step will take about 20-min to run."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "9987b016703e"
+ },
+ "source": [
+ "### 5. Prompt Optimization Results\n",
+ "\n",
+ "View the Optimization Results.\n",
+ "\n",
+ "The last run will be shown at the top of the screen. Pick this from the dropdown menu: \n",
+ "\n",
+ "\n",
+ "\n",
+ "Review the results and select the highest scoring version and copy the instruction."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "f15bce27fc94"
+ },
+ "source": [
+ "### 6. Navigate Back to Prompt for New Version\n",
+ "\n",
+ "Load your existing prompt from before.\n",
+ "\n",
+ "📋 Paste your new instructions from the prompt optimizer, and save new version."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "68957d00c89e"
+ },
+ "source": [
+ "### 7. Run new Evaluation\n",
+ "\n",
+ "Repeat step 3 with your new version."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {
+ "id": "c2e6c524d69a"
+ },
+ "source": [
+ "### 8. View the Records\n",
+ "\n",
+ "Navigate to the leaderboard and load the results."
+ ]
+ }
+ ],
+ "metadata": {
+ "colab": {
+ "name": "prompt-management-tutorial.ipynb",
+ "toc_visible": true
+ },
+ "kernelspec": {
+ "display_name": "Python 3",
+ "name": "python3"
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 0
+}
diff --git a/tools/llmevalkit/requirements.txt b/tools/llmevalkit/requirements.txt
new file mode 100644
index 00000000000..0d7619f1a5c
--- /dev/null
+++ b/tools/llmevalkit/requirements.txt
@@ -0,0 +1,25 @@
+ipykernel
+google-cloud-storage
+streamlit
+pandas
+plotly
+matplotlib
+db-dtypes
+nbformat>=4.2.0
+streamlit-tags
+dotenv
+google-cloud-pipeline-components
+gcsfs
+pipfile
+ipython
+asyncio
+tqdm
+tenacity
+etils
+importlib-resources
+fsspec
+ipywidgets
+google-cloud-aiplatform
+tensorflow
+pylint
+google-genai
diff --git a/tools/llmevalkit/src/.env.example b/tools/llmevalkit/src/.env.example
new file mode 100644
index 00000000000..94a92a45575
--- /dev/null
+++ b/tools/llmevalkit/src/.env.example
@@ -0,0 +1,25 @@
+# Project Configuration
+BUCKET = ""
+PROMPT_PREFIX = "prompts"
+PROJECT_ID = ""
+LOCATION = ""
+
+START_PROMPT_PATH = "src/initial_local_prompt_empty.json"
+
+## Constants for Eval
+JITTER_AMOUNT = 0.1
+PLOTLY_RENDERER = "colab"
+DEFAULT_HUMAN_RATING_COL = "human_rating"
+DEFAULT_SCORE_COL = "score"
+EXPERIMENT_NAME = "eval-open-judge-test"
+
+## Constants for Optimize
+OPTIMIZATION_MODE = "instruction"
+EVAL_METRIC = "exact_match"
+NUM_INST_OPTIMIZATION_STEPS = 10
+TARGET_MODEL_QPS = 3.0
+EVAL_QPS = 3.0
+RESPONSE_MIME_TYPE = "text/plain"
+TARGET_LANGUAGE = "English"
+SERVICE_ACCOUNT = ""
+APD_CONTAINER_URI = "us-docker.pkg.dev/vertex-ai-restricted/builtin-algorithm/apd:preview_v1_0"
diff --git a/tools/llmevalkit/src/__init__.py b/tools/llmevalkit/src/__init__.py
new file mode 100644
index 00000000000..0a2669d7a25
--- /dev/null
+++ b/tools/llmevalkit/src/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2025 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
+#
+# http://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.
diff --git a/tools/llmevalkit/src/gcp_dataset.py b/tools/llmevalkit/src/gcp_dataset.py
new file mode 100644
index 00000000000..8131353695e
--- /dev/null
+++ b/tools/llmevalkit/src/gcp_dataset.py
@@ -0,0 +1,95 @@
+# Copyright 2025 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
+#
+# http://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.
+
+
+"""This module provides utility functions for managing and processing datasets
+stored on Google Cloud Storage (GCS).
+
+It includes functionalities to:
+- List available datasets from a specified GCS bucket.
+- Fetch and read CSV files from GCS into pandas DataFrames.
+- Process raw data from CSVs to generate structured user prompts and expected
+ outcomes for model evaluation.
+- A helper function to escape special characters in text to be used with the
+ Gemini API.
+
+The primary purpose is to abstract the GCS interactions and data preprocessing
+steps required for the prompt management and evaluation application.
+"""
+
+import io
+import logging
+import os
+
+import pandas as pd
+from dotenv import load_dotenv
+from google.cloud import storage
+
+logging.basicConfig(
+ level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+load_dotenv("src/.env")
+
+
+PREFIX = "datasets_meta"
+
+
+def get_existing_datasets() -> list[str]:
+ """Gets the list of existing dataset names from the GCS bucket.
+
+ Returns:
+ list[str]: A list of dataset names.
+ """
+ storage_client = storage.Client()
+ blobs = storage_client.list_blobs(os.getenv("BUCKET_NAME"), prefix=PREFIX)
+ return [i.name.split("/")[-1] for i in blobs]
+
+
+def escape_special_characters(text: str) -> str:
+ """Escapes special characters for Gemini Flash API."""
+ if not isinstance(text, str):
+ return text
+ text = text.replace("\\", "\\")
+ text = text.replace("\n", "\n")
+ text = text.replace("\r", "\r")
+ text = text.replace("\t", "\t")
+ return text.replace('"', '"')
+
+
+def process_csv_from_gcs(bucket_name: str, file_path: str) -> pd.DataFrame:
+ """Reads a CSV file from GCS, processes each row to create user prompts
+ and expected results, and returns a pandas DataFrame.
+
+ Args:
+ bucket_name: Name of the GCS bucket.
+ file_path: Path to the CSV file within the bucket.
+
+ Returns:
+ pandas.DataFrame: DataFrame with original data and added user_prompt
+ and expected_result columns.
+ """
+ try:
+ client = storage.Client()
+ bucket = client.get_bucket(bucket_name)
+ blob = bucket.blob(file_path)
+
+ content = blob.download_as_bytes()
+
+ return pd.read_csv(io.BytesIO(content))
+
+ except (OSError, ValueError) as e:
+ print(f"An error occurred: {e}")
+ return pd.DataFrame()
diff --git a/tools/llmevalkit/src/gcp_evaluation.py b/tools/llmevalkit/src/gcp_evaluation.py
new file mode 100644
index 00000000000..45f7c67dff6
--- /dev/null
+++ b/tools/llmevalkit/src/gcp_evaluation.py
@@ -0,0 +1,401 @@
+# Copyright 2025 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
+#
+# http://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.
+
+"""Provides tools for evaluating and visualizing model performance against human ratings.
+
+This module contains functions to process evaluation data, which includes both
+human-provided scores and model-generated scores. It prepares the data into a
+pandas DataFrame, extracts key metrics like the confusion matrix, and generates
+a series of Plotly visualizations to compare the two sets of scores.
+
+The primary functions are:
+- prepare_dataframe: Cleans and formats the raw evaluation data.
+- extract_completeness_metrics: Pulls confusion matrix data from metrics logs.
+- plot_distribution_comparison: Compares the distribution of human vs. model scores.
+- plot_confusion_matrix: Creates a heatmap to show agreement and disagreement.
+- plot_jitter_scatter: Visualizes the alignment of individual data points.
+- run_visual_analysis: An orchestrator function that runs the full analysis and
+ returns a summary report along with the generated figures.
+"""
+
+import ast
+import logging
+import os
+from typing import Any
+
+import numpy as np
+import pandas as pd
+import plotly.graph_objects as go
+from dotenv import load_dotenv
+from plotly.subplots import make_subplots
+
+load_dotenv("src/.env")
+
+
+logging.basicConfig(
+ level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+
+def format_user_content(user_content: str, tokenizer: Any, **kwargs: Any) -> str | None:
+ """Applies tokenizer.apply_chat_template to user content string.
+ Assumes user_content is the text for the 'user' role.
+ """
+ message = [
+ {"role": "user", "content": user_content},
+ ]
+ kwargs.setdefault("tokenize", False)
+ kwargs.setdefault("add_generation_prompt", True)
+
+ try:
+ return tokenizer.apply_chat_template(message, **kwargs)
+ except Exception as e:
+ print(
+ f"Error applying chat template to content: '{user_content[:50]}...'. Error: {e}"
+ )
+ return None
+
+
+def prepare_dataframe(
+ raw_data: Any,
+ human_col_name: str,
+ score_col_name: str,
+) -> pd.DataFrame:
+ """Prepares the DataFrame from raw data, renames columns, and converts types."""
+ try:
+ if isinstance(raw_data, pd.DataFrame):
+ df = raw_data.copy()
+ if len(df.columns) >= 2:
+ # If it's a DataFrame, check if desired columns exist, else use first two
+ if human_col_name not in df.columns or score_col_name not in df.columns:
+ original_cols = df.columns
+ df = df[
+ [original_cols[0], original_cols[1]]
+ ].copy() # Use first two
+ df.columns = [human_col_name, score_col_name]
+ else:
+ # Ensure we only keep the needed columns if more exist
+ pass
+ if human_col_name not in df.columns or score_col_name not in df.columns:
+ original_cols = df.columns
+ df = df[[original_cols[0], original_cols[1]]].copy()
+ df.columns = [human_col_name, score_col_name]
+ else:
+ df = df[[human_col_name, score_col_name]].copy()
+ else:
+ print(
+ f"Warning: Input DataFrame has < 2 columns. Expected at least '{human_col_name}' and '{score_col_name}'."
+ )
+ return pd.DataFrame(columns=[human_col_name, score_col_name])
+ else:
+ df = pd.DataFrame(raw_data)
+ if df.empty:
+ print("Warning: Created empty DataFrame from raw_data.")
+ return pd.DataFrame(columns=[human_col_name, score_col_name])
+
+ if len(df.columns) >= 2:
+ df = df.iloc[:, :2]
+ df.columns = [human_col_name, score_col_name]
+ else:
+ print(
+ f"Warning: DataFrame from raw_data has < 2 columns. Cannot set '{human_col_name}' and '{score_col_name}'."
+ )
+ return pd.DataFrame(columns=[human_col_name, score_col_name])
+
+ df[human_col_name] = pd.to_numeric(df[human_col_name], errors="coerce")
+ df[score_col_name] = pd.to_numeric(df[score_col_name], errors="coerce")
+ df.dropna(subset=[human_col_name, score_col_name], inplace=True)
+ df[human_col_name] = df[human_col_name].astype(float)
+ df[score_col_name] = df[score_col_name].astype(float)
+ return df
+
+ except Exception as e:
+ print(f"Error preparing DataFrame: {e}")
+ return pd.DataFrame(columns=[human_col_name, score_col_name])
+
+
+def extract_completeness_metrics(
+ metrics_data: list[dict[str, Any]] | None,
+) -> tuple[list[list[Any]] | None, list[str] | None, list[float] | None]:
+ """Extracts confusion matrix info from metrics data (expected at index 0)."""
+ if not metrics_data or not isinstance(metrics_data, list) or len(metrics_data) == 0:
+ print("Warning: Metrics data is empty or not a list.")
+ return None, None, None
+
+ try:
+ completeness_metrics = metrics_data[0]
+ if (
+ "confusion_matrix" not in completeness_metrics
+ or "confusion_matrix_labels" not in completeness_metrics
+ ):
+ print(
+ "Warning: 'confusion_matrix' or 'confusion_matrix_labels' not in first metrics item."
+ )
+ return None, None, None
+
+ cm = completeness_metrics["confusion_matrix"]
+ cm_labels = completeness_metrics["confusion_matrix_labels"]
+ cm_labels_numeric = [float(cl) for cl in cm_labels]
+ return cm, cm_labels, cm_labels_numeric
+ except (KeyError, IndexError, ValueError, TypeError) as e:
+ print(f"Warning: Could not extract confusion matrix info from metrics: {e}")
+ return None, None, None
+
+
+def plot_distribution_comparison(
+ df: pd.DataFrame,
+ human_col: str,
+ score_col: str,
+) -> go.Figure:
+ """Generates bar charts comparing distributions of human ratings and model scores."""
+ human_counts = df[human_col].value_counts().sort_index()
+ score_counts = df[score_col].value_counts().sort_index()
+
+ fig = make_subplots(
+ rows=1,
+ cols=2,
+ subplot_titles=("Human Rating Distribution", "Model Score Distribution"),
+ )
+
+ fig.add_trace(
+ go.Bar(
+ x=human_counts.index,
+ y=human_counts.values,
+ name="Human Rating",
+ marker_color="indianred",
+ ),
+ row=1,
+ col=1,
+ )
+
+ fig.add_trace(
+ go.Bar(
+ x=score_counts.index,
+ y=score_counts.values,
+ name="Model Score",
+ marker_color="lightsalmon",
+ ),
+ row=1,
+ col=2,
+ )
+
+ fig.update_layout(
+ title_text="Distribution of Human Ratings vs. Model Scores",
+ bargap=0.2,
+ xaxis1_title="Rating Value",
+ yaxis1_title="Count",
+ xaxis2_title="Score Value",
+ yaxis2_title="Count",
+ xaxis1_type="category",
+ xaxis2_type="category",
+ xaxis1={
+ "categoryorder": "array",
+ "categoryarray": sorted(human_counts.index.unique()),
+ },
+ xaxis2={
+ "categoryorder": "array",
+ "categoryarray": sorted(score_counts.index.unique()),
+ },
+ height=400,
+ )
+ return fig
+
+
+def plot_confusion_matrix(
+ cm: list[list[Any]] | None, cm_labels: list[str] | None
+) -> go.Figure | None:
+ """Generates a heatmap for the confusion matrix."""
+ if cm is None or cm_labels is None:
+ print("Skipping confusion matrix plot: missing data.")
+ return None
+
+ fig = go.Figure(
+ data=go.Heatmap(
+ z=cm,
+ x=cm_labels,
+ y=cm_labels,
+ hoverongaps=False,
+ colorscale="Blues",
+ text=cm,
+ texttemplate="%{text}",
+ zmin=0,
+ )
+ )
+
+ fig.update_layout(
+ title="Confusion Matrix: Human Rating vs. Model Score (Completeness)",
+ xaxis_title="Predicted (Model Score)",
+ yaxis_title="True (Human Rating)",
+ yaxis={
+ "type": "category",
+ "categoryorder": "array",
+ "categoryarray": cm_labels,
+ },
+ xaxis={
+ "type": "category",
+ "categoryorder": "array",
+ "categoryarray": cm_labels,
+ },
+ height=600,
+ width=600,
+ )
+ return fig
+
+
+def plot_jitter_scatter(
+ df: pd.DataFrame,
+ cm_labels: list[str] | None,
+ cm_labels_numeric: list[float] | None,
+ human_col: str,
+ score_col: str,
+) -> go.Figure:
+ """Generates a jitter scatter plot comparing individual scores and ratings."""
+ df_jitter = df[[human_col, score_col]].copy()
+
+ if cm_labels is None or cm_labels_numeric is None:
+ print(
+ "Using data range for jitter plot axes due to missing confusion matrix labels."
+ )
+ min_val: float = (
+ min(df_jitter[human_col].min(), df_jitter[score_col].min()) - 0.5
+ )
+ max_val: float = (
+ max(df_jitter[human_col].max(), df_jitter[score_col].max()) + 0.5
+ )
+ plot_range = [min_val, max_val]
+ tick_vals = sorted(
+ df_jitter[human_col].unique()
+ ) # Use unique human ratings for ticks if available
+ tick_text = [str(int(v)) if v == int(v) else str(v) for v in tick_vals]
+ else:
+ plot_range = [min(cm_labels_numeric) - 0.5, max(cm_labels_numeric) + 0.5]
+ tick_vals = cm_labels_numeric
+ tick_text = cm_labels
+
+ df_jitter[f"{human_col}_jitter"] = df_jitter[human_col] + np.random.uniform(
+ -ast.literal_eval(os.getenv("JITTER_AMOUNT")),
+ ast.literal_eval(os.getenv("JITTER_AMOUNT")),
+ size=len(df_jitter),
+ )
+ df_jitter[f"{score_col}_jitter"] = df_jitter[score_col] + np.random.uniform(
+ -ast.literal_eval(os.getenv("JITTER_AMOUNT")),
+ ast.literal_eval(os.getenv("JITTER_AMOUNT")),
+ size=len(df_jitter),
+ )
+
+ fig = go.Figure()
+
+ fig.add_trace(
+ go.Scatter(
+ x=df_jitter[f"{score_col}_jitter"],
+ y=df_jitter[f"{human_col}_jitter"],
+ mode="markers",
+ marker={
+ "color": "rgba(0, 100, 200, 0.7)",
+ "size": 10,
+ "line": {"width": 1, "color": "DarkSlateGrey"},
+ },
+ text=[
+ f"HR: {hr:.1f}, Score: {s:.1f}"
+ for hr, s in zip(
+ df_jitter[human_col], df_jitter[score_col], strict=False
+ )
+ ],
+ hoverinfo="text",
+ name="Ratings",
+ )
+ )
+
+ fig.add_trace(
+ go.Scatter(
+ x=plot_range,
+ y=plot_range,
+ mode="lines",
+ name="Ideal Alignment (Score = Human Rating)",
+ line={"color": "red", "dash": "dash"},
+ )
+ )
+
+ fig.update_layout(
+ title="Model Score vs. Human Rating (with Jitter)",
+ xaxis_title="Model Score (Jittered)",
+ yaxis_title="Human Rating (Jittered)",
+ xaxis={"range": plot_range, "tickvals": list(tick_vals), "ticktext": tick_text},
+ yaxis={"range": plot_range, "tickvals": list(tick_vals), "ticktext": tick_text},
+ width=600,
+ height=600,
+ showlegend=True,
+ hovermode="closest",
+ )
+ return fig
+
+
+def run_visual_analysis(
+ df_data: Any,
+ metrics: list[dict[str, Any]] | None,
+ human_col_name: str,
+ score_col_name: str,
+) -> tuple[str, go.Figure | None, go.Figure | None, go.Figure | None]:
+ """Runs the visual analysis comparing model scores and human ratings.
+ Returns a markdown string and the three figure objects.
+ """
+ result_str = ""
+
+ df = prepare_dataframe(
+ df_data, human_col_name=human_col_name, score_col_name=score_col_name
+ )
+ cm, cm_labels, cm_labels_numeric = extract_completeness_metrics(metrics)
+
+ if df.empty:
+ result_str = f"Could not create valid DataFrame with columns '{human_col_name}' and '{score_col_name}' from input data. Stopping analysis."
+ return result_str, None, None, None
+
+ if human_col_name not in df.columns or score_col_name not in df.columns:
+ result_str = f"Expected columns '{human_col_name}' and '{score_col_name}' not found in DataFrame. Stopping analysis."
+ return result_str, None, None, None
+
+ result_str += "# Visual Analysis: Model Score vs. Human Rating Alignment \n"
+
+ result_str += "## 1. Distributions Comparison\n"
+ fig_dist: go.Figure | None = plot_distribution_comparison(
+ df, human_col=human_col_name, score_col=score_col_name
+ )
+ if fig_dist:
+ result_str += "*Shows if the model score distribution mirrors human ratings.*\n"
+ else:
+ result_str += "*Could not generate distribution comparison plot.*\n"
+
+ result_str += "## 2. Confusion Matrix (Human vs. Model)\n"
+ fig_cm: go.Figure | None = plot_confusion_matrix(cm, cm_labels)
+ if fig_cm:
+ result_str += "*Visualizes agreement and disagreement between discrete human ratings and model scores. Ideal alignment is along the diagonal.*\n"
+
+ else:
+ result_str += "*Confusion Matrix could not be generated (check metrics data or ensure it's at index 0).**\n"
+
+ result_str += "## 3. Score vs. Rating Alignment (Jitter Plot)\n"
+ fig_scatter: go.Figure | None = plot_jitter_scatter(
+ df,
+ cm_labels,
+ cm_labels_numeric,
+ human_col=human_col_name,
+ score_col=score_col_name,
+ )
+ if fig_scatter:
+ result_str += "*Shows individual item alignment. Points close to the red dashed line indicate good agreement.*\n"
+ else:
+ result_str += "*Could not generate jitter scatter plot.*"
+
+ return result_str, fig_dist, fig_cm, fig_scatter
diff --git a/tools/llmevalkit/src/gcp_prompt.py b/tools/llmevalkit/src/gcp_prompt.py
new file mode 100644
index 00000000000..578741ed662
--- /dev/null
+++ b/tools/llmevalkit/src/gcp_prompt.py
@@ -0,0 +1,273 @@
+# Copyright 2025 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
+#
+# http://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.
+
+"""Manages the lifecycle of prompts using Google Cloud Platform services.
+
+This module provides a `gcp_prompt` class that facilitates the creation,
+storage, retrieval, and execution of prompts with Vertex AI and Cloud Storage.
+It handles the interaction with the Vertex AI SDK to manage prompt versions
+and uses Cloud Storage to persist prompt metadata.
+
+Key functionalities include:
+- Initializing a connection to Google Cloud Platform services (Vertex AI, Cloud Storage).
+- Caching and refreshing a list of existing prompts.
+- Saving new or updated prompts, including their metadata, to both the
+ Vertex AI prompt registry and a Cloud Storage bucket.
+- Loading existing prompts from the registry and their associated metadata
+ from the bucket.
+- Generating responses from a specified model using a loaded prompt and
+ a given set of variables.
+- Helper functions for escaping special characters in prompt text.
+
+The `model_response` Pydantic model defines the expected structure of the
+response from the generative model.
+"""
+
+import json
+import logging
+import os
+from typing import Any
+
+from dotenv import load_dotenv
+from google import genai
+from google.cloud import aiplatform, storage
+from vertexai.generative_models import GenerationConfig
+from vertexai.preview import prompts
+from vertexai.preview.prompts import Prompt
+
+load_dotenv("src/.env")
+
+# Configure logging to the console
+logging.basicConfig(
+ level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+)
+logger = logging.getLogger(__name__)
+
+aiplatform.init(
+ project=os.getenv("PROJECT_ID"),
+ location=os.getenv("LOCATION"),
+ staging_bucket=os.getenv("BUCKET"),
+)
+
+# --- Constants ---
+PROMPT_PREFIX = os.getenv("PROMPT_PREFIX", "prompts_meta")
+
+
+class GcpPrompt:
+ """A wrapper class for the Vertex AI Prompt Management service."""
+
+ def __init__(self) -> None:
+ """Initializes the GcpPrompt client."""
+ self.storage_client = storage.Client()
+
+ self.refresh_prompt_cache()
+ logger.info("Found %d existing prompts.", len(self.existing_prompts))
+
+ # Load a default/initial state for the prompt metadata.
+ self.prompt_meta: dict[str, Any] = {}
+
+ self.prompt_to_run: Prompt = Prompt()
+ self.refresh_bucket_cache()
+
+ def refresh_prompt_cache(self) -> None:
+ """Refreshes the local cache of existing prompts from the service."""
+ self.existing_prompts = {p.display_name: p.prompt_id for p in prompts.list()}
+
+ def refresh_bucket_cache(self) -> None:
+ """Refreshes the list of metadata files from GCS."""
+ blobs = self.storage_client.list_blobs(
+ os.getenv("BUCKET"), prefix=PROMPT_PREFIX
+ )
+ self.bucket_cache_index = [i.name for i in blobs]
+ logger.debug("Found %d metadata files in GCS.", len(self.bucket_cache_index))
+
+ def _get_metadata_blob_name(self) -> str:
+ """Constructs the GCS blob name for the prompt metadata file."""
+ return (
+ f"{PROMPT_PREFIX}/{self.prompt_to_run.prompt_name}_"
+ f"{self.prompt_to_run._prompt_name}_{self.prompt_to_run._version_id}_"
+ f"{self.prompt_to_run._version_name}.json"
+ )
+
+ def save_prompt(self, check_existing: bool = False) -> str:
+ """Saves the current prompt.
+
+ If check_existing is True, it ensures a prompt with the same name does not
+ already exist before creating version "1".
+ If check_existing is False, it creates a new version of an existing prompt.
+
+ Args:
+ check_existing: If True, raises an error if a prompt with the same
+ display name already exists.
+
+ Returns:
+ A string with the details of the saved prompt version.
+ """
+ logger.info("Attempting to save prompt: %s", self.prompt_to_run.prompt_name)
+
+ if check_existing and self.prompt_to_run.prompt_name in self.existing_prompts:
+ raise ValueError(
+ f"Prompt with name '{self.prompt_to_run.prompt_name}' already exists. "
+ "To create a new version, load the prompt and save from the "
+ "'Existing Prompt' page."
+ )
+
+ if "generation_config" in self.prompt_meta and isinstance(
+ self.prompt_meta["generation_config"], dict
+ ):
+ self.prompt_to_run.generation_config = GenerationConfig(
+ **self.prompt_meta["generation_config"]
+ )
+ elif not isinstance(self.prompt_to_run.generation_config, GenerationConfig):
+ logger.warning("No valid generation_config found. Using default.")
+ self.prompt_to_run.generation_config = GenerationConfig()
+
+ logger.debug("Prompt object being sent to SDK: %s", self.prompt_to_run)
+
+ self.prompt_to_run = prompts.create_version(prompt=self.prompt_to_run)
+
+ sdk_details = (
+ f"SDK Prompt Version Details received:\n"
+ f"- Resource Name: {self.prompt_to_run._prompt_name}\n"
+ f"- Version ID: {self.prompt_to_run._version_id}\n"
+ f"- Version Name: {self.prompt_to_run._version_name}\n"
+ f"- Display Name: {self.prompt_to_run.prompt_name}"
+ )
+ logger.info(sdk_details)
+
+ self.prompt_meta["name"] = self.prompt_to_run.prompt_name
+ self.write_to_bucket()
+
+ return sdk_details
+
+ def load_prompt(self, prompt_id: str, prompt_name: str, version_id: str) -> None:
+ """Loads a specific version of a prompt and its associated metadata."""
+ self.prompt_to_run = prompts.get(prompt_id, version_id)
+ self.prompt_to_run.prompt_name = prompt_name
+
+ blob_name = self._get_metadata_blob_name()
+
+ none_blob_name = (
+ f"{PROMPT_PREFIX}/{None}_{None}_"
+ f"{self.prompt_to_run._version_id}_{self.prompt_to_run._version_name}.json"
+ )
+
+ self.refresh_bucket_cache()
+ bucket = self.storage_client.bucket(os.getenv("BUCKET"))
+
+ if blob_name in self.bucket_cache_index:
+ blob = bucket.blob(blob_name)
+ self.prompt_meta = json.loads(blob.download_as_string())
+ logger.info("Loaded metadata from %s", blob_name)
+ elif none_blob_name in self.bucket_cache_index:
+ logger.warning(
+ "Metadata not found at %s, using fallback %s", blob_name, none_blob_name
+ )
+ blob = bucket.blob(none_blob_name)
+ self.prompt_meta = json.loads(blob.download_as_string())
+
+ self.write_to_bucket()
+ else:
+ raise FileNotFoundError(
+ f"Could not find metadata file in GCS for prompt '{prompt_name}' "
+ f"version '{version_id}'. Looked for '{blob_name}' and '{none_blob_name}'."
+ )
+
+ def write_to_bucket(self) -> None:
+ """Writes the current prompt_meta to a GCS blob."""
+ blob_name = self._get_metadata_blob_name()
+ bucket = self.storage_client.bucket(os.getenv("BUCKET"))
+ blob = bucket.blob(blob_name)
+ blob.upload_from_string(
+ json.dumps(self.prompt_meta, indent=2), content_type="application/json"
+ )
+ logger.info("Wrote metadata to gs://%s/%s", os.getenv("BUCKET"), blob_name)
+
+ def generate_response(self, variables: dict[str, Any]) -> str | None:
+ """Generates a response from the currently loaded prompt and variables,
+ handling image uploads.
+ """
+ if not self.prompt_to_run.prompt_data:
+ raise ValueError("Prompt data is not loaded. Cannot generate response.")
+
+ client = genai.Client(
+ vertexai=True,
+ project=os.getenv("PROJECT_ID"),
+ location=os.getenv("LOCATION"),
+ )
+ updated_contents: list[Any] = []
+ for key, value in variables.items():
+ logger.info("Iterating through variables")
+ if key == "image":
+ if "jpg" in value or "jpeg" in value:
+ mime_type = "image/jpeg"
+ elif "png" in value:
+ mime_type = "image/png"
+ else:
+ logger.warning(
+ f"Unsupported image type for {value}, attempting to process as is."
+ )
+ updated_contents.append(value)
+ continue
+
+ try:
+ # Create Image Part from URI
+ updated_contents.append(
+ genai.Part.from_uri(mime_type=mime_type, uri=value)
+ )
+ except Exception as e:
+ logger.error(
+ "Error creating Part from URI for variable '%s': %s",
+ key,
+ e,
+ exc_info=True,
+ )
+ updated_contents.append(value)
+ else:
+ updated_contents.append(value)
+
+ # Append the prompt text if it's not already included in variables
+ prompt_text = self.prompt_to_run.prompt_data
+ if not any(
+ isinstance(item, str) and item == prompt_text for item in updated_contents
+ ):
+ updated_contents.append(prompt_text)
+
+ logger.info("Generating response with contents: %s", updated_contents)
+
+ model = self.prompt_to_run.model_name
+ response = client.models.generate_content(
+ model=model, contents=updated_contents
+ )
+
+ return response.text if response and response.text else None
+
+ def get_generation_config_dict(self):
+ """Returns the generation configuration as a dictionary."""
+ return self.prompt_to_run.generation_config.to_dict()
+
+
+def escape_special_characters(text: str) -> str:
+ """Escapes special characters for embedding in a string.
+
+ Note: This function is also present in gcp_dataset.py and could be moved
+ to a shared utility module in the future.
+ """
+ if not isinstance(text, str):
+ return text
+ text = text.replace("\\", "\\\\")
+ text = text.replace("\n", "\\n")
+ text = text.replace("\r", "\\r")
+ text = text.replace("\t", "\\t")
+ return text.replace('"', '\\"')
diff --git a/tools/llmevalkit/src/vapo_lib.py b/tools/llmevalkit/src/vapo_lib.py
new file mode 100644
index 00000000000..0eb512f6e38
--- /dev/null
+++ b/tools/llmevalkit/src/vapo_lib.py
@@ -0,0 +1,20 @@
+# 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.
+
+"""Utility functions and classes for the VAPO notebook.
+
+This file imports all the functions and classes from the original vapo_lib.py file.
+"""
+
+from gemini.prompts.prompt_optimizer.vapo_lib import * # noqa: F403