Skip to main content

Overview

GEPA integrates with Weights & Biases (W&B) to provide real-time experiment tracking, visualization, and collaboration features. The integration automatically logs metrics, parameters, and optimization progress to your W&B workspace.

Setup

Install Weights & Biases:
pip install wandb
Log in to your W&B account:
wandb login

Basic Usage

Enable W&B tracking by setting use_wandb=True in your optimization call:
import gepa

result = gepa.optimize(
    seed_candidate={"system_prompt": "You are a helpful assistant."},
    trainset=trainset,
    valset=valset,
    task_lm="openai/gpt-4o-mini",
    reflection_lm="openai/gpt-4o",
    max_metric_calls=100,
    use_wandb=True,  # Enable W&B tracking
)

Configuration Options

The gepa.optimize() function provides several W&B-specific parameters:

use_wandb

  • Type: bool
  • Default: False
  • Description: Enable Weights & Biases experiment tracking

wandb_api_key

  • Type: str | None
  • Default: None
  • Description: W&B API key for authentication. If not provided, uses the key from wandb login or environment variables.
result = gepa.optimize(
    ...,
    use_wandb=True,
    wandb_api_key="your-api-key-here",
)

wandb_init_kwargs

  • Type: dict[str, Any] | None
  • Default: None
  • Description: Additional keyword arguments passed to wandb.init(). Use this to configure project name, entity, tags, notes, and other W&B settings.
result = gepa.optimize(
    ...,
    use_wandb=True,
    wandb_init_kwargs={
        "project": "prompt-optimization",
        "entity": "my-team",
        "tags": ["gpt-4o", "math"],
        "notes": "Optimizing math problem solving",
        "name": "math-optimization-run-1",
    },
)

Complete Example

import gepa
import os

# Sample data
trainset = [
    {"question": "What is 2+2?", "answer": "4"},
    {"question": "What is 5*3?", "answer": "15"},
    {"question": "What is 12/4?", "answer": "3"},
]

valset = [
    {"question": "What is 10-3?", "answer": "7"},
    {"question": "What is 8*2?", "answer": "16"},
]

# Optimize with W&B tracking
result = gepa.optimize(
    seed_candidate={
        "system_prompt": "You are a math tutor. Answer questions clearly."
    },
    trainset=trainset,
    valset=valset,
    task_lm="openai/gpt-4o-mini",
    reflection_lm="openai/gpt-4o",
    max_metric_calls=50,
    use_wandb=True,
    wandb_init_kwargs={
        "project": "gepa-math-optimization",
        "tags": ["math", "gpt-4o-mini"],
        "notes": "Optimizing system prompt for math question answering",
    },
)

print("Best prompt:", result.best_candidate["system_prompt"])

Logged Metrics

GEPA automatically logs the following to W&B during optimization:
  • Validation scores: Performance metrics on the validation set
  • Training scores: Performance on training minibatches
  • Iteration count: Current optimization iteration number
  • Metric calls: Total number of evaluations performed
  • Best score: Highest validation score achieved so far
  • Candidate history: Evolution of candidates over time

Authentication Methods

W&B supports multiple authentication methods:

1. Interactive Login

wandb login

2. Environment Variable

export WANDB_API_KEY="your-api-key"

3. Programmatic Login

result = gepa.optimize(
    ...,
    use_wandb=True,
    wandb_api_key=os.environ["WANDB_API_KEY"],
)

4. Netrc File

Add credentials to ~/.netrc:
machine api.wandb.ai
  login user
  password your-api-key

Combined Logging

You can use both W&B and MLflow simultaneously:
result = gepa.optimize(
    ...,
    use_wandb=True,
    wandb_init_kwargs={"project": "my-project"},
    use_mlflow=True,
    mlflow_experiment_name="my-experiment",
)

Viewing Results

After running your optimization, view results in the W&B dashboard:
  1. Navigate to your W&B project at https://wandb.ai/<entity>/<project>
  2. Select your run to see detailed metrics, charts, and system information
  3. Compare multiple runs to analyze optimization strategies
  4. Share results with team members using W&B’s collaboration features

Advanced Features

Custom Metrics

If you need to log custom metrics beyond what GEPA provides automatically, you can access the W&B run object through callbacks:
from gepa.core.callbacks import GEPACallback
import wandb

class CustomLoggingCallback(GEPACallback):
    def on_iteration_end(self, state, **kwargs):
        # Log custom metrics
        wandb.log({"custom_metric": compute_custom_metric(state)})

result = gepa.optimize(
    ...,
    use_wandb=True,
    callbacks=[CustomLoggingCallback()],
)

External Resources

Weights & Biases Documentation

Official W&B documentation and guides

W&B Python Library Reference

Complete Python API reference for wandb

Build docs developers (and LLMs) love