Few-shot prompt optimizer

The few-shot prompt optimizer helps you refine system instructions by analyzing examples where a model's response didn't meet expectations. You provide specific examples of prompts, model responses, and feedback on those responses, and the optimizer improves the system instructions. You can provide feedback in one of two ways: rubrics and rubric-based evaluation or target responses.

Few-shot optimization requires you to provide examples in a Pandas DataFrame and is available through the Vertex AI SDK.

Optimize using rubrics-based evaluation

This method is for scenarios where you have rubrics (also known as grading criteria) for each example and a record of whether the model met those criteria. Your DataFrame must contain the following four columns:

  • prompt: The original user prompt.
  • model_response: The actual output generated by the model.
  • rubrics: A list of specific criteria for that example, as a string representation of a list of strings.
  • rubrics_evaluations: A list of booleans indicating if each rubric was met, as a string representation of a list of booleans.

The rubrics can be consistent across all rows or tailored to each example.

import pandas as pd
import vertexai
from vertexai import types

PROJECT_NAME = "YOUR_PROJECT_NAME"
LOCATION = "YOUR_LOCATION" # e.g. "us-central1"
client = vertexai.Client(project=PROJECT_NAME, location=LOCATION)

system_instructions = "You are an AI assistant skilled in analyzing articles. Your task is to extract key arguments, evaluate evidence quality, identify potential biases, and summarize core findings into a concise, actionable, and objective report. Make sure the response is less than 50 words"
df = pd.DataFrame(
    [
        {
            "prompt": "prompt1",
            "model_response": "response1",
            "rubrics": "['Response is in English', 'Under 50 words']",
            "rubrics_evaluations": "[True, True]",
        },
        {
            "prompt": "prompt2",
            "model_response": "response2",
            "rubrics": "['Response is in English', 'Under 50 words']",
            "rubrics_evaluations": "[True, False]",
        },
    ]
)

config = vertexai.types.OptimizeConfig(
    optimization_target=vertexai.types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_RUBRICS,
    examples_dataframe=df,
)

response = client.prompts.optimize(prompt=system_instructions, config=config)

print(response.parsed_response.suggested_prompt)

Optimize using target responses

If you don't have specific rubrics but have an ideal or "gold standard" response for each prompt, you can use target response-based optimization. The optimizer compares the model's output against the ideal output and adjusts the instructions to bridge the gap.

Your DataFrame must include these three columns:

  • prompt: The original user instruction.
  • model_response: The current model output that needs improvement.
  • target_response: The ideal response you want the model to generate.
import pandas as pd
import vertexai
from vertexai import types

PROJECT_NAME = "PROJECT"
LOCATION = "YOUR_LOCATION" # e.g. "us-central1"
client = vertexai.Client(project=YOUR_PROJECT_NAME, location=LOCATION)

system_instructions = "You are an AI assistant skilled in analyzing articles. Your task is to extract key arguments, evaluate evidence quality, identify potential biases, and summarize core findings into a concise, actionable, and objective report. Make sure the response is less than 50 words"
df = pd.DataFrame(
    {
        "prompt": ["prompt1", "prompt2"],
        "model_response": ["response1", "response2"],
        "target_response": ["target1", "target2"],
    }
)
config = vertexai.types.OptimizeConfig(
    optimization_target=vertexai.types.OptimizeTarget.OPTIMIZATION_TARGET_FEW_SHOT_TARGET_RESPONSE,
    examples_dataframe=df,
)
response = client.prompts.optimize(
    prompt=system_instructions,
    config=config,
)

print(response.parsed_response.suggested_prompt)