Custom Datasets for Optimization

agent-opt expects a dataset as a plain Python list of dictionaries, one dict per row. Every source format you load ends up in this shape.

in_memory_dataset = [
    {
        "question": "What is the capital of France?",
        "context": "France is a country in Western Europe. Its capital and largest city is Paris.",
        "answer": "Paris",
    },
    {
        "question": "Who painted the Mona Lisa?",
        "context": "The Mona Lisa is a half-length portrait by the Italian artist Leonardo da Vinci.",
        "answer": "Leonardo da Vinci",
    },
]

You should see a list of two dicts, both sharing the same keys. Every row in a dataset must carry the same set of keys so the data mapper can address a column by name across the full dataset.

This puts loading, mapping, and optimizing together against the answer_similarity eval template.

from fi.opt.base import Evaluator
from fi.opt.datamappers import BasicDataMapper
from fi.opt.generators import LiteLLMGenerator
from fi.opt.optimizers import RandomSearchOptimizer

dataset = [
    {"question": "What is the capital of France?", "answer": "Paris"},
    {"question": "Who painted the Mona Lisa?", "answer": "Leonardo da Vinci"},
]

evaluator = Evaluator(
    eval_template="answer_similarity",
    eval_model_name="turing_flash",
)

data_mapper = BasicDataMapper(
    key_map={
        "response": "generated_output",
        "expected_response": "answer",
    }
)

initial_generator = LiteLLMGenerator(
    model="gpt-4o-mini",
    prompt_template="Q: {question}\nA:",
)

optimizer = RandomSearchOptimizer(
    generator=initial_generator,
    teacher_model="gpt-4o",
    num_variations=3,
)

result = optimizer.optimize(
    evaluator=evaluator,
    data_mapper=data_mapper,
    dataset=dataset,
)

print(f"Best prompt:\n{result.best_generator.get_prompt_template()}")
print(f"Final score: {result.final_score:.4f}")

You should see a best prompt printed along with a final score between 0 and 1. The exact wording and score vary by run since the optimizer samples variations.

For a data.csv file with a header row:

question,context,answer
"What is the capital of France?","France is a country in Western Europe.","Paris"
"Who painted the Mona Lisa?","A portrait by the Italian artist.","Leonardo da Vinci"

import pandas as pd

df = pd.read_csv("data.csv")
dataset_from_csv = df.to_dict(orient="records")

print(dataset_from_csv[0])

You should see:

{'question': 'What is the capital of France?', 'context': 'France is a country in Western Europe.', 'answer': 'Paris'}

to_dict(orient="records") is what turns the dataframe into the list-of-dicts shape agent-opt needs.

A data.json file holding a list of objects loads the same way:

import pandas as pd

df = pd.read_json("data.json", orient="records")
dataset_from_json = df.to_dict(orient="records")

A data.jsonl file, one JSON object per line, needs lines=True. Skipping it is the most common mistake, so trigger it on purpose first:

import pandas as pd

df = pd.read_json("data.jsonl")

You should see:

ValueError: Trailing data

Pandas tries to parse the whole file as a single JSON document and chokes on the second line. Add lines=True to parse it as JSON Lines instead:

import pandas as pd

df = pd.read_json("data.jsonl", lines=True)
dataset_from_jsonl = df.to_dict(orient="records")

You should see the same list-of-dicts shape as the CSV step above, this time with no error. If dataset_from_jsonl still comes back empty, double check the file has one JSON object per line rather than one array.

The Evaluator expects fixed input keys like response and expected_response. Your dataset’s column names rarely match those, so BasicDataMapper translates between the two with a key_map.

from fi.opt.datamappers import BasicDataMapper

# key_map = { evaluator's expected key: your dataset's column name }
data_mapper = BasicDataMapper(
    key_map={
        "response": "generated_output",   # reserved: the Generator's output
        "expected_response": "answer",     # your dataset's ground-truth column
    }
)

print(data_mapper.map("Paris", in_memory_dataset[0]))

You should see:

{'response': 'Paris', 'expected_response': 'Paris'}

generated_output is a reserved key: it always refers to the text the Generator under optimization produces, not a column that exists in your dataset file. Confirming the mapped output here is what tells you the key_map direction is right before it feeds into optimize().

Optimization scores every row for every candidate prompt, so a dataset of thousands of rows makes each trial slow and expensive. Draw a representative sample instead.

import random

import pandas as pd

df = pd.read_csv("large_dataset.csv")
full_dataset = df.to_dict(orient="records")

sample_size = 100
if len(full_dataset) > sample_size:
    optimization_dataset = random.sample(full_dataset, sample_size)
else:
    optimization_dataset = full_dataset

print(f"Using {len(optimization_dataset)} rows for optimization.")

You should see the sampled row count printed, capped at sample_size. 30 to 200 examples is enough signal for most optimizers without running up a large model bill. Pass optimization_dataset as the dataset argument in the optimize() call from Step 2.