wehub-resource-sync
457 lines
16 KiB
Python
457 lines
16 KiB
Python
# 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.5-pro",
|
|
"gemini-2.5-flash",
|
|
"gemini-2.0-flash-001",
|
|
"gemini-2.0-flash-lite-001",
|
|
"gemini-1.5-pro-002",
|
|
"gemini-1.5-flash-002",
|
|
]
|
|
|
|
|
|
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:
|
|
if not st.session_state.get("selected_dataset") or not st.session_state.get(
|
|
"selected_file_from_dataset"
|
|
):
|
|
st.warning("Please select a dataset and a file first.")
|
|
else:
|
|
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,
|
|
target_model,
|
|
target_qps=1.0,
|
|
optimizer_qps=1.0,
|
|
eval_qps=1.0,
|
|
data_limit=10,
|
|
):
|
|
"""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=target_model,
|
|
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=target_qps,
|
|
optimizer_model_qps=optimizer_qps,
|
|
eval_qps=eval_qps,
|
|
source_model_qps="",
|
|
response_mime_type=response_mime_type,
|
|
response_schema=response_schema_arg,
|
|
language="English",
|
|
placeholder_to_content=json.loads("{}"),
|
|
data_limit=data_limit,
|
|
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,
|
|
st.session_state.target_model_optimization,
|
|
st.session_state.target_qps,
|
|
st.session_state.optimizer_qps,
|
|
st.session_state.eval_qps,
|
|
st.session_state.data_limit,
|
|
)
|
|
|
|
with st.expander("Prompt Optimization 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",
|
|
)
|
|
|
|
with st.expander("Advanced Settings"):
|
|
st.session_state.target_qps = st.number_input(
|
|
"Target Model QPS", min_value=0.1, max_value=10.0, value=1.0, step=0.1
|
|
)
|
|
st.session_state.optimizer_qps = st.number_input(
|
|
"Optimizer Model QPS", min_value=0.1, max_value=10.0, value=1.0, step=0.1
|
|
)
|
|
st.session_state.eval_qps = st.number_input(
|
|
"Evaluation QPS", min_value=0.1, max_value=10.0, value=1.0, step=0.1
|
|
)
|
|
st.session_state.data_limit = st.number_input(
|
|
"Data Limit (Sample Size)", min_value=1, max_value=1000, value=10, step=1
|
|
)
|
|
|
|
prompt_selection()
|
|
dataset_selection()
|
|
start_optimization()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|