Skip to content

Callbacks

Callback classes for logging, saving, and tracking optimization progress.

BestPromptCallback

Bases: Callback

Callback for tracking the best prompt during optimization.

This callback keeps track of the prompt with the highest score.

Attributes:

Name Type Description
best_prompt str

The prompt with the highest score so far.
best_score float

The highest score achieved so far.
Source code in promptolution/callbacks.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
class BestPromptCallback(Callback):
    """Callback for tracking the best prompt during optimization.

    This callback keeps track of the prompt with the highest score.

    Attributes:
        best_prompt (str): The prompt with the highest score so far.
        best_score (float): The highest score achieved so far.
    """

    def __init__(self):
        """Initialize the BestPromptCallback."""
        self.best_prompt = ""
        self.best_score = -99999

    def on_step_end(self, optimizer):
        """Update the best prompt and score if a new high score is achieved.

        Args:
        optimizer: The optimizer object that called the callback.
        """
        if optimizer.scores[0] > self.best_score:
            self.best_score = optimizer.scores[0]
            self.best_prompt = optimizer.prompts[0]

        return True

    def get_best_prompt(self):
        """Get the best prompt and score achieved during optimization.

        Returns:
        Tuple[str, float]: The best prompt and score.
        """
        return self.best_prompt, self.best_score

__init__()

Initialize the BestPromptCallback.

Source code in promptolution/callbacks.py 
183
184
185
186
def __init__(self):
    """Initialize the BestPromptCallback."""
    self.best_prompt = ""
    self.best_score = -99999

get_best_prompt()

Get the best prompt and score achieved during optimization.

Returns: Tuple[str, float]: The best prompt and score.

Source code in promptolution/callbacks.py 
200
201
202
203
204
205
206
def get_best_prompt(self):
    """Get the best prompt and score achieved during optimization.

    Returns:
    Tuple[str, float]: The best prompt and score.
    """
    return self.best_prompt, self.best_score

on_step_end(optimizer)

Update the best prompt and score if a new high score is achieved.

Args: optimizer: The optimizer object that called the callback.

Source code in promptolution/callbacks.py 
188
189
190
191
192
193
194
195
196
197
198
def on_step_end(self, optimizer):
    """Update the best prompt and score if a new high score is achieved.

    Args:
    optimizer: The optimizer object that called the callback.
    """
    if optimizer.scores[0] > self.best_score:
        self.best_score = optimizer.scores[0]
        self.best_prompt = optimizer.prompts[0]

    return True

CSVCallback

Bases: Callback

Callback for saving optimization progress to a CSV file.

This callback saves prompts and scores at each step to a CSV file.

Attributes:

Name Type Description
dir str

Directory the CSV file is saved to.
step int

The current step number.
Source code in promptolution/callbacks.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
class CSVCallback(Callback):
    """Callback for saving optimization progress to a CSV file.

    This callback saves prompts and scores at each step to a CSV file.

    Attributes:
        dir (str): Directory the CSV file is saved to.
        step (int): The current step number.
    """

    def __init__(self, dir):
        """Initialize the CSVCallback.

        Args:
        dir (str): Directory the CSV file is saved to.
        """
        if not os.path.exists(dir):
            os.makedirs(dir)

        self.dir = dir
        self.dir = dir
        self.step = 0
        self.input_tokens = 0
        self.output_tokens = 0
        self.start_time = datetime.now()
        self.step_time = datetime.now()

    def on_step_end(self, optimizer):
        """Save prompts and scores to csv.

        Args:
        optimizer: The optimizer object that called the callback
        """
        self.step += 1
        df = pd.DataFrame(
            {
                "step": [self.step] * len(optimizer.prompts),
                "input_tokens": [optimizer.meta_llm.input_token_count - self.input_tokens] * len(optimizer.prompts),
                "output_tokens": [optimizer.meta_llm.output_token_count - self.output_tokens] * len(optimizer.prompts),
                "time_elapsed": [(datetime.now() - self.step_time).total_seconds()] * len(optimizer.prompts),
                "score": optimizer.scores,
                "prompt": optimizer.prompts,
            }
        )
        self.step_time = datetime.now()
        self.input_tokens = optimizer.meta_llm.input_token_count
        self.output_tokens = optimizer.meta_llm.output_token_count

        if not os.path.exists(self.dir + "step_results.csv"):
            df.to_csv(self.dir + "step_results.csv", index=False)
        else:
            df.to_csv(self.dir + "step_results.csv", mode="a", header=False, index=False)

        return True

    def on_train_end(self, optimizer):
        """Called at the end of training.

        Args:
        optimizer: The optimizer object that called the callback.
        """
        df = pd.DataFrame(
            dict(
                steps=self.step,
                input_tokens=optimizer.meta_llm.input_token_count,
                output_tokens=optimizer.meta_llm.output_token_count,
                time_elapsed=(datetime.now() - self.start_time).total_seconds(),
                time=datetime.now(),
                score=np.array(optimizer.scores).mean(),
                best_prompts=str(optimizer.prompts),
            ),
            index=[0],
        )

        if not os.path.exists(self.dir + "train_results.csv"):
            df.to_csv(self.dir + "train_results.csv", index=False)
        else:
            df.to_csv(self.dir + "train_results.csv", mode="a", header=False, index=False)

        return True

__init__(dir)

Initialize the CSVCallback.

Args: dir (str): Directory the CSV file is saved to.

Source code in promptolution/callbacks.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def __init__(self, dir):
    """Initialize the CSVCallback.

    Args:
    dir (str): Directory the CSV file is saved to.
    """
    if not os.path.exists(dir):
        os.makedirs(dir)

    self.dir = dir
    self.dir = dir
    self.step = 0
    self.input_tokens = 0
    self.output_tokens = 0
    self.start_time = datetime.now()
    self.step_time = datetime.now()

on_step_end(optimizer)

Save prompts and scores to csv.

Args: optimizer: The optimizer object that called the callback

Source code in promptolution/callbacks.py 
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
def on_step_end(self, optimizer):
    """Save prompts and scores to csv.

    Args:
    optimizer: The optimizer object that called the callback
    """
    self.step += 1
    df = pd.DataFrame(
        {
            "step": [self.step] * len(optimizer.prompts),
            "input_tokens": [optimizer.meta_llm.input_token_count - self.input_tokens] * len(optimizer.prompts),
            "output_tokens": [optimizer.meta_llm.output_token_count - self.output_tokens] * len(optimizer.prompts),
            "time_elapsed": [(datetime.now() - self.step_time).total_seconds()] * len(optimizer.prompts),
            "score": optimizer.scores,
            "prompt": optimizer.prompts,
        }
    )
    self.step_time = datetime.now()
    self.input_tokens = optimizer.meta_llm.input_token_count
    self.output_tokens = optimizer.meta_llm.output_token_count

    if not os.path.exists(self.dir + "step_results.csv"):
        df.to_csv(self.dir + "step_results.csv", index=False)
    else:
        df.to_csv(self.dir + "step_results.csv", mode="a", header=False, index=False)

    return True

on_train_end(optimizer)

Called at the end of training.

Args: optimizer: The optimizer object that called the callback.

Source code in promptolution/callbacks.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
def on_train_end(self, optimizer):
    """Called at the end of training.

    Args:
    optimizer: The optimizer object that called the callback.
    """
    df = pd.DataFrame(
        dict(
            steps=self.step,
            input_tokens=optimizer.meta_llm.input_token_count,
            output_tokens=optimizer.meta_llm.output_token_count,
            time_elapsed=(datetime.now() - self.start_time).total_seconds(),
            time=datetime.now(),
            score=np.array(optimizer.scores).mean(),
            best_prompts=str(optimizer.prompts),
        ),
        index=[0],
    )

    if not os.path.exists(self.dir + "train_results.csv"):
        df.to_csv(self.dir + "train_results.csv", index=False)
    else:
        df.to_csv(self.dir + "train_results.csv", mode="a", header=False, index=False)

    return True

Callback

Base class for optimization callbacks.

Source code in promptolution/callbacks.py 
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
class Callback:
    """Base class for optimization callbacks."""

    def on_step_end(self, optimizer):
        """Called at the end of each optimization step.

        Args:
        optimizer: The optimizer object that called the callback.

        Returns:
            Bool: True if the optimization should continue, False if it should stop.
        """
        return True

    def on_epoch_end(self, optimizer):
        """Called at the end of each optimization epoch.

        Args:
        optimizer: The optimizer object that called the callback.

        Returns:
            Bool: True if the optimization should continue, False if it should stop.
        """
        return True

    def on_train_end(self, optimizer):
        """Called at the end of the entire optimization process.

        Args:
        optimizer: The optimizer object that called the callback.

        Returns:
            Bool: True if the optimization should continue, False if it should stop.
        """
        return True

on_epoch_end(optimizer)

Called at the end of each optimization epoch.

Args: optimizer: The optimizer object that called the callback.

Returns:

Name Type Description
Bool

True if the optimization should continue, False if it should stop.
Source code in promptolution/callbacks.py 
26
27
28
29
30
31
32
33
34
35
def on_epoch_end(self, optimizer):
    """Called at the end of each optimization epoch.

    Args:
    optimizer: The optimizer object that called the callback.

    Returns:
        Bool: True if the optimization should continue, False if it should stop.
    """
    return True

on_step_end(optimizer)

Called at the end of each optimization step.

Args: optimizer: The optimizer object that called the callback.

Returns:

Name Type Description
Bool

True if the optimization should continue, False if it should stop.
Source code in promptolution/callbacks.py 
15
16
17
18
19
20
21
22
23
24
def on_step_end(self, optimizer):
    """Called at the end of each optimization step.

    Args:
    optimizer: The optimizer object that called the callback.

    Returns:
        Bool: True if the optimization should continue, False if it should stop.
    """
    return True

on_train_end(optimizer)

Called at the end of the entire optimization process.

Args: optimizer: The optimizer object that called the callback.

Returns:

Name Type Description
Bool

True if the optimization should continue, False if it should stop.
Source code in promptolution/callbacks.py 
37
38
39
40
41
42
43
44
45
46
def on_train_end(self, optimizer):
    """Called at the end of the entire optimization process.

    Args:
    optimizer: The optimizer object that called the callback.

    Returns:
        Bool: True if the optimization should continue, False if it should stop.
    """
    return True

LoggerCallback

Bases: Callback

Callback for logging optimization progress.

This callback logs information about each step, epoch, and the end of training.

Attributes:

Name Type Description
logger

The logger object to use for logging.
step int

The current step number.
Source code in promptolution/callbacks.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
class LoggerCallback(Callback):
    """Callback for logging optimization progress.

    This callback logs information about each step, epoch, and the end of training.

    Attributes:
        logger: The logger object to use for logging.
        step (int): The current step number.
    """

    def __init__(self, logger):
        """Initialize the LoggerCallback."""
        self.logger = logger
        self.step = 0

    def on_step_end(self, optimizer):
        """Log information about the current step."""
        self.step += 1
        time = datetime.now().strftime("%d-%m-%y %H:%M:%S:%f")
        self.logger.critical(f"{time} - ✨Step {self.step} ended✨")
        for i, (prompt, score) in enumerate(zip(optimizer.prompts, optimizer.scores)):
            self.logger.critical(f"*** Prompt {i}: Score: {score}")
            self.logger.critical(f"{prompt}")

        return True

    def on_train_end(self, optimizer, logs=None):
        """Log information at the end of training.

        Args:
        optimizer: The optimizer object that called the callback.
        logs: Additional information to log.
        """
        time = datetime.now().strftime("%d-%m-%y %H:%M:%S:%f")
        if logs is None:
            self.logger.critical(f"{time} - Training ended")
        else:
            self.logger.critical(f"{time} - Training ended - {logs}")

        return True

__init__(logger)

Initialize the LoggerCallback.

Source code in promptolution/callbacks.py 
59
60
61
62
def __init__(self, logger):
    """Initialize the LoggerCallback."""
    self.logger = logger
    self.step = 0

on_step_end(optimizer)

Log information about the current step.

Source code in promptolution/callbacks.py 
64
65
66
67
68
69
70
71
72
73
def on_step_end(self, optimizer):
    """Log information about the current step."""
    self.step += 1
    time = datetime.now().strftime("%d-%m-%y %H:%M:%S:%f")
    self.logger.critical(f"{time} - ✨Step {self.step} ended✨")
    for i, (prompt, score) in enumerate(zip(optimizer.prompts, optimizer.scores)):
        self.logger.critical(f"*** Prompt {i}: Score: {score}")
        self.logger.critical(f"{prompt}")

    return True

on_train_end(optimizer, logs=None)

Log information at the end of training.

Args: optimizer: The optimizer object that called the callback. logs: Additional information to log.

Source code in promptolution/callbacks.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
def on_train_end(self, optimizer, logs=None):
    """Log information at the end of training.

    Args:
    optimizer: The optimizer object that called the callback.
    logs: Additional information to log.
    """
    time = datetime.now().strftime("%d-%m-%y %H:%M:%S:%f")
    if logs is None:
        self.logger.critical(f"{time} - Training ended")
    else:
        self.logger.critical(f"{time} - Training ended - {logs}")

    return True

ProgressBarCallback

Bases: Callback

Callback for displaying a progress bar during optimization.

This callback uses tqdm to display a progress bar that updates at each step.

Attributes:

Name Type Description
pbar tqdm

The tqdm progress bar object.
Source code in promptolution/callbacks.py 
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
class ProgressBarCallback(Callback):
    """Callback for displaying a progress bar during optimization.

    This callback uses tqdm to display a progress bar that updates at each step.

    Attributes:
        pbar (tqdm): The tqdm progress bar object.
    """

    def __init__(self, total_steps):
        """Initialize the ProgressBarCallback.

        Args:
        total_steps (int): The total number of steps in the optimization process.
        """
        self.pbar = tqdm(total=total_steps)

    def on_step_end(self, optimizer):
        """Update the progress bar at the end of each step.

        Args:
        optimizer: The optimizer object that called the callback.
        """
        self.pbar.update(1)

        return True

    def on_train_end(self, optimizer):
        """Close the progress bar at the end of training.

        Args:
        optimizer: The optimizer object that called the callback.
        """
        self.pbar.close()

        return True

__init__(total_steps)

Initialize the ProgressBarCallback.

Args: total_steps (int): The total number of steps in the optimization process.

Source code in promptolution/callbacks.py 
218
219
220
221
222
223
224
def __init__(self, total_steps):
    """Initialize the ProgressBarCallback.

    Args:
    total_steps (int): The total number of steps in the optimization process.
    """
    self.pbar = tqdm(total=total_steps)

on_step_end(optimizer)

Update the progress bar at the end of each step.

Args: optimizer: The optimizer object that called the callback.

Source code in promptolution/callbacks.py 
226
227
228
229
230
231
232
233
234
def on_step_end(self, optimizer):
    """Update the progress bar at the end of each step.

    Args:
    optimizer: The optimizer object that called the callback.
    """
    self.pbar.update(1)

    return True

on_train_end(optimizer)

Close the progress bar at the end of training.

Args: optimizer: The optimizer object that called the callback.

Source code in promptolution/callbacks.py 
236
237
238
239
240
241
242
243
244
def on_train_end(self, optimizer):
    """Close the progress bar at the end of training.

    Args:
    optimizer: The optimizer object that called the callback.
    """
    self.pbar.close()

    return True

TokenCountCallback

Bases: Callback

Callback for stopping optimization based on the total token count.

Source code in promptolution/callbacks.py 
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
class TokenCountCallback(Callback):
    """Callback for stopping optimization based on the total token count."""

    def __init__(
        self,
        max_tokens_for_termination: int,
        token_type_for_termination: Literal["input_tokens", "output_tokens", "total_tokens"],
    ):
        """Initialize the TokenCountCallback.

        Args:
        max_tokens_for_termination (int): Maximum number of tokens which is allowed befor the algorithm is stopped.
        token_type_for_termination (str): Can be one of either "input_tokens", "output_tokens" or "total_tokens".
        """
        self.max_tokens_for_termination = max_tokens_for_termination
        self.token_type_for_termination = token_type_for_termination

    def on_step_end(self, optimizer):
        """Check if the total token count exceeds the maximum allowed. If so, stop the optimization."""
        token_counts = optimizer.predictor.llm.get_token_count()

        if token_counts[self.token_type_for_termination] > self.max_tokens_for_termination:
            return False

        return True

__init__(max_tokens_for_termination, token_type_for_termination)

Initialize the TokenCountCallback.

Args: max_tokens_for_termination (int): Maximum number of tokens which is allowed befor the algorithm is stopped. token_type_for_termination (str): Can be one of either "input_tokens", "output_tokens" or "total_tokens".

Source code in promptolution/callbacks.py 
250
251
252
253
254
255
256
257
258
259
260
261
262
def __init__(
    self,
    max_tokens_for_termination: int,
    token_type_for_termination: Literal["input_tokens", "output_tokens", "total_tokens"],
):
    """Initialize the TokenCountCallback.

    Args:
    max_tokens_for_termination (int): Maximum number of tokens which is allowed befor the algorithm is stopped.
    token_type_for_termination (str): Can be one of either "input_tokens", "output_tokens" or "total_tokens".
    """
    self.max_tokens_for_termination = max_tokens_for_termination
    self.token_type_for_termination = token_type_for_termination

on_step_end(optimizer)

Check if the total token count exceeds the maximum allowed. If so, stop the optimization.

Source code in promptolution/callbacks.py 
264
265
266
267
268
269
270
271
def on_step_end(self, optimizer):
    """Check if the total token count exceeds the maximum allowed. If so, stop the optimization."""
    token_counts = optimizer.predictor.llm.get_token_count()

    if token_counts[self.token_type_for_termination] > self.max_tokens_for_termination:
        return False

    return True