Skip to content

Tasks

Module for task-related functions and classes.

ClassificationTask

Bases: BaseTask

A class representing a classification task in the promptolution library.

This class handles the loading and management of classification datasets, as well as the evaluation of predictors on these datasets.

Source code in promptolution/tasks/classification_tasks.py 
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
47
48
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
class ClassificationTask(BaseTask):
    """A class representing a classification task in the promptolution library.

    This class handles the loading and management of classification datasets,
    as well as the evaluation of predictors on these datasets.
    """

    def __init__(
        self,
        df: pd.DataFrame,
        task_description: Optional[str] = None,
        x_column: str = "x",
        y_column: str = "y",
        n_subsamples: int = 30,
        eval_strategy: Literal["full", "subsample", "sequential_block", "random_block"] = "full",
        seed: int = 42,
        metric: Callable[[Any, Any], float] = accuracy_score,
        config: Optional["ExperimentConfig"] = None,
    ) -> None:
        """Initialize the ClassificationTask from a pandas DataFrame.

        Args:
            df (pd.DataFrame): Input DataFrame containing the data
            task_description (str): Description of the task
            x_column (str, optional): Name of the column containing input texts. Defaults to "x".
            y_column (str, optional): Name of the column containing labels. Defaults to "y".
            n_subsamples (int, optional): Number of subsamples to use. No subsampling if None. Defaults to None.
            eval_strategy (str, optional): Subsampling strategy to use. Options:
                - "full": Uses the entire dataset for evaluation.
                - "evaluated": Uses only previously evaluated datapoints from the cache.
                - "subsample": Randomly selects n_subsamples datapoints without replacement.
                - "sequential_block": Uses a block of block_size consecutive datapoints, advancing through blocks sequentially.
                - "random_block": Randomly selects a block of block_size consecutive datapoints.
                Defaults to "full".
            seed (int, optional): Random seed for reproducibility. Defaults to 42.
            metric (Callable, optional): Metric to use for evaluation. Defaults to accuracy_score.
            config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
        """
        self.metric = metric
        super().__init__(
            df=df,
            x_column=x_column,
            y_column=y_column,
            task_description=task_description,
            n_subsamples=n_subsamples,
            eval_strategy=eval_strategy,
            seed=seed,
            config=config,
        )
        self.task_type = "classification"
        self.ys: List[str] = (
            df[self.y_column].str.lower().values.tolist()
        )  # Ensure y values are lowercase for consistent comparison
        self.classes = np.unique(self.ys)

    def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarray:
        """Calculate the score for a single prediction."""
        scores = [self.metric([y], [pred]) for pred, y in zip(preds, ys)]
        return np.asarray(scores, dtype=float)

__init__(df, task_description=None, x_column='x', y_column='y', n_subsamples=30, eval_strategy='full', seed=42, metric=accuracy_score, config=None)

Initialize the ClassificationTask from a pandas DataFrame.

Parameters:

Name Type Description Default
df DataFrame

Input DataFrame containing the data

 required
task_description str

Description of the task

 None
x_column str

Name of the column containing input texts. Defaults to "x".

 'x'
y_column str

Name of the column containing labels. Defaults to "y".

 'y'
n_subsamples int

Number of subsamples to use. No subsampling if None. Defaults to None.

 30
eval_strategy str

Subsampling strategy to use. Options: - "full": Uses the entire dataset for evaluation. - "evaluated": Uses only previously evaluated datapoints from the cache. - "subsample": Randomly selects n_subsamples datapoints without replacement. - "sequential_block": Uses a block of block_size consecutive datapoints, advancing through blocks sequentially. - "random_block": Randomly selects a block of block_size consecutive datapoints. Defaults to "full".

 'full'
seed int

Random seed for reproducibility. Defaults to 42.

 42
metric Callable

Metric to use for evaluation. Defaults to accuracy_score.

 accuracy_score
config ExperimentConfig

Configuration for the task, overriding defaults.

 None
Source code in promptolution/tasks/classification_tasks.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def __init__(
    self,
    df: pd.DataFrame,
    task_description: Optional[str] = None,
    x_column: str = "x",
    y_column: str = "y",
    n_subsamples: int = 30,
    eval_strategy: Literal["full", "subsample", "sequential_block", "random_block"] = "full",
    seed: int = 42,
    metric: Callable[[Any, Any], float] = accuracy_score,
    config: Optional["ExperimentConfig"] = None,
) -> None:
    """Initialize the ClassificationTask from a pandas DataFrame.

    Args:
        df (pd.DataFrame): Input DataFrame containing the data
        task_description (str): Description of the task
        x_column (str, optional): Name of the column containing input texts. Defaults to "x".
        y_column (str, optional): Name of the column containing labels. Defaults to "y".
        n_subsamples (int, optional): Number of subsamples to use. No subsampling if None. Defaults to None.
        eval_strategy (str, optional): Subsampling strategy to use. Options:
            - "full": Uses the entire dataset for evaluation.
            - "evaluated": Uses only previously evaluated datapoints from the cache.
            - "subsample": Randomly selects n_subsamples datapoints without replacement.
            - "sequential_block": Uses a block of block_size consecutive datapoints, advancing through blocks sequentially.
            - "random_block": Randomly selects a block of block_size consecutive datapoints.
            Defaults to "full".
        seed (int, optional): Random seed for reproducibility. Defaults to 42.
        metric (Callable, optional): Metric to use for evaluation. Defaults to accuracy_score.
        config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
    """
    self.metric = metric
    super().__init__(
        df=df,
        x_column=x_column,
        y_column=y_column,
        task_description=task_description,
        n_subsamples=n_subsamples,
        eval_strategy=eval_strategy,
        seed=seed,
        config=config,
    )
    self.task_type = "classification"
    self.ys: List[str] = (
        df[self.y_column].str.lower().values.tolist()
    )  # Ensure y values are lowercase for consistent comparison
    self.classes = np.unique(self.ys)

JudgeTask

Bases: BaseTask

Task that evaluates a predictor using an LLM-as-a-judge, optionally accepting a ground truth.

Source code in promptolution/tasks/judge_tasks.py 
 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
 89
 90
 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
class JudgeTask(BaseTask):
    """Task that evaluates a predictor using an LLM-as-a-judge, optionally accepting a ground truth."""

    def __init__(
        self,
        df: pd.DataFrame,
        judge_llm: "BaseLLM",
        x_column: str = "x",
        y_column: Optional[str] = None,
        task_description: Optional[str] = None,
        n_subsamples: int = 30,
        eval_strategy: "EvalStrategy" = "full",
        seed: int = 42,
        judge_prompt: Optional[str] = None,
        min_score: float = -5.0,
        max_score: float = 5.0,
        config: "ExperimentConfig" = None,
    ):
        """Initialize the JudgeTask.

        Args:
            df (pd.DataFrame): The input DataFrame containing the data.
            judge_llm (BaseLLM): The LLM judging the predictions.
            x_column (str): Name of the column containing input texts.
            y_column (Optional[str]): Name of the column containing labels/ground truth (if applicable).
            task_description (Optional[str]): Description of the task, parsed to the Judge-LLM and Meta-LLM.
            n_subsamples (int): Number of subsamples to use for evaluation.
            eval_strategy (EvalStrategy): Subsampling strategy to use for evaluation.
            seed (int): Random seed for reproducibility.
            judge_prompt (Optional[str]): Custom prompt for the judge. Note: The score of the Judge will be extracted inside <final_score> tags.
            min_score (float): Minimum score for evaluation.
            max_score (float): Maximum score for evaluation.
            config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
        """
        if judge_prompt is None:
            judge_prompt = JUDGE_PROMPT_WITH_GROUND_TRUTH if y_column else JUDGE_PROMPT_WITHOUT_GROUND_TRUTH
        self.judge_prompt = judge_prompt
        self.min_score = min_score
        self.max_score = max_score

        super().__init__(
            df=df,
            x_column=x_column,
            y_column=y_column,
            task_description=task_description,
            n_subsamples=n_subsamples,
            eval_strategy=eval_strategy,
            seed=seed,
            config=config,
        )
        self.judge_llm = judge_llm
        self.task_type = "judge"

    def _construct_judge_prompt(self, x: str, pred: str, y: Optional[str] = None) -> str:
        """Construct the judge prompt based on whether ground truth is available."""
        if y is not None:
            prompt = self.judge_prompt.replace("{ground_truth}", str(y))
        else:
            prompt = self.judge_prompt

        task_description = self.task_description or ""
        prompt = prompt.replace("{task}", task_description).replace("{input}", x).replace("{prediction}", pred)
        return prompt

    def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarray:
        """Calculate the score for a single prediction using the LLM judge."""
        prompts: List[str] = []
        for x, y, pred in zip(xs, ys, preds):
            judge_prompt = self._construct_judge_prompt(x, pred, y)
            prompts.append(judge_prompt)
        judge_responses = self.judge_llm.get_response(prompts)
        scores_str = extract_from_tag(judge_responses, "<final_score>", "</final_score>")
        scores = []
        for score_str in scores_str:
            try:
                # only numeric chars, - or . are allowed
                score_str = "".join(filter(lambda c: c.isdigit() or c in "-.", score_str))
                score = float(score_str)
                # normalize from [min_score, max_score] to [0, 1]
                score = (score - self.min_score) / (self.max_score - self.min_score)
                score = max(0.0, min(1.0, score))
            except ValueError:
                logger.warning(f"Failed to parse score '{score_str}' as float. Defaulting to a score 0.0.")
                score = 0.0

            scores.append(score)

        return np.asarray(scores, dtype=float)

__init__(df, judge_llm, x_column='x', y_column=None, task_description=None, n_subsamples=30, eval_strategy='full', seed=42, judge_prompt=None, min_score=-5.0, max_score=5.0, config=None)

Initialize the JudgeTask.

Parameters:

Name Type Description Default
df DataFrame

The input DataFrame containing the data.

 required
judge_llm BaseLLM

The LLM judging the predictions.

 required
x_column str

Name of the column containing input texts.

 'x'
y_column Optional[str]

Name of the column containing labels/ground truth (if applicable).

 None
task_description Optional[str]

Description of the task, parsed to the Judge-LLM and Meta-LLM.

 None
n_subsamples int

Number of subsamples to use for evaluation.

 30
eval_strategy EvalStrategy

Subsampling strategy to use for evaluation.

 'full'
seed int

Random seed for reproducibility.

 42
judge_prompt Optional[str]

Custom prompt for the judge. Note: The score of the Judge will be extracted inside tags.

 None
min_score float

Minimum score for evaluation.

 -5.0
max_score float

Maximum score for evaluation.

 5.0
config ExperimentConfig

Configuration for the task, overriding defaults.

 None
Source code in promptolution/tasks/judge_tasks.py 
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def __init__(
    self,
    df: pd.DataFrame,
    judge_llm: "BaseLLM",
    x_column: str = "x",
    y_column: Optional[str] = None,
    task_description: Optional[str] = None,
    n_subsamples: int = 30,
    eval_strategy: "EvalStrategy" = "full",
    seed: int = 42,
    judge_prompt: Optional[str] = None,
    min_score: float = -5.0,
    max_score: float = 5.0,
    config: "ExperimentConfig" = None,
):
    """Initialize the JudgeTask.

    Args:
        df (pd.DataFrame): The input DataFrame containing the data.
        judge_llm (BaseLLM): The LLM judging the predictions.
        x_column (str): Name of the column containing input texts.
        y_column (Optional[str]): Name of the column containing labels/ground truth (if applicable).
        task_description (Optional[str]): Description of the task, parsed to the Judge-LLM and Meta-LLM.
        n_subsamples (int): Number of subsamples to use for evaluation.
        eval_strategy (EvalStrategy): Subsampling strategy to use for evaluation.
        seed (int): Random seed for reproducibility.
        judge_prompt (Optional[str]): Custom prompt for the judge. Note: The score of the Judge will be extracted inside <final_score> tags.
        min_score (float): Minimum score for evaluation.
        max_score (float): Maximum score for evaluation.
        config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
    """
    if judge_prompt is None:
        judge_prompt = JUDGE_PROMPT_WITH_GROUND_TRUTH if y_column else JUDGE_PROMPT_WITHOUT_GROUND_TRUTH
    self.judge_prompt = judge_prompt
    self.min_score = min_score
    self.max_score = max_score

    super().__init__(
        df=df,
        x_column=x_column,
        y_column=y_column,
        task_description=task_description,
        n_subsamples=n_subsamples,
        eval_strategy=eval_strategy,
        seed=seed,
        config=config,
    )
    self.judge_llm = judge_llm
    self.task_type = "judge"

MultiObjectiveTask

Bases: BaseTask

A task that aggregates evaluations across multiple underlying tasks.

Source code in promptolution/tasks/multi_objective_task.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 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
 89
 90
 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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
class MultiObjectiveTask(BaseTask):
    """A task that aggregates evaluations across multiple underlying tasks."""

    def __init__(
        self,
        tasks: List[BaseTask],
        eval_strategy: Optional[EvalStrategy] = None,
    ) -> None:
        """Initialize with a list of tasks sharing subsampling and seed settings."""
        if not tasks:
            raise ValueError("tasks must be a non-empty list")

        primary = tasks[0]
        for t in tasks[1:]:
            assert t.n_subsamples == primary.n_subsamples, "All tasks must share n_subsamples"
            assert t.seed == primary.seed, "All tasks must share seed"
            assert t.eval_strategy == primary.eval_strategy, "All tasks must share eval_strategy"

        combined_description = "This task is a combination of the following tasks:\n" + "\n".join(
            [f"Task: {t.task_description}" for t in tasks if t.task_description]
        )

        super().__init__(
            df=primary.df,
            x_column=primary.x_column,
            y_column=primary.y_column,
            task_description=combined_description,
            n_subsamples=primary.n_subsamples,
            eval_strategy=eval_strategy or primary.eval_strategy,
            seed=primary.seed,
            config=None,
        )
        self.task_type = "multi"
        self.tasks = tasks
        self._scalarized_objective: bool = False

    def activate_scalarized_objective(self) -> None:
        """Force single-objective behavior by equally averaging task scores."""
        self._scalarized_objective = True

    def evaluate(  # type: ignore
        self,
        prompts: Prompt | List[Prompt],
        predictor,
        system_prompts: Optional[str | List[str]] = None,
        eval_strategy: Optional[EvalStrategy] = None,
    ) -> MultiObjectiveEvalResult | EvalResult:
        """Run prediction once, then score via each task's _evaluate."""
        prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts)
        strategy = eval_strategy or self.eval_strategy

        # Keep block alignment across tasks so block-based strategies stay in sync.
        for task in self.tasks:
            task.block_idx = self.block_idx

        xs, ys = self.subsample(eval_strategy=strategy)

        # Collect all uncached prompt/x/y triples across tasks to predict only once.
        prompts_to_evaluate: List[str] = []
        xs_to_evaluate: List[str] = []
        ys_to_evaluate: List[str] = []
        key_to_index: Dict[Tuple[str, str, str], int] = {}
        cache_keys: List[Tuple[str, str, str]] = []

        for task in self.tasks:
            t_prompts, t_xs, t_ys, t_keys = task._prepare_batch(prompts_list, xs, ys, eval_strategy=strategy)
            for prompt_str, x_val, y_val, key in zip(t_prompts, t_xs, t_ys, t_keys):
                if key in key_to_index:
                    continue
                key_to_index[key] = len(prompts_to_evaluate)
                prompts_to_evaluate.append(prompt_str)
                xs_to_evaluate.append(x_val)
                ys_to_evaluate.append(y_val)
                cache_keys.append(key)

        preds: List[str] = []
        pred_seqs: List[str] = []
        if prompts_to_evaluate:
            preds, pred_seqs = predictor.predict(
                prompts=prompts_to_evaluate,
                xs=xs_to_evaluate,
                system_prompts=system_prompts,
            )

        # Map predictions back to each task and populate caches via _evaluate.
        key_to_pred: Dict[Tuple[str, str, str], Tuple[str, str]] = {
            key: (preds[idx], pred_seqs[idx]) for key, idx in key_to_index.items()
        }

        per_task_results: List[EvalResult] = []
        for task in self.tasks:
            if cache_keys:
                xs_eval = [k[1] for k in cache_keys]
                ys_eval = [k[2] for k in cache_keys]
                preds_eval = [key_to_pred[k][0] for k in cache_keys]
                scores = task._evaluate(xs_eval, ys_eval, preds_eval)
                for score, cache_key in zip(scores, cache_keys):
                    task.eval_cache[cache_key] = score
                    task.seq_cache[cache_key] = key_to_pred[cache_key][1]

            scores_array, agg_scores, seqs = task._collect_results_from_cache(prompts_list, xs, ys)
            input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = task._compute_costs(
                prompts_list, xs, ys, predictor
            )

            per_task_results.append(
                EvalResult(
                    scores=scores_array,
                    agg_scores=agg_scores,
                    sequences=seqs,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    agg_input_tokens=agg_input_tokens,
                    agg_output_tokens=agg_output_tokens,
                )
            )

        stacked_scores = [r.scores for r in per_task_results]
        stacked_agg_scores = [r.agg_scores for r in per_task_results]

        # Record evaluated blocks for this evaluation (mirroring BaseTask behavior)
        for prompt in prompts_list:
            # Use self.block_idx (the MultiObjectiveTask's block_idx) if in a block strategy
            if strategy in ["sequential_block", "random_block"]:
                if isinstance(self.block_idx, list):
                    self.prompt_evaluated_blocks.setdefault(prompt, []).extend(self.block_idx)
                else:
                    self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx)
            elif strategy == "full":
                self.prompt_evaluated_blocks.setdefault(prompt, []).extend(list(range(self.n_blocks)))

        # Use first task's result for sequences and token counts (they're all the same across tasks)
        first_result = per_task_results[0]

        if self._scalarized_objective:
            return EvalResult(
                scores=np.mean(stacked_scores, axis=0),
                agg_scores=np.mean(stacked_agg_scores, axis=0),
                sequences=first_result.sequences,
                input_tokens=first_result.input_tokens,
                output_tokens=first_result.output_tokens,
                agg_input_tokens=first_result.agg_input_tokens,
                agg_output_tokens=first_result.agg_output_tokens,
            )

        return MultiObjectiveEvalResult(
            scores=stacked_scores,
            agg_scores=stacked_agg_scores,
            sequences=first_result.sequences,
            input_tokens=first_result.input_tokens,
            output_tokens=first_result.output_tokens,
            agg_input_tokens=first_result.agg_input_tokens,
            agg_output_tokens=first_result.agg_output_tokens,
        )

    def _evaluate(self, xs, ys, preds):  # pragma: no cover
        raise NotImplementedError("MultiObjectiveTask overrides evaluate directly")

__init__(tasks, eval_strategy=None)

Initialize with a list of tasks sharing subsampling and seed settings.

Source code in promptolution/tasks/multi_objective_task.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def __init__(
    self,
    tasks: List[BaseTask],
    eval_strategy: Optional[EvalStrategy] = None,
) -> None:
    """Initialize with a list of tasks sharing subsampling and seed settings."""
    if not tasks:
        raise ValueError("tasks must be a non-empty list")

    primary = tasks[0]
    for t in tasks[1:]:
        assert t.n_subsamples == primary.n_subsamples, "All tasks must share n_subsamples"
        assert t.seed == primary.seed, "All tasks must share seed"
        assert t.eval_strategy == primary.eval_strategy, "All tasks must share eval_strategy"

    combined_description = "This task is a combination of the following tasks:\n" + "\n".join(
        [f"Task: {t.task_description}" for t in tasks if t.task_description]
    )

    super().__init__(
        df=primary.df,
        x_column=primary.x_column,
        y_column=primary.y_column,
        task_description=combined_description,
        n_subsamples=primary.n_subsamples,
        eval_strategy=eval_strategy or primary.eval_strategy,
        seed=primary.seed,
        config=None,
    )
    self.task_type = "multi"
    self.tasks = tasks
    self._scalarized_objective: bool = False

activate_scalarized_objective()

Force single-objective behavior by equally averaging task scores.

Source code in promptolution/tasks/multi_objective_task.py 
64
65
66
def activate_scalarized_objective(self) -> None:
    """Force single-objective behavior by equally averaging task scores."""
    self._scalarized_objective = True

evaluate(prompts, predictor, system_prompts=None, eval_strategy=None)

Run prediction once, then score via each task's _evaluate.

Source code in promptolution/tasks/multi_objective_task.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 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
171
172
173
174
175
176
177
178
179
180
181
def evaluate(  # type: ignore
    self,
    prompts: Prompt | List[Prompt],
    predictor,
    system_prompts: Optional[str | List[str]] = None,
    eval_strategy: Optional[EvalStrategy] = None,
) -> MultiObjectiveEvalResult | EvalResult:
    """Run prediction once, then score via each task's _evaluate."""
    prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts)
    strategy = eval_strategy or self.eval_strategy

    # Keep block alignment across tasks so block-based strategies stay in sync.
    for task in self.tasks:
        task.block_idx = self.block_idx

    xs, ys = self.subsample(eval_strategy=strategy)

    # Collect all uncached prompt/x/y triples across tasks to predict only once.
    prompts_to_evaluate: List[str] = []
    xs_to_evaluate: List[str] = []
    ys_to_evaluate: List[str] = []
    key_to_index: Dict[Tuple[str, str, str], int] = {}
    cache_keys: List[Tuple[str, str, str]] = []

    for task in self.tasks:
        t_prompts, t_xs, t_ys, t_keys = task._prepare_batch(prompts_list, xs, ys, eval_strategy=strategy)
        for prompt_str, x_val, y_val, key in zip(t_prompts, t_xs, t_ys, t_keys):
            if key in key_to_index:
                continue
            key_to_index[key] = len(prompts_to_evaluate)
            prompts_to_evaluate.append(prompt_str)
            xs_to_evaluate.append(x_val)
            ys_to_evaluate.append(y_val)
            cache_keys.append(key)

    preds: List[str] = []
    pred_seqs: List[str] = []
    if prompts_to_evaluate:
        preds, pred_seqs = predictor.predict(
            prompts=prompts_to_evaluate,
            xs=xs_to_evaluate,
            system_prompts=system_prompts,
        )

    # Map predictions back to each task and populate caches via _evaluate.
    key_to_pred: Dict[Tuple[str, str, str], Tuple[str, str]] = {
        key: (preds[idx], pred_seqs[idx]) for key, idx in key_to_index.items()
    }

    per_task_results: List[EvalResult] = []
    for task in self.tasks:
        if cache_keys:
            xs_eval = [k[1] for k in cache_keys]
            ys_eval = [k[2] for k in cache_keys]
            preds_eval = [key_to_pred[k][0] for k in cache_keys]
            scores = task._evaluate(xs_eval, ys_eval, preds_eval)
            for score, cache_key in zip(scores, cache_keys):
                task.eval_cache[cache_key] = score
                task.seq_cache[cache_key] = key_to_pred[cache_key][1]

        scores_array, agg_scores, seqs = task._collect_results_from_cache(prompts_list, xs, ys)
        input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = task._compute_costs(
            prompts_list, xs, ys, predictor
        )

        per_task_results.append(
            EvalResult(
                scores=scores_array,
                agg_scores=agg_scores,
                sequences=seqs,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                agg_input_tokens=agg_input_tokens,
                agg_output_tokens=agg_output_tokens,
            )
        )

    stacked_scores = [r.scores for r in per_task_results]
    stacked_agg_scores = [r.agg_scores for r in per_task_results]

    # Record evaluated blocks for this evaluation (mirroring BaseTask behavior)
    for prompt in prompts_list:
        # Use self.block_idx (the MultiObjectiveTask's block_idx) if in a block strategy
        if strategy in ["sequential_block", "random_block"]:
            if isinstance(self.block_idx, list):
                self.prompt_evaluated_blocks.setdefault(prompt, []).extend(self.block_idx)
            else:
                self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx)
        elif strategy == "full":
            self.prompt_evaluated_blocks.setdefault(prompt, []).extend(list(range(self.n_blocks)))

    # Use first task's result for sequences and token counts (they're all the same across tasks)
    first_result = per_task_results[0]

    if self._scalarized_objective:
        return EvalResult(
            scores=np.mean(stacked_scores, axis=0),
            agg_scores=np.mean(stacked_agg_scores, axis=0),
            sequences=first_result.sequences,
            input_tokens=first_result.input_tokens,
            output_tokens=first_result.output_tokens,
            agg_input_tokens=first_result.agg_input_tokens,
            agg_output_tokens=first_result.agg_output_tokens,
        )

    return MultiObjectiveEvalResult(
        scores=stacked_scores,
        agg_scores=stacked_agg_scores,
        sequences=first_result.sequences,
        input_tokens=first_result.input_tokens,
        output_tokens=first_result.output_tokens,
        agg_input_tokens=first_result.agg_input_tokens,
        agg_output_tokens=first_result.agg_output_tokens,
    )

RewardTask

Bases: BaseTask

A task that evaluates a predictor using a reward function.

This task takes a DataFrame, a column name for input data, and a reward function. The reward function takes in a prediction as input and returns a scalar reward.

Source code in promptolution/tasks/reward_tasks.py 
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
47
48
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
class RewardTask(BaseTask):
    """A task that evaluates a predictor using a reward function.

    This task takes a DataFrame, a column name for input data, and a reward function.
    The reward function takes in a prediction as input and returns a scalar reward.
    """

    def __init__(
        self,
        df: pd.DataFrame,
        reward_function: Callable[[str], float],
        x_column: str = "x",
        y_column: Optional[str] = None,
        reward_columns: Optional[List[str]] = None,
        task_description: Optional[str] = None,
        n_subsamples: int = 30,
        eval_strategy: "EvalStrategy" = "full",
        seed: int = 42,
        config: Optional["ExperimentConfig"] = None,
    ) -> None:
        """Initialize the RewardTask.

        Args:
            df (pd.DataFrame): Input DataFrame containing the data.
            reward_function (Callable): Function that takes a prediction, potential keyword arguments from the dataframe, and returns a reward score. Note: The optimizers aim to maximize.
            x_column (str, optional): Name of the column containing input texts. Defaults to "x".
            y_column (str, optional): Name of the column containing target texts if available. Defaults to None.
            reward_columns (List[str], optional): Additional dataframe columns to pass as keyword args to reward_function.
            task_description (str, optional): Description of the task.
            n_subsamples (int, optional): Number of subsamples to use. Defaults to 30.
            eval_strategy (str, optional): Subsampling strategy to use. Defaults to "full".
            seed (int, optional): Random seed for reproducibility. Defaults to 42.
            config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
        """
        self.reward_function = reward_function
        self.reward_columns = reward_columns or []
        super().__init__(
            df=df,
            x_column=x_column,
            y_column=y_column,
            task_description=task_description,
            n_subsamples=n_subsamples,
            eval_strategy=eval_strategy,
            seed=seed,
            config=config,
        )
        self.task_type = "reward"
        # x -> kwargs to reward function
        km = self.df.set_index(x_column)[self.reward_columns].to_dict("index")
        self.kwargs_map = defaultdict(dict, km)

    def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarray:
        """Calculate reward for each prediction, passing configured columns as kwargs."""
        kwargs_list = [self.kwargs_map[x] for x in xs]
        rewards = [self.reward_function(pred, **kwargs) for pred, kwargs in zip(preds, kwargs_list)]
        return np.asarray(rewards, dtype=float)

__init__(df, reward_function, x_column='x', y_column=None, reward_columns=None, task_description=None, n_subsamples=30, eval_strategy='full', seed=42, config=None)

Initialize the RewardTask.

Parameters:

Name Type Description Default
df DataFrame

Input DataFrame containing the data.

 required
reward_function Callable

Function that takes a prediction, potential keyword arguments from the dataframe, and returns a reward score. Note: The optimizers aim to maximize.

 required
x_column str

Name of the column containing input texts. Defaults to "x".

 'x'
y_column str

Name of the column containing target texts if available. Defaults to None.

 None
reward_columns List[str]

Additional dataframe columns to pass as keyword args to reward_function.

 None
task_description str

Description of the task.

 None
n_subsamples int

Number of subsamples to use. Defaults to 30.

 30
eval_strategy str

Subsampling strategy to use. Defaults to "full".

 'full'
seed int

Random seed for reproducibility. Defaults to 42.

 42
config ExperimentConfig

Configuration for the task, overriding defaults.

 None
Source code in promptolution/tasks/reward_tasks.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(
    self,
    df: pd.DataFrame,
    reward_function: Callable[[str], float],
    x_column: str = "x",
    y_column: Optional[str] = None,
    reward_columns: Optional[List[str]] = None,
    task_description: Optional[str] = None,
    n_subsamples: int = 30,
    eval_strategy: "EvalStrategy" = "full",
    seed: int = 42,
    config: Optional["ExperimentConfig"] = None,
) -> None:
    """Initialize the RewardTask.

    Args:
        df (pd.DataFrame): Input DataFrame containing the data.
        reward_function (Callable): Function that takes a prediction, potential keyword arguments from the dataframe, and returns a reward score. Note: The optimizers aim to maximize.
        x_column (str, optional): Name of the column containing input texts. Defaults to "x".
        y_column (str, optional): Name of the column containing target texts if available. Defaults to None.
        reward_columns (List[str], optional): Additional dataframe columns to pass as keyword args to reward_function.
        task_description (str, optional): Description of the task.
        n_subsamples (int, optional): Number of subsamples to use. Defaults to 30.
        eval_strategy (str, optional): Subsampling strategy to use. Defaults to "full".
        seed (int, optional): Random seed for reproducibility. Defaults to 42.
        config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
    """
    self.reward_function = reward_function
    self.reward_columns = reward_columns or []
    super().__init__(
        df=df,
        x_column=x_column,
        y_column=y_column,
        task_description=task_description,
        n_subsamples=n_subsamples,
        eval_strategy=eval_strategy,
        seed=seed,
        config=config,
    )
    self.task_type = "reward"
    # x -> kwargs to reward function
    km = self.df.set_index(x_column)[self.reward_columns].to_dict("index")
    self.kwargs_map = defaultdict(dict, km)

base_task

Base module for tasks.

BaseTask

Bases: ABC

Abstract base class for tasks in the promptolution library.

Source code in promptolution/tasks/base_task.py 
 40
 41
 42
 43
 44
 45
 46
 47
 48
 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
 89
 90
 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
171
172
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
207
208
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
245
246
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
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
class BaseTask(ABC):
    """Abstract base class for tasks in the promptolution library."""

    def __init__(
        self,
        df: pd.DataFrame,
        x_column: str,
        y_column: Optional[str] = None,
        task_description: Optional[str] = None,
        n_subsamples: int = 30,
        eval_strategy: "EvalStrategy" = "full",
        seed: int = 42,
        config: Optional["ExperimentConfig"] = None,
    ) -> None:
        """Initialize the BaseTask.

        Args:
            df (pd.DataFrame): The input DataFrame containing the data.
            x_column (str): Name of the column containing input texts.
            y_column (Optional[str]): Name of the column containing labels/ground truth (if applicable).
            task_description (str): Description of the task.
            n_subsamples (int): Number of subsamples to use for evaluation.
            eval_strategy (Literal): Subsampling strategy ("full", "subsample", "sequential_block", "random_block", "evaluated").
            seed (int): Random seed for reproducibility.
            config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
        """
        self.df = df.drop_duplicates(subset=[x_column])
        if len(self.df) != len(df):
            logger.warning(
                f"Duplicate entries detected for x_column '{x_column}' - dropped {len(df) - len(self.df)} rows to enforce uniqueness."
            )
        self.x_column: str = x_column
        self.y_column: Optional[str] = y_column
        self.task_type: TaskType | None = None
        self.task_description: Optional[str] = task_description
        self.n_subsamples: int = n_subsamples
        self.eval_strategy: EvalStrategy = eval_strategy
        self.seed: int = seed

        super().__init__()
        if config is not None:
            config.apply_to(self)

        self.xs: List[str] = self.df[self.x_column].values.astype(str).tolist()
        self.has_y: bool = y_column is not None
        if self.has_y and y_column is not None:
            self.ys: List[str] = self.df[y_column].values.astype(str).tolist()
        else:
            # If no y_column is provided, create a dummy y array
            self.ys = [""] * len(self.xs)

        self.block_idx: int = 0
        self.n_blocks: int = len(self.xs) // self.n_subsamples if self.n_subsamples > 0 else 1
        self.rng = np.random.default_rng(seed)

        self.eval_cache: Dict[Tuple[str, str, str], float] = {}  # (prompt, x, y): scores per datapoint
        self.seq_cache: Dict[Tuple[str, str, str], str] = {}  # (prompt, x, y): raw model output per datapoint

        self.prompt_evaluated_blocks: Dict[Prompt, List[int]] = {}  # prompt_str: set of evaluated block indices

    def subsample(
        self, eval_strategy: Optional["EvalStrategy"] = None, block_idx: List[int] | None = None
    ) -> Tuple[List[str], List[str]]:
        """Subsample the dataset based on the specified parameters.

        Args:
            eval_strategy (EvalStrategy, optional): Subsampling strategy to use instead of self.eval_strategy. Defaults to None.
            block_idx (List[int] | None, optional): Specific block index or indices to evaluate, overriding eval_strategy. Defaults to None.

        Returns:
            Tuple[List[str], List[str]]: Subsampled input data and labels.
        """
        if block_idx is not None:
            indices: List[int] = []
            for idx in block_idx:
                start_idx = idx * self.n_subsamples
                end_idx = min((idx + 1) * self.n_subsamples, len(self.xs))
                indices.extend(range(start_idx, end_idx))

            return [self.xs[i] for i in indices], [self.ys[i] for i in indices]

        if eval_strategy is None:
            eval_strategy = self.eval_strategy

        if eval_strategy in ["full", "evaluated"]:
            return self.xs, self.ys
        elif eval_strategy == "subsample":
            indices = self.rng.choice(len(self.xs), min(self.n_subsamples, len(self.xs)), replace=False)
            return [self.xs[i] for i in indices], [self.ys[i] for i in indices]
        elif eval_strategy == "random_block":
            block_id = self.rng.integers(0, self.n_blocks)
            start_idx = block_id * self.n_subsamples
            end_idx = min((block_id + 1) * self.n_subsamples, len(self.xs))
            indices = np.arange(start_idx, end_idx)
            return [self.xs[i] for i in indices], [self.ys[i] for i in indices]
        elif eval_strategy == "sequential_block":
            # Handle case where self.block_idx is a list
            if isinstance(self.block_idx, list):
                indices_list: List[int] = []
                for idx in self.block_idx:
                    start_idx = idx * self.n_subsamples
                    end_idx = min((idx + 1) * self.n_subsamples, len(self.xs))
                    indices_list.extend(range(start_idx, end_idx))
                return [self.xs[i] for i in indices_list], [self.ys[i] for i in indices_list]
            else:
                start_idx = self.block_idx * self.n_subsamples
                end_idx = min((self.block_idx + 1) * self.n_subsamples, len(self.xs))
                indices = np.arange(start_idx, end_idx)
                return [self.xs[i] for i in indices], [self.ys[i] for i in indices]
        else:
            raise ValueError(f"Unknown subsampling strategy: '{eval_strategy}'")

    def _prepare_batch(
        self,
        prompts: List[Prompt],
        xs: List[str],
        ys: List[str],
        eval_strategy: Literal["full", "subsample", "sequential_block", "random_block", "evaluated"] = "full",
    ) -> Tuple[List[str], List[str], List[str], List[Tuple[str, str, str]]]:
        """Return uncached prompt/x/y triples for prediction and their cache keys."""
        if eval_strategy == "evaluated":
            return [], [], [], []

        prompts_to_predict: List[str] = []
        xs_to_predict: List[str] = []
        ys_to_predict: List[str] = []
        keys_to_predict: List[Tuple[str, str, str]] = []

        for prompt in prompts:
            for x, y in zip(xs, ys):
                cache_key = (str(prompt), x, str(y))
                if cache_key in self.eval_cache:
                    continue
                prompts_to_predict.append(str(prompt))
                xs_to_predict.append(x)
                ys_to_predict.append(str(y))
                keys_to_predict.append(cache_key)

        return prompts_to_predict, xs_to_predict, ys_to_predict, keys_to_predict

    @staticmethod
    def _cache_key(prompt: Prompt, x: str, y: str) -> Tuple[str, str, str]:
        return (prompt.construct_prompt(), x, y)

    def _collect_results_from_cache(
        self, prompts: List[Prompt], xs: List[str], ys: List[str]
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """Collect cached scores and sequences for provided prompts/xs/ys."""
        score_rows: List[List[float]] = []
        seq_rows: List[List[str]] = []

        for prompt in prompts:
            datapoint_scores: List[float] = []
            datapoint_seqs: List[str] = []
            for x, y in zip(xs, ys):
                cache_key = self._cache_key(prompt, x, str(y))
                if cache_key not in self.eval_cache:
                    datapoint_scores.append(np.nan)  # Fill with NaN instead of skipping
                    datapoint_seqs.append("")
                else:
                    datapoint_score = self.eval_cache[cache_key]
                    datapoint_scores.append(datapoint_score)
                    datapoint_seqs.append(self.seq_cache.get(cache_key, ""))
            score_rows.append(datapoint_scores)
            seq_rows.append(datapoint_seqs)

        scores_array = np.array(score_rows, dtype=float)
        agg_scores = np.nanmean(scores_array, axis=1) if scores_array.size else np.array([])
        seqs_array = np.array(seq_rows, dtype=object)
        return scores_array, agg_scores, seqs_array

    def _compute_costs(
        self,
        prompts: List[Prompt],
        xs: List[str],
        ys: List[str],
        predictor: "BasePredictor",
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        token_counter = get_token_counter(predictor.llm)

        per_prompt_inputs: List[np.ndarray] = []
        per_prompt_outputs: List[np.ndarray] = []

        for prompt in prompts:
            prompt_token_count = token_counter(prompt.construct_prompt())
            seq_token_counts: List[float] = []
            input_token_counts = []
            for x, y in zip(xs, ys):
                cache_key = self._cache_key(prompt, x, str(y))
                if cache_key not in self.seq_cache:
                    # Use NaN for missing datapoints instead of skipping
                    seq_token_counts.append(np.nan)
                    input_token_counts.append(np.nan)
                    continue
                seq_text = self.seq_cache[cache_key]
                seq_token_counts.append(token_counter(seq_text))
                input_token_counts.append(token_counter(x))

            prompt_input_tokens = np.array(input_token_counts, dtype=float) + prompt_token_count
            output_token_counts = np.array(seq_token_counts, dtype=float) - np.array(input_token_counts, dtype=float)

            per_prompt_inputs.append(np.asarray(prompt_input_tokens, dtype=float))
            per_prompt_outputs.append(output_token_counts)

        inputs_array = np.vstack(per_prompt_inputs)
        outputs_array = np.vstack(per_prompt_outputs)

        agg_input_tokens = np.nanmean(inputs_array, axis=1)
        agg_output_tokens = np.nanmean(outputs_array, axis=1)

        return inputs_array, outputs_array, agg_input_tokens, agg_output_tokens

    @abstractmethod
    def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarray:
        """Abstract method to calculate the score for a predictions.

        This method should be implemented by subclasses based on their specific evaluation logic.
        """
        raise NotImplementedError

    def activate_scalarized_objective(self) -> None:
        """Activate scalarized objective for multi-objective tasks."""
        raise NotImplementedError

    def evaluate(
        self,
        prompts: Union[Prompt, List[Prompt]],
        predictor: "BasePredictor",
        system_prompts: Optional[Union[str, List[str]]] = None,
        eval_strategy: Optional["EvalStrategy"] = None,
        block_idx: int | list[int] | None = None,
    ) -> EvalResult:
        """Evaluate a set of prompts using a given predictor.

        This method orchestrates subsampling, prediction, caching, and result collection.
        Sequences, token costs, raw scores, and aggregated scores are always returned.

        Args:
            prompts (Union[Prompt, List[Prompt]]): A single prompt or a list of prompts to evaluate. Results will be returned in the same order.
            predictor (BasePredictor): The predictor to evaluate the prompts with.
            system_prompts (Optional[Union[str, List[str]]], optional): Optional system prompts to parse to the predictor.
            eval_strategy (Optional[EvalStrategy], optional): Subsampling strategy to use instead of self.eval_strategy. Defaults to None, which uses self.eval_strategy.
            block_idx (Optional[int | list[int]], optional): Specific block index or indices to evaluate, overriding eval_strategy. Defaults to None.
        """
        prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts)
        eval_strategy = eval_strategy or self.eval_strategy

        if block_idx is not None and isinstance(block_idx, int):
            block_idx = [block_idx]

        xs, ys = self.subsample(eval_strategy=eval_strategy, block_idx=block_idx)
        (
            prompts_to_evaluate,
            xs_to_evaluate,
            ys_to_evaluate,
            cache_keys,
        ) = self._prepare_batch(prompts_list, xs, ys, eval_strategy=eval_strategy)

        preds, pred_seqs = predictor.predict(
            prompts=prompts_to_evaluate,
            xs=xs_to_evaluate,
            system_prompts=system_prompts,
        )

        scores = self._evaluate(xs_to_evaluate, ys_to_evaluate, preds)
        for i, cache_key in enumerate(cache_keys):
            self.eval_cache[cache_key] = scores[i]
            self.seq_cache[cache_key] = str(pred_seqs[i])

        scores, agg_scores, seqs = self._collect_results_from_cache(
            prompts_list,
            xs,
            ys,
        )

        # Record evaluated block for block strategies
        for prompt in prompts_list:
            if block_idx is not None:
                self.prompt_evaluated_blocks.setdefault(prompt, []).extend(block_idx)
            elif eval_strategy in ["sequential_block", "random_block"]:
                # Handle case where self.block_idx is a list
                if isinstance(self.block_idx, list):
                    self.prompt_evaluated_blocks.setdefault(prompt, []).extend(self.block_idx)
                else:
                    self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx)
            elif eval_strategy == "full":
                self.prompt_evaluated_blocks.setdefault(prompt, []).extend(list(range(self.n_blocks)))

        input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = self._compute_costs(
            prompts_list, xs, ys, predictor
        )

        return EvalResult(
            scores=scores,
            agg_scores=agg_scores,
            sequences=seqs,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            agg_input_tokens=agg_input_tokens,
            agg_output_tokens=agg_output_tokens,
        )

    def pop_datapoints(self, n: Optional[int] = None, frac: Optional[float] = None) -> pd.DataFrame:
        """Pop a number of datapoints from the dataset.

        Args:
            n (int, optional): Number of datapoints to pop. Defaults to None.
            frac (float, optional): Fraction of datapoints to pop. Defaults to None.

        Returns:
            pd.DataFrame: DataFrame containing the popped datapoints.
        """
        assert n is None or frac is None, "Only one of n or frac can be specified."
        if n is not None:
            indices = self.rng.choice(len(self.xs), n, replace=False)
        elif frac is not None:
            indices = self.rng.choice(len(self.xs), int(len(self.xs) * frac), replace=False)
        else:
            raise ValueError("Either n or frac must be specified.")

        popped_xs = [self.xs[i] for i in indices]
        popped_ys = [self.ys[i] for i in indices]
        df_popped = pd.DataFrame({self.x_column: popped_xs, self.y_column: popped_ys})

        self.xs = [x for i, x in enumerate(self.xs) if i not in indices]
        self.ys = [y for i, y in enumerate(self.ys) if i not in indices]

        # Update n_blocks and block_idx based on the new dataset size
        self.n_blocks = len(self.xs) // self.n_subsamples if self.n_subsamples > 0 else 1
        self.block_idx = min(self.block_idx, self.n_blocks - 1) if self.n_blocks > 0 else 0

        # Clear cache for popped items (optional, but good practice if memory is a concern)
        keys_to_remove = []
        for key in self.eval_cache:
            if key[1] in popped_xs and key[2] in popped_ys:  # Check if the x and y correspond to popped data
                keys_to_remove.append(key)
        for key in keys_to_remove:
            self.eval_cache.pop(key, None)
            self.seq_cache.pop(key, None)

        return df_popped

    def increment_block_idx(self) -> None:
        """Increment the block index for subsampling.

        Raises:
            ValueError: If the eval_strategy does not contain "block".
        """
        if "block" not in self.eval_strategy:
            raise ValueError("Block increment is only valid for block subsampling.")
        assert isinstance(self.block_idx, int), "Block index must be an integer to increment."
        self.block_idx += 1
        if self.n_blocks > 0:  # Ensure n_blocks is not zero to avoid division by zero
            self.block_idx %= self.n_blocks
        else:
            self.block_idx = 0  # If no blocks, reset to 0

    def reset_block_idx(self) -> None:
        """Reset the block index for subsampling.

        Raises:
            ValueError: If the eval_strategy does not contain "block".
        """
        if "block" not in self.eval_strategy:
            raise ValueError("Block reset is only valid for block subsampling.")
        self.block_idx = 0

    def set_block_idx(self, idx: int) -> None:
        """Set the block index (or indices) for block subsampling strategies."""
        if "block" not in self.eval_strategy:
            raise ValueError("Block assignment is only valid for block subsampling.")

        assert isinstance(idx, int), "Block index must be an integer"

        self.block_idx = idx

    def get_evaluated_blocks(self, prompts: Union[Prompt, List[Prompt]]) -> Dict[Prompt, List[int]]:
        """Return mapping of prompt string to evaluated block indices."""
        prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts)
        return {p: list(self.prompt_evaluated_blocks.get(p, [])) for p in prompts_list}

__init__(df, x_column, y_column=None, task_description=None, n_subsamples=30, eval_strategy='full', seed=42, config=None)

Initialize the BaseTask.

Parameters:

Name Type Description Default
df DataFrame

The input DataFrame containing the data.

 required
x_column str

Name of the column containing input texts.

 required
y_column Optional[str]

Name of the column containing labels/ground truth (if applicable).

 None
task_description str

Description of the task.

 None
n_subsamples int

Number of subsamples to use for evaluation.

 30
eval_strategy Literal

Subsampling strategy ("full", "subsample", "sequential_block", "random_block", "evaluated").

 'full'
seed int

Random seed for reproducibility.

 42
config ExperimentConfig

Configuration for the task, overriding defaults.

 None
Source code in promptolution/tasks/base_task.py 
43
44
45
46
47
48
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
89
90
91
92
93
94
95
96
97
98
def __init__(
    self,
    df: pd.DataFrame,
    x_column: str,
    y_column: Optional[str] = None,
    task_description: Optional[str] = None,
    n_subsamples: int = 30,
    eval_strategy: "EvalStrategy" = "full",
    seed: int = 42,
    config: Optional["ExperimentConfig"] = None,
) -> None:
    """Initialize the BaseTask.

    Args:
        df (pd.DataFrame): The input DataFrame containing the data.
        x_column (str): Name of the column containing input texts.
        y_column (Optional[str]): Name of the column containing labels/ground truth (if applicable).
        task_description (str): Description of the task.
        n_subsamples (int): Number of subsamples to use for evaluation.
        eval_strategy (Literal): Subsampling strategy ("full", "subsample", "sequential_block", "random_block", "evaluated").
        seed (int): Random seed for reproducibility.
        config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
    """
    self.df = df.drop_duplicates(subset=[x_column])
    if len(self.df) != len(df):
        logger.warning(
            f"Duplicate entries detected for x_column '{x_column}' - dropped {len(df) - len(self.df)} rows to enforce uniqueness."
        )
    self.x_column: str = x_column
    self.y_column: Optional[str] = y_column
    self.task_type: TaskType | None = None
    self.task_description: Optional[str] = task_description
    self.n_subsamples: int = n_subsamples
    self.eval_strategy: EvalStrategy = eval_strategy
    self.seed: int = seed

    super().__init__()
    if config is not None:
        config.apply_to(self)

    self.xs: List[str] = self.df[self.x_column].values.astype(str).tolist()
    self.has_y: bool = y_column is not None
    if self.has_y and y_column is not None:
        self.ys: List[str] = self.df[y_column].values.astype(str).tolist()
    else:
        # If no y_column is provided, create a dummy y array
        self.ys = [""] * len(self.xs)

    self.block_idx: int = 0
    self.n_blocks: int = len(self.xs) // self.n_subsamples if self.n_subsamples > 0 else 1
    self.rng = np.random.default_rng(seed)

    self.eval_cache: Dict[Tuple[str, str, str], float] = {}  # (prompt, x, y): scores per datapoint
    self.seq_cache: Dict[Tuple[str, str, str], str] = {}  # (prompt, x, y): raw model output per datapoint

    self.prompt_evaluated_blocks: Dict[Prompt, List[int]] = {}  # prompt_str: set of evaluated block indices

activate_scalarized_objective()

Activate scalarized objective for multi-objective tasks.

Source code in promptolution/tasks/base_task.py 
260
261
262
def activate_scalarized_objective(self) -> None:
    """Activate scalarized objective for multi-objective tasks."""
    raise NotImplementedError

evaluate(prompts, predictor, system_prompts=None, eval_strategy=None, block_idx=None)

Evaluate a set of prompts using a given predictor.

This method orchestrates subsampling, prediction, caching, and result collection. Sequences, token costs, raw scores, and aggregated scores are always returned.

Parameters:

Name Type Description Default
prompts Union[Prompt, List[Prompt]]

A single prompt or a list of prompts to evaluate. Results will be returned in the same order.

 required
predictor BasePredictor

The predictor to evaluate the prompts with.

 required
system_prompts Optional[Union[str, List[str]]]

Optional system prompts to parse to the predictor.

 None
eval_strategy Optional[EvalStrategy]

Subsampling strategy to use instead of self.eval_strategy. Defaults to None, which uses self.eval_strategy.

 None
block_idx Optional[int | list[int]]

Specific block index or indices to evaluate, overriding eval_strategy. Defaults to None.

 None
Source code in promptolution/tasks/base_task.py 
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
def evaluate(
    self,
    prompts: Union[Prompt, List[Prompt]],
    predictor: "BasePredictor",
    system_prompts: Optional[Union[str, List[str]]] = None,
    eval_strategy: Optional["EvalStrategy"] = None,
    block_idx: int | list[int] | None = None,
) -> EvalResult:
    """Evaluate a set of prompts using a given predictor.

    This method orchestrates subsampling, prediction, caching, and result collection.
    Sequences, token costs, raw scores, and aggregated scores are always returned.

    Args:
        prompts (Union[Prompt, List[Prompt]]): A single prompt or a list of prompts to evaluate. Results will be returned in the same order.
        predictor (BasePredictor): The predictor to evaluate the prompts with.
        system_prompts (Optional[Union[str, List[str]]], optional): Optional system prompts to parse to the predictor.
        eval_strategy (Optional[EvalStrategy], optional): Subsampling strategy to use instead of self.eval_strategy. Defaults to None, which uses self.eval_strategy.
        block_idx (Optional[int | list[int]], optional): Specific block index or indices to evaluate, overriding eval_strategy. Defaults to None.
    """
    prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts)
    eval_strategy = eval_strategy or self.eval_strategy

    if block_idx is not None and isinstance(block_idx, int):
        block_idx = [block_idx]

    xs, ys = self.subsample(eval_strategy=eval_strategy, block_idx=block_idx)
    (
        prompts_to_evaluate,
        xs_to_evaluate,
        ys_to_evaluate,
        cache_keys,
    ) = self._prepare_batch(prompts_list, xs, ys, eval_strategy=eval_strategy)

    preds, pred_seqs = predictor.predict(
        prompts=prompts_to_evaluate,
        xs=xs_to_evaluate,
        system_prompts=system_prompts,
    )

    scores = self._evaluate(xs_to_evaluate, ys_to_evaluate, preds)
    for i, cache_key in enumerate(cache_keys):
        self.eval_cache[cache_key] = scores[i]
        self.seq_cache[cache_key] = str(pred_seqs[i])

    scores, agg_scores, seqs = self._collect_results_from_cache(
        prompts_list,
        xs,
        ys,
    )

    # Record evaluated block for block strategies
    for prompt in prompts_list:
        if block_idx is not None:
            self.prompt_evaluated_blocks.setdefault(prompt, []).extend(block_idx)
        elif eval_strategy in ["sequential_block", "random_block"]:
            # Handle case where self.block_idx is a list
            if isinstance(self.block_idx, list):
                self.prompt_evaluated_blocks.setdefault(prompt, []).extend(self.block_idx)
            else:
                self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx)
        elif eval_strategy == "full":
            self.prompt_evaluated_blocks.setdefault(prompt, []).extend(list(range(self.n_blocks)))

    input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = self._compute_costs(
        prompts_list, xs, ys, predictor
    )

    return EvalResult(
        scores=scores,
        agg_scores=agg_scores,
        sequences=seqs,
        input_tokens=input_tokens,
        output_tokens=output_tokens,
        agg_input_tokens=agg_input_tokens,
        agg_output_tokens=agg_output_tokens,
    )

get_evaluated_blocks(prompts)

Return mapping of prompt string to evaluated block indices.

Source code in promptolution/tasks/base_task.py 
416
417
418
419
def get_evaluated_blocks(self, prompts: Union[Prompt, List[Prompt]]) -> Dict[Prompt, List[int]]:
    """Return mapping of prompt string to evaluated block indices."""
    prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts)
    return {p: list(self.prompt_evaluated_blocks.get(p, [])) for p in prompts_list}

increment_block_idx()

Increment the block index for subsampling.

Raises:

Type Description
ValueError

If the eval_strategy does not contain "block".
Source code in promptolution/tasks/base_task.py 
382
383
384
385
386
387
388
389
390
391
392
393
394
395
def increment_block_idx(self) -> None:
    """Increment the block index for subsampling.

    Raises:
        ValueError: If the eval_strategy does not contain "block".
    """
    if "block" not in self.eval_strategy:
        raise ValueError("Block increment is only valid for block subsampling.")
    assert isinstance(self.block_idx, int), "Block index must be an integer to increment."
    self.block_idx += 1
    if self.n_blocks > 0:  # Ensure n_blocks is not zero to avoid division by zero
        self.block_idx %= self.n_blocks
    else:
        self.block_idx = 0  # If no blocks, reset to 0

pop_datapoints(n=None, frac=None)

Pop a number of datapoints from the dataset.

Parameters:

Name Type Description Default
n int

Number of datapoints to pop. Defaults to None.

 None
frac float

Fraction of datapoints to pop. Defaults to None.

 None

Returns:

Type Description
DataFrame

pd.DataFrame: DataFrame containing the popped datapoints.
Source code in promptolution/tasks/base_task.py 
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
def pop_datapoints(self, n: Optional[int] = None, frac: Optional[float] = None) -> pd.DataFrame:
    """Pop a number of datapoints from the dataset.

    Args:
        n (int, optional): Number of datapoints to pop. Defaults to None.
        frac (float, optional): Fraction of datapoints to pop. Defaults to None.

    Returns:
        pd.DataFrame: DataFrame containing the popped datapoints.
    """
    assert n is None or frac is None, "Only one of n or frac can be specified."
    if n is not None:
        indices = self.rng.choice(len(self.xs), n, replace=False)
    elif frac is not None:
        indices = self.rng.choice(len(self.xs), int(len(self.xs) * frac), replace=False)
    else:
        raise ValueError("Either n or frac must be specified.")

    popped_xs = [self.xs[i] for i in indices]
    popped_ys = [self.ys[i] for i in indices]
    df_popped = pd.DataFrame({self.x_column: popped_xs, self.y_column: popped_ys})

    self.xs = [x for i, x in enumerate(self.xs) if i not in indices]
    self.ys = [y for i, y in enumerate(self.ys) if i not in indices]

    # Update n_blocks and block_idx based on the new dataset size
    self.n_blocks = len(self.xs) // self.n_subsamples if self.n_subsamples > 0 else 1
    self.block_idx = min(self.block_idx, self.n_blocks - 1) if self.n_blocks > 0 else 0

    # Clear cache for popped items (optional, but good practice if memory is a concern)
    keys_to_remove = []
    for key in self.eval_cache:
        if key[1] in popped_xs and key[2] in popped_ys:  # Check if the x and y correspond to popped data
            keys_to_remove.append(key)
    for key in keys_to_remove:
        self.eval_cache.pop(key, None)
        self.seq_cache.pop(key, None)

    return df_popped

reset_block_idx()

Reset the block index for subsampling.

Raises:

Type Description
ValueError

If the eval_strategy does not contain "block".
Source code in promptolution/tasks/base_task.py 
397
398
399
400
401
402
403
404
405
def reset_block_idx(self) -> None:
    """Reset the block index for subsampling.

    Raises:
        ValueError: If the eval_strategy does not contain "block".
    """
    if "block" not in self.eval_strategy:
        raise ValueError("Block reset is only valid for block subsampling.")
    self.block_idx = 0

set_block_idx(idx)

Set the block index (or indices) for block subsampling strategies.

Source code in promptolution/tasks/base_task.py 
407
408
409
410
411
412
413
414
def set_block_idx(self, idx: int) -> None:
    """Set the block index (or indices) for block subsampling strategies."""
    if "block" not in self.eval_strategy:
        raise ValueError("Block assignment is only valid for block subsampling.")

    assert isinstance(idx, int), "Block index must be an integer"

    self.block_idx = idx

subsample(eval_strategy=None, block_idx=None)

Subsample the dataset based on the specified parameters.

Parameters:

Name Type Description Default
eval_strategy EvalStrategy

Subsampling strategy to use instead of self.eval_strategy. Defaults to None.

 None
block_idx List[int] | None

Specific block index or indices to evaluate, overriding eval_strategy. Defaults to None.

 None

Returns:

Type Description
Tuple[List[str], List[str]]

Tuple[List[str], List[str]]: Subsampled input data and labels.
Source code in promptolution/tasks/base_task.py 
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
def subsample(
    self, eval_strategy: Optional["EvalStrategy"] = None, block_idx: List[int] | None = None
) -> Tuple[List[str], List[str]]:
    """Subsample the dataset based on the specified parameters.

    Args:
        eval_strategy (EvalStrategy, optional): Subsampling strategy to use instead of self.eval_strategy. Defaults to None.
        block_idx (List[int] | None, optional): Specific block index or indices to evaluate, overriding eval_strategy. Defaults to None.

    Returns:
        Tuple[List[str], List[str]]: Subsampled input data and labels.
    """
    if block_idx is not None:
        indices: List[int] = []
        for idx in block_idx:
            start_idx = idx * self.n_subsamples
            end_idx = min((idx + 1) * self.n_subsamples, len(self.xs))
            indices.extend(range(start_idx, end_idx))

        return [self.xs[i] for i in indices], [self.ys[i] for i in indices]

    if eval_strategy is None:
        eval_strategy = self.eval_strategy

    if eval_strategy in ["full", "evaluated"]:
        return self.xs, self.ys
    elif eval_strategy == "subsample":
        indices = self.rng.choice(len(self.xs), min(self.n_subsamples, len(self.xs)), replace=False)
        return [self.xs[i] for i in indices], [self.ys[i] for i in indices]
    elif eval_strategy == "random_block":
        block_id = self.rng.integers(0, self.n_blocks)
        start_idx = block_id * self.n_subsamples
        end_idx = min((block_id + 1) * self.n_subsamples, len(self.xs))
        indices = np.arange(start_idx, end_idx)
        return [self.xs[i] for i in indices], [self.ys[i] for i in indices]
    elif eval_strategy == "sequential_block":
        # Handle case where self.block_idx is a list
        if isinstance(self.block_idx, list):
            indices_list: List[int] = []
            for idx in self.block_idx:
                start_idx = idx * self.n_subsamples
                end_idx = min((idx + 1) * self.n_subsamples, len(self.xs))
                indices_list.extend(range(start_idx, end_idx))
            return [self.xs[i] for i in indices_list], [self.ys[i] for i in indices_list]
        else:
            start_idx = self.block_idx * self.n_subsamples
            end_idx = min((self.block_idx + 1) * self.n_subsamples, len(self.xs))
            indices = np.arange(start_idx, end_idx)
            return [self.xs[i] for i in indices], [self.ys[i] for i in indices]
    else:
        raise ValueError(f"Unknown subsampling strategy: '{eval_strategy}'")

EvalResult dataclass

Evaluation outputs including scores, sequences, and costs.

Source code in promptolution/tasks/base_task.py 
27
28
29
30
31
32
33
34
35
36
37
@dataclass
class EvalResult:
    """Evaluation outputs including scores, sequences, and costs."""

    scores: np.ndarray  # shape: (n_prompts, n_datapoints)
    agg_scores: np.ndarray  # shape: (n_prompts,) - mean over datapoints
    sequences: np.ndarray  # shape: (n_prompts, n_datapoints)
    input_tokens: np.ndarray  # shape: (n_prompts, n_datapoints)
    output_tokens: np.ndarray  # shape: (n_prompts, n_datapoints)
    agg_input_tokens: np.ndarray  # shape: (n_prompts,) - mean over datapoints
    agg_output_tokens: np.ndarray  # shape: (n_prompts,) - mean over datapoints

classification_tasks

Module for classification tasks.

ClassificationTask

Bases: BaseTask

A class representing a classification task in the promptolution library.

This class handles the loading and management of classification datasets, as well as the evaluation of predictors on these datasets.

Source code in promptolution/tasks/classification_tasks.py 
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
47
48
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
class ClassificationTask(BaseTask):
    """A class representing a classification task in the promptolution library.

    This class handles the loading and management of classification datasets,
    as well as the evaluation of predictors on these datasets.
    """

    def __init__(
        self,
        df: pd.DataFrame,
        task_description: Optional[str] = None,
        x_column: str = "x",
        y_column: str = "y",
        n_subsamples: int = 30,
        eval_strategy: Literal["full", "subsample", "sequential_block", "random_block"] = "full",
        seed: int = 42,
        metric: Callable[[Any, Any], float] = accuracy_score,
        config: Optional["ExperimentConfig"] = None,
    ) -> None:
        """Initialize the ClassificationTask from a pandas DataFrame.

        Args:
            df (pd.DataFrame): Input DataFrame containing the data
            task_description (str): Description of the task
            x_column (str, optional): Name of the column containing input texts. Defaults to "x".
            y_column (str, optional): Name of the column containing labels. Defaults to "y".
            n_subsamples (int, optional): Number of subsamples to use. No subsampling if None. Defaults to None.
            eval_strategy (str, optional): Subsampling strategy to use. Options:
                - "full": Uses the entire dataset for evaluation.
                - "evaluated": Uses only previously evaluated datapoints from the cache.
                - "subsample": Randomly selects n_subsamples datapoints without replacement.
                - "sequential_block": Uses a block of block_size consecutive datapoints, advancing through blocks sequentially.
                - "random_block": Randomly selects a block of block_size consecutive datapoints.
                Defaults to "full".
            seed (int, optional): Random seed for reproducibility. Defaults to 42.
            metric (Callable, optional): Metric to use for evaluation. Defaults to accuracy_score.
            config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
        """
        self.metric = metric
        super().__init__(
            df=df,
            x_column=x_column,
            y_column=y_column,
            task_description=task_description,
            n_subsamples=n_subsamples,
            eval_strategy=eval_strategy,
            seed=seed,
            config=config,
        )
        self.task_type = "classification"
        self.ys: List[str] = (
            df[self.y_column].str.lower().values.tolist()
        )  # Ensure y values are lowercase for consistent comparison
        self.classes = np.unique(self.ys)

    def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarray:
        """Calculate the score for a single prediction."""
        scores = [self.metric([y], [pred]) for pred, y in zip(preds, ys)]
        return np.asarray(scores, dtype=float)

__init__(df, task_description=None, x_column='x', y_column='y', n_subsamples=30, eval_strategy='full', seed=42, metric=accuracy_score, config=None)

Initialize the ClassificationTask from a pandas DataFrame.

Parameters:

Name Type Description Default
df DataFrame

Input DataFrame containing the data

 required
task_description str

Description of the task

 None
x_column str

Name of the column containing input texts. Defaults to "x".

 'x'
y_column str

Name of the column containing labels. Defaults to "y".

 'y'
n_subsamples int

Number of subsamples to use. No subsampling if None. Defaults to None.

 30
eval_strategy str

Subsampling strategy to use. Options: - "full": Uses the entire dataset for evaluation. - "evaluated": Uses only previously evaluated datapoints from the cache. - "subsample": Randomly selects n_subsamples datapoints without replacement. - "sequential_block": Uses a block of block_size consecutive datapoints, advancing through blocks sequentially. - "random_block": Randomly selects a block of block_size consecutive datapoints. Defaults to "full".

 'full'
seed int

Random seed for reproducibility. Defaults to 42.

 42
metric Callable

Metric to use for evaluation. Defaults to accuracy_score.

 accuracy_score
config ExperimentConfig

Configuration for the task, overriding defaults.

 None
Source code in promptolution/tasks/classification_tasks.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def __init__(
    self,
    df: pd.DataFrame,
    task_description: Optional[str] = None,
    x_column: str = "x",
    y_column: str = "y",
    n_subsamples: int = 30,
    eval_strategy: Literal["full", "subsample", "sequential_block", "random_block"] = "full",
    seed: int = 42,
    metric: Callable[[Any, Any], float] = accuracy_score,
    config: Optional["ExperimentConfig"] = None,
) -> None:
    """Initialize the ClassificationTask from a pandas DataFrame.

    Args:
        df (pd.DataFrame): Input DataFrame containing the data
        task_description (str): Description of the task
        x_column (str, optional): Name of the column containing input texts. Defaults to "x".
        y_column (str, optional): Name of the column containing labels. Defaults to "y".
        n_subsamples (int, optional): Number of subsamples to use. No subsampling if None. Defaults to None.
        eval_strategy (str, optional): Subsampling strategy to use. Options:
            - "full": Uses the entire dataset for evaluation.
            - "evaluated": Uses only previously evaluated datapoints from the cache.
            - "subsample": Randomly selects n_subsamples datapoints without replacement.
            - "sequential_block": Uses a block of block_size consecutive datapoints, advancing through blocks sequentially.
            - "random_block": Randomly selects a block of block_size consecutive datapoints.
            Defaults to "full".
        seed (int, optional): Random seed for reproducibility. Defaults to 42.
        metric (Callable, optional): Metric to use for evaluation. Defaults to accuracy_score.
        config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
    """
    self.metric = metric
    super().__init__(
        df=df,
        x_column=x_column,
        y_column=y_column,
        task_description=task_description,
        n_subsamples=n_subsamples,
        eval_strategy=eval_strategy,
        seed=seed,
        config=config,
    )
    self.task_type = "classification"
    self.ys: List[str] = (
        df[self.y_column].str.lower().values.tolist()
    )  # Ensure y values are lowercase for consistent comparison
    self.classes = np.unique(self.ys)

judge_tasks

Module for judge tasks.

JudgeTask

Bases: BaseTask

Task that evaluates a predictor using an LLM-as-a-judge, optionally accepting a ground truth.

Source code in promptolution/tasks/judge_tasks.py 
 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
 89
 90
 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
class JudgeTask(BaseTask):
    """Task that evaluates a predictor using an LLM-as-a-judge, optionally accepting a ground truth."""

    def __init__(
        self,
        df: pd.DataFrame,
        judge_llm: "BaseLLM",
        x_column: str = "x",
        y_column: Optional[str] = None,
        task_description: Optional[str] = None,
        n_subsamples: int = 30,
        eval_strategy: "EvalStrategy" = "full",
        seed: int = 42,
        judge_prompt: Optional[str] = None,
        min_score: float = -5.0,
        max_score: float = 5.0,
        config: "ExperimentConfig" = None,
    ):
        """Initialize the JudgeTask.

        Args:
            df (pd.DataFrame): The input DataFrame containing the data.
            judge_llm (BaseLLM): The LLM judging the predictions.
            x_column (str): Name of the column containing input texts.
            y_column (Optional[str]): Name of the column containing labels/ground truth (if applicable).
            task_description (Optional[str]): Description of the task, parsed to the Judge-LLM and Meta-LLM.
            n_subsamples (int): Number of subsamples to use for evaluation.
            eval_strategy (EvalStrategy): Subsampling strategy to use for evaluation.
            seed (int): Random seed for reproducibility.
            judge_prompt (Optional[str]): Custom prompt for the judge. Note: The score of the Judge will be extracted inside <final_score> tags.
            min_score (float): Minimum score for evaluation.
            max_score (float): Maximum score for evaluation.
            config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
        """
        if judge_prompt is None:
            judge_prompt = JUDGE_PROMPT_WITH_GROUND_TRUTH if y_column else JUDGE_PROMPT_WITHOUT_GROUND_TRUTH
        self.judge_prompt = judge_prompt
        self.min_score = min_score
        self.max_score = max_score

        super().__init__(
            df=df,
            x_column=x_column,
            y_column=y_column,
            task_description=task_description,
            n_subsamples=n_subsamples,
            eval_strategy=eval_strategy,
            seed=seed,
            config=config,
        )
        self.judge_llm = judge_llm
        self.task_type = "judge"

    def _construct_judge_prompt(self, x: str, pred: str, y: Optional[str] = None) -> str:
        """Construct the judge prompt based on whether ground truth is available."""
        if y is not None:
            prompt = self.judge_prompt.replace("{ground_truth}", str(y))
        else:
            prompt = self.judge_prompt

        task_description = self.task_description or ""
        prompt = prompt.replace("{task}", task_description).replace("{input}", x).replace("{prediction}", pred)
        return prompt

    def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarray:
        """Calculate the score for a single prediction using the LLM judge."""
        prompts: List[str] = []
        for x, y, pred in zip(xs, ys, preds):
            judge_prompt = self._construct_judge_prompt(x, pred, y)
            prompts.append(judge_prompt)
        judge_responses = self.judge_llm.get_response(prompts)
        scores_str = extract_from_tag(judge_responses, "<final_score>", "</final_score>")
        scores = []
        for score_str in scores_str:
            try:
                # only numeric chars, - or . are allowed
                score_str = "".join(filter(lambda c: c.isdigit() or c in "-.", score_str))
                score = float(score_str)
                # normalize from [min_score, max_score] to [0, 1]
                score = (score - self.min_score) / (self.max_score - self.min_score)
                score = max(0.0, min(1.0, score))
            except ValueError:
                logger.warning(f"Failed to parse score '{score_str}' as float. Defaulting to a score 0.0.")
                score = 0.0

            scores.append(score)

        return np.asarray(scores, dtype=float)

__init__(df, judge_llm, x_column='x', y_column=None, task_description=None, n_subsamples=30, eval_strategy='full', seed=42, judge_prompt=None, min_score=-5.0, max_score=5.0, config=None)

Initialize the JudgeTask.

Parameters:

Name Type Description Default
df DataFrame

The input DataFrame containing the data.

 required
judge_llm BaseLLM

The LLM judging the predictions.

 required
x_column str

Name of the column containing input texts.

 'x'
y_column Optional[str]

Name of the column containing labels/ground truth (if applicable).

 None
task_description Optional[str]

Description of the task, parsed to the Judge-LLM and Meta-LLM.

 None
n_subsamples int

Number of subsamples to use for evaluation.

 30
eval_strategy EvalStrategy

Subsampling strategy to use for evaluation.

 'full'
seed int

Random seed for reproducibility.

 42
judge_prompt Optional[str]

Custom prompt for the judge. Note: The score of the Judge will be extracted inside tags.

 None
min_score float

Minimum score for evaluation.

 -5.0
max_score float

Maximum score for evaluation.

 5.0
config ExperimentConfig

Configuration for the task, overriding defaults.

 None
Source code in promptolution/tasks/judge_tasks.py 
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
def __init__(
    self,
    df: pd.DataFrame,
    judge_llm: "BaseLLM",
    x_column: str = "x",
    y_column: Optional[str] = None,
    task_description: Optional[str] = None,
    n_subsamples: int = 30,
    eval_strategy: "EvalStrategy" = "full",
    seed: int = 42,
    judge_prompt: Optional[str] = None,
    min_score: float = -5.0,
    max_score: float = 5.0,
    config: "ExperimentConfig" = None,
):
    """Initialize the JudgeTask.

    Args:
        df (pd.DataFrame): The input DataFrame containing the data.
        judge_llm (BaseLLM): The LLM judging the predictions.
        x_column (str): Name of the column containing input texts.
        y_column (Optional[str]): Name of the column containing labels/ground truth (if applicable).
        task_description (Optional[str]): Description of the task, parsed to the Judge-LLM and Meta-LLM.
        n_subsamples (int): Number of subsamples to use for evaluation.
        eval_strategy (EvalStrategy): Subsampling strategy to use for evaluation.
        seed (int): Random seed for reproducibility.
        judge_prompt (Optional[str]): Custom prompt for the judge. Note: The score of the Judge will be extracted inside <final_score> tags.
        min_score (float): Minimum score for evaluation.
        max_score (float): Maximum score for evaluation.
        config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
    """
    if judge_prompt is None:
        judge_prompt = JUDGE_PROMPT_WITH_GROUND_TRUTH if y_column else JUDGE_PROMPT_WITHOUT_GROUND_TRUTH
    self.judge_prompt = judge_prompt
    self.min_score = min_score
    self.max_score = max_score

    super().__init__(
        df=df,
        x_column=x_column,
        y_column=y_column,
        task_description=task_description,
        n_subsamples=n_subsamples,
        eval_strategy=eval_strategy,
        seed=seed,
        config=config,
    )
    self.judge_llm = judge_llm
    self.task_type = "judge"

multi_objective_task

Multi-objective task wrapper that evaluates prompts across multiple tasks.

MultiObjectiveEvalResult dataclass

Container for per-task evaluation outputs in multi-objective runs.

Source code in promptolution/tasks/multi_objective_task.py 
15
16
17
18
19
20
21
22
23
24
25
@dataclass
class MultiObjectiveEvalResult:
    """Container for per-task evaluation outputs in multi-objective runs."""

    scores: List[np.ndarray]
    agg_scores: List[np.ndarray]
    sequences: np.ndarray
    input_tokens: np.ndarray
    output_tokens: np.ndarray
    agg_input_tokens: np.ndarray
    agg_output_tokens: np.ndarray

MultiObjectiveTask

Bases: BaseTask

A task that aggregates evaluations across multiple underlying tasks.

Source code in promptolution/tasks/multi_objective_task.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 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
 89
 90
 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
171
172
173
174
175
176
177
178
179
180
181
182
183
184
class MultiObjectiveTask(BaseTask):
    """A task that aggregates evaluations across multiple underlying tasks."""

    def __init__(
        self,
        tasks: List[BaseTask],
        eval_strategy: Optional[EvalStrategy] = None,
    ) -> None:
        """Initialize with a list of tasks sharing subsampling and seed settings."""
        if not tasks:
            raise ValueError("tasks must be a non-empty list")

        primary = tasks[0]
        for t in tasks[1:]:
            assert t.n_subsamples == primary.n_subsamples, "All tasks must share n_subsamples"
            assert t.seed == primary.seed, "All tasks must share seed"
            assert t.eval_strategy == primary.eval_strategy, "All tasks must share eval_strategy"

        combined_description = "This task is a combination of the following tasks:\n" + "\n".join(
            [f"Task: {t.task_description}" for t in tasks if t.task_description]
        )

        super().__init__(
            df=primary.df,
            x_column=primary.x_column,
            y_column=primary.y_column,
            task_description=combined_description,
            n_subsamples=primary.n_subsamples,
            eval_strategy=eval_strategy or primary.eval_strategy,
            seed=primary.seed,
            config=None,
        )
        self.task_type = "multi"
        self.tasks = tasks
        self._scalarized_objective: bool = False

    def activate_scalarized_objective(self) -> None:
        """Force single-objective behavior by equally averaging task scores."""
        self._scalarized_objective = True

    def evaluate(  # type: ignore
        self,
        prompts: Prompt | List[Prompt],
        predictor,
        system_prompts: Optional[str | List[str]] = None,
        eval_strategy: Optional[EvalStrategy] = None,
    ) -> MultiObjectiveEvalResult | EvalResult:
        """Run prediction once, then score via each task's _evaluate."""
        prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts)
        strategy = eval_strategy or self.eval_strategy

        # Keep block alignment across tasks so block-based strategies stay in sync.
        for task in self.tasks:
            task.block_idx = self.block_idx

        xs, ys = self.subsample(eval_strategy=strategy)

        # Collect all uncached prompt/x/y triples across tasks to predict only once.
        prompts_to_evaluate: List[str] = []
        xs_to_evaluate: List[str] = []
        ys_to_evaluate: List[str] = []
        key_to_index: Dict[Tuple[str, str, str], int] = {}
        cache_keys: List[Tuple[str, str, str]] = []

        for task in self.tasks:
            t_prompts, t_xs, t_ys, t_keys = task._prepare_batch(prompts_list, xs, ys, eval_strategy=strategy)
            for prompt_str, x_val, y_val, key in zip(t_prompts, t_xs, t_ys, t_keys):
                if key in key_to_index:
                    continue
                key_to_index[key] = len(prompts_to_evaluate)
                prompts_to_evaluate.append(prompt_str)
                xs_to_evaluate.append(x_val)
                ys_to_evaluate.append(y_val)
                cache_keys.append(key)

        preds: List[str] = []
        pred_seqs: List[str] = []
        if prompts_to_evaluate:
            preds, pred_seqs = predictor.predict(
                prompts=prompts_to_evaluate,
                xs=xs_to_evaluate,
                system_prompts=system_prompts,
            )

        # Map predictions back to each task and populate caches via _evaluate.
        key_to_pred: Dict[Tuple[str, str, str], Tuple[str, str]] = {
            key: (preds[idx], pred_seqs[idx]) for key, idx in key_to_index.items()
        }

        per_task_results: List[EvalResult] = []
        for task in self.tasks:
            if cache_keys:
                xs_eval = [k[1] for k in cache_keys]
                ys_eval = [k[2] for k in cache_keys]
                preds_eval = [key_to_pred[k][0] for k in cache_keys]
                scores = task._evaluate(xs_eval, ys_eval, preds_eval)
                for score, cache_key in zip(scores, cache_keys):
                    task.eval_cache[cache_key] = score
                    task.seq_cache[cache_key] = key_to_pred[cache_key][1]

            scores_array, agg_scores, seqs = task._collect_results_from_cache(prompts_list, xs, ys)
            input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = task._compute_costs(
                prompts_list, xs, ys, predictor
            )

            per_task_results.append(
                EvalResult(
                    scores=scores_array,
                    agg_scores=agg_scores,
                    sequences=seqs,
                    input_tokens=input_tokens,
                    output_tokens=output_tokens,
                    agg_input_tokens=agg_input_tokens,
                    agg_output_tokens=agg_output_tokens,
                )
            )

        stacked_scores = [r.scores for r in per_task_results]
        stacked_agg_scores = [r.agg_scores for r in per_task_results]

        # Record evaluated blocks for this evaluation (mirroring BaseTask behavior)
        for prompt in prompts_list:
            # Use self.block_idx (the MultiObjectiveTask's block_idx) if in a block strategy
            if strategy in ["sequential_block", "random_block"]:
                if isinstance(self.block_idx, list):
                    self.prompt_evaluated_blocks.setdefault(prompt, []).extend(self.block_idx)
                else:
                    self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx)
            elif strategy == "full":
                self.prompt_evaluated_blocks.setdefault(prompt, []).extend(list(range(self.n_blocks)))

        # Use first task's result for sequences and token counts (they're all the same across tasks)
        first_result = per_task_results[0]

        if self._scalarized_objective:
            return EvalResult(
                scores=np.mean(stacked_scores, axis=0),
                agg_scores=np.mean(stacked_agg_scores, axis=0),
                sequences=first_result.sequences,
                input_tokens=first_result.input_tokens,
                output_tokens=first_result.output_tokens,
                agg_input_tokens=first_result.agg_input_tokens,
                agg_output_tokens=first_result.agg_output_tokens,
            )

        return MultiObjectiveEvalResult(
            scores=stacked_scores,
            agg_scores=stacked_agg_scores,
            sequences=first_result.sequences,
            input_tokens=first_result.input_tokens,
            output_tokens=first_result.output_tokens,
            agg_input_tokens=first_result.agg_input_tokens,
            agg_output_tokens=first_result.agg_output_tokens,
        )

    def _evaluate(self, xs, ys, preds):  # pragma: no cover
        raise NotImplementedError("MultiObjectiveTask overrides evaluate directly")

__init__(tasks, eval_strategy=None)

Initialize with a list of tasks sharing subsampling and seed settings.

Source code in promptolution/tasks/multi_objective_task.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def __init__(
    self,
    tasks: List[BaseTask],
    eval_strategy: Optional[EvalStrategy] = None,
) -> None:
    """Initialize with a list of tasks sharing subsampling and seed settings."""
    if not tasks:
        raise ValueError("tasks must be a non-empty list")

    primary = tasks[0]
    for t in tasks[1:]:
        assert t.n_subsamples == primary.n_subsamples, "All tasks must share n_subsamples"
        assert t.seed == primary.seed, "All tasks must share seed"
        assert t.eval_strategy == primary.eval_strategy, "All tasks must share eval_strategy"

    combined_description = "This task is a combination of the following tasks:\n" + "\n".join(
        [f"Task: {t.task_description}" for t in tasks if t.task_description]
    )

    super().__init__(
        df=primary.df,
        x_column=primary.x_column,
        y_column=primary.y_column,
        task_description=combined_description,
        n_subsamples=primary.n_subsamples,
        eval_strategy=eval_strategy or primary.eval_strategy,
        seed=primary.seed,
        config=None,
    )
    self.task_type = "multi"
    self.tasks = tasks
    self._scalarized_objective: bool = False

activate_scalarized_objective()

Force single-objective behavior by equally averaging task scores.

Source code in promptolution/tasks/multi_objective_task.py 
64
65
66
def activate_scalarized_objective(self) -> None:
    """Force single-objective behavior by equally averaging task scores."""
    self._scalarized_objective = True

evaluate(prompts, predictor, system_prompts=None, eval_strategy=None)

Run prediction once, then score via each task's _evaluate.

Source code in promptolution/tasks/multi_objective_task.py 
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 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
171
172
173
174
175
176
177
178
179
180
181
def evaluate(  # type: ignore
    self,
    prompts: Prompt | List[Prompt],
    predictor,
    system_prompts: Optional[str | List[str]] = None,
    eval_strategy: Optional[EvalStrategy] = None,
) -> MultiObjectiveEvalResult | EvalResult:
    """Run prediction once, then score via each task's _evaluate."""
    prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts)
    strategy = eval_strategy or self.eval_strategy

    # Keep block alignment across tasks so block-based strategies stay in sync.
    for task in self.tasks:
        task.block_idx = self.block_idx

    xs, ys = self.subsample(eval_strategy=strategy)

    # Collect all uncached prompt/x/y triples across tasks to predict only once.
    prompts_to_evaluate: List[str] = []
    xs_to_evaluate: List[str] = []
    ys_to_evaluate: List[str] = []
    key_to_index: Dict[Tuple[str, str, str], int] = {}
    cache_keys: List[Tuple[str, str, str]] = []

    for task in self.tasks:
        t_prompts, t_xs, t_ys, t_keys = task._prepare_batch(prompts_list, xs, ys, eval_strategy=strategy)
        for prompt_str, x_val, y_val, key in zip(t_prompts, t_xs, t_ys, t_keys):
            if key in key_to_index:
                continue
            key_to_index[key] = len(prompts_to_evaluate)
            prompts_to_evaluate.append(prompt_str)
            xs_to_evaluate.append(x_val)
            ys_to_evaluate.append(y_val)
            cache_keys.append(key)

    preds: List[str] = []
    pred_seqs: List[str] = []
    if prompts_to_evaluate:
        preds, pred_seqs = predictor.predict(
            prompts=prompts_to_evaluate,
            xs=xs_to_evaluate,
            system_prompts=system_prompts,
        )

    # Map predictions back to each task and populate caches via _evaluate.
    key_to_pred: Dict[Tuple[str, str, str], Tuple[str, str]] = {
        key: (preds[idx], pred_seqs[idx]) for key, idx in key_to_index.items()
    }

    per_task_results: List[EvalResult] = []
    for task in self.tasks:
        if cache_keys:
            xs_eval = [k[1] for k in cache_keys]
            ys_eval = [k[2] for k in cache_keys]
            preds_eval = [key_to_pred[k][0] for k in cache_keys]
            scores = task._evaluate(xs_eval, ys_eval, preds_eval)
            for score, cache_key in zip(scores, cache_keys):
                task.eval_cache[cache_key] = score
                task.seq_cache[cache_key] = key_to_pred[cache_key][1]

        scores_array, agg_scores, seqs = task._collect_results_from_cache(prompts_list, xs, ys)
        input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = task._compute_costs(
            prompts_list, xs, ys, predictor
        )

        per_task_results.append(
            EvalResult(
                scores=scores_array,
                agg_scores=agg_scores,
                sequences=seqs,
                input_tokens=input_tokens,
                output_tokens=output_tokens,
                agg_input_tokens=agg_input_tokens,
                agg_output_tokens=agg_output_tokens,
            )
        )

    stacked_scores = [r.scores for r in per_task_results]
    stacked_agg_scores = [r.agg_scores for r in per_task_results]

    # Record evaluated blocks for this evaluation (mirroring BaseTask behavior)
    for prompt in prompts_list:
        # Use self.block_idx (the MultiObjectiveTask's block_idx) if in a block strategy
        if strategy in ["sequential_block", "random_block"]:
            if isinstance(self.block_idx, list):
                self.prompt_evaluated_blocks.setdefault(prompt, []).extend(self.block_idx)
            else:
                self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx)
        elif strategy == "full":
            self.prompt_evaluated_blocks.setdefault(prompt, []).extend(list(range(self.n_blocks)))

    # Use first task's result for sequences and token counts (they're all the same across tasks)
    first_result = per_task_results[0]

    if self._scalarized_objective:
        return EvalResult(
            scores=np.mean(stacked_scores, axis=0),
            agg_scores=np.mean(stacked_agg_scores, axis=0),
            sequences=first_result.sequences,
            input_tokens=first_result.input_tokens,
            output_tokens=first_result.output_tokens,
            agg_input_tokens=first_result.agg_input_tokens,
            agg_output_tokens=first_result.agg_output_tokens,
        )

    return MultiObjectiveEvalResult(
        scores=stacked_scores,
        agg_scores=stacked_agg_scores,
        sequences=first_result.sequences,
        input_tokens=first_result.input_tokens,
        output_tokens=first_result.output_tokens,
        agg_input_tokens=first_result.agg_input_tokens,
        agg_output_tokens=first_result.agg_output_tokens,
    )

reward_tasks

Module for Reward tasks.

RewardTask

Bases: BaseTask

A task that evaluates a predictor using a reward function.

This task takes a DataFrame, a column name for input data, and a reward function. The reward function takes in a prediction as input and returns a scalar reward.

Source code in promptolution/tasks/reward_tasks.py 
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
47
48
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
class RewardTask(BaseTask):
    """A task that evaluates a predictor using a reward function.

    This task takes a DataFrame, a column name for input data, and a reward function.
    The reward function takes in a prediction as input and returns a scalar reward.
    """

    def __init__(
        self,
        df: pd.DataFrame,
        reward_function: Callable[[str], float],
        x_column: str = "x",
        y_column: Optional[str] = None,
        reward_columns: Optional[List[str]] = None,
        task_description: Optional[str] = None,
        n_subsamples: int = 30,
        eval_strategy: "EvalStrategy" = "full",
        seed: int = 42,
        config: Optional["ExperimentConfig"] = None,
    ) -> None:
        """Initialize the RewardTask.

        Args:
            df (pd.DataFrame): Input DataFrame containing the data.
            reward_function (Callable): Function that takes a prediction, potential keyword arguments from the dataframe, and returns a reward score. Note: The optimizers aim to maximize.
            x_column (str, optional): Name of the column containing input texts. Defaults to "x".
            y_column (str, optional): Name of the column containing target texts if available. Defaults to None.
            reward_columns (List[str], optional): Additional dataframe columns to pass as keyword args to reward_function.
            task_description (str, optional): Description of the task.
            n_subsamples (int, optional): Number of subsamples to use. Defaults to 30.
            eval_strategy (str, optional): Subsampling strategy to use. Defaults to "full".
            seed (int, optional): Random seed for reproducibility. Defaults to 42.
            config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
        """
        self.reward_function = reward_function
        self.reward_columns = reward_columns or []
        super().__init__(
            df=df,
            x_column=x_column,
            y_column=y_column,
            task_description=task_description,
            n_subsamples=n_subsamples,
            eval_strategy=eval_strategy,
            seed=seed,
            config=config,
        )
        self.task_type = "reward"
        # x -> kwargs to reward function
        km = self.df.set_index(x_column)[self.reward_columns].to_dict("index")
        self.kwargs_map = defaultdict(dict, km)

    def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarray:
        """Calculate reward for each prediction, passing configured columns as kwargs."""
        kwargs_list = [self.kwargs_map[x] for x in xs]
        rewards = [self.reward_function(pred, **kwargs) for pred, kwargs in zip(preds, kwargs_list)]
        return np.asarray(rewards, dtype=float)

__init__(df, reward_function, x_column='x', y_column=None, reward_columns=None, task_description=None, n_subsamples=30, eval_strategy='full', seed=42, config=None)

Initialize the RewardTask.

Parameters:

Name Type Description Default
df DataFrame

Input DataFrame containing the data.

 required
reward_function Callable

Function that takes a prediction, potential keyword arguments from the dataframe, and returns a reward score. Note: The optimizers aim to maximize.

 required
x_column str

Name of the column containing input texts. Defaults to "x".

 'x'
y_column str

Name of the column containing target texts if available. Defaults to None.

 None
reward_columns List[str]

Additional dataframe columns to pass as keyword args to reward_function.

 None
task_description str

Description of the task.

 None
n_subsamples int

Number of subsamples to use. Defaults to 30.

 30
eval_strategy str

Subsampling strategy to use. Defaults to "full".

 'full'
seed int

Random seed for reproducibility. Defaults to 42.

 42
config ExperimentConfig

Configuration for the task, overriding defaults.

 None
Source code in promptolution/tasks/reward_tasks.py 
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def __init__(
    self,
    df: pd.DataFrame,
    reward_function: Callable[[str], float],
    x_column: str = "x",
    y_column: Optional[str] = None,
    reward_columns: Optional[List[str]] = None,
    task_description: Optional[str] = None,
    n_subsamples: int = 30,
    eval_strategy: "EvalStrategy" = "full",
    seed: int = 42,
    config: Optional["ExperimentConfig"] = None,
) -> None:
    """Initialize the RewardTask.

    Args:
        df (pd.DataFrame): Input DataFrame containing the data.
        reward_function (Callable): Function that takes a prediction, potential keyword arguments from the dataframe, and returns a reward score. Note: The optimizers aim to maximize.
        x_column (str, optional): Name of the column containing input texts. Defaults to "x".
        y_column (str, optional): Name of the column containing target texts if available. Defaults to None.
        reward_columns (List[str], optional): Additional dataframe columns to pass as keyword args to reward_function.
        task_description (str, optional): Description of the task.
        n_subsamples (int, optional): Number of subsamples to use. Defaults to 30.
        eval_strategy (str, optional): Subsampling strategy to use. Defaults to "full".
        seed (int, optional): Random seed for reproducibility. Defaults to 42.
        config (ExperimentConfig, optional): Configuration for the task, overriding defaults.
    """
    self.reward_function = reward_function
    self.reward_columns = reward_columns or []
    super().__init__(
        df=df,
        x_column=x_column,
        y_column=y_column,
        task_description=task_description,
        n_subsamples=n_subsamples,
        eval_strategy=eval_strategy,
        seed=seed,
        config=config,
    )
    self.task_type = "reward"
    # x -> kwargs to reward function
    km = self.df.set_index(x_column)[self.reward_columns].to_dict("index")
    self.kwargs_map = defaultdict(dict, km)