From ca6dc91e38875ecfeb6fdfe15a8ac0f272848b0e Mon Sep 17 00:00:00 2001 From: finitearth Date: Sat, 27 Dec 2025 12:42:19 +0100 Subject: [PATCH 01/53] set up capoeira --- promptolution/optimizers/capoeira.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 promptolution/optimizers/capoeira.py diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py new file mode 100644 index 00000000..e69de29b From 1bc1c72993fc64e70ef680ac1c9cdd50562f0afc Mon Sep 17 00:00:00 2001 From: finitearth Date: Sat, 27 Dec 2025 13:28:10 +0100 Subject: [PATCH 02/53] copy pased optimizer from research repo --- promptolution/optimizers/capoeira.py | 1159 ++++++++++++++++++++++++++ 1 file changed, 1159 insertions(+) diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index e69de29b..125fcc8d 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -0,0 +1,1159 @@ +""" +Implementation of the MO-CAPO (Multi-Objective Cost-Aware Prompt Optimization) algorithm. +Contains the MOCAPOptimizer class, which manages the prompt optimization process using +intensification techniques for multi-objective optimization. + +This is the new multi-objective version. The original single-objective CAPO remains +in capo.py for backward compatibility and comparative experiments. +""" + +import random +from logging import getLogger +from typing import Any, Callable, Dict, List, Tuple + +import numpy as np +import pandas as pd +from promptolution.llms.base_llm import BaseLLM +from promptolution.optimizers.base_optimizer import BaseOptimizer +from promptolution.predictors.base_predictor import BasePredictor +from promptolution.tasks.base_task import BaseTask +from promptolution.utils.prompt import Prompt + +from capo.mo_task import MOCAPOClassificationTask +from capo.runhistory import RunHistory +from capo.templates import CROSSOVER_TEMPLATE, FEWSHOT_TEMPLATE, MUTATION_TEMPLATE +from capo.utils import seed_everything + +# ### HELPER FUNCTIONS FOR MULTI-OBJECTIVE OPTIMIZATION ### + + +def fast_non_dominated_sort(obj_vectors: np.ndarray) -> list[list[int]]: + """ + Performs a fast non-dominated sort on a set of objective vectors. + This is a standard algorithm from NSGA-II. + + Args: + obj_vectors: A numpy array of shape (n_solutions, n_objectives). + + Returns: + A list of fronts, where each front is a list of indices corresponding + to the input obj_vectors. The first front is the Pareto-optimal set. + """ + num_solutions = obj_vectors.shape[0] + if num_solutions == 0: + return [] + + domination_counts = np.zeros(num_solutions, dtype=int) + dominated_solutions = [[] for _ in range(num_solutions)] + fronts = [[]] # The first front + + for i in range(num_solutions): + for j in range(i + 1, num_solutions): + # Assumes minimization for all objectives + is_i_dom_j = np.all(obj_vectors[i] <= obj_vectors[j]) and np.any( + obj_vectors[i] < obj_vectors[j] + ) + is_j_dom_i = np.all(obj_vectors[j] <= obj_vectors[i]) and np.any( + obj_vectors[j] < obj_vectors[i] + ) + + if is_i_dom_j: + dominated_solutions[i].append(j) + domination_counts[j] += 1 + elif is_j_dom_i: + dominated_solutions[j].append(i) + domination_counts[i] += 1 + + # Identify the first front (solutions with domination_count == 0) + for i in range(num_solutions): + if domination_counts[i] == 0: + fronts[0].append(i) + + # Build subsequent fronts + front_idx = 0 + while fronts and front_idx < len(fronts) and fronts[front_idx]: + next_front = [] + for i in fronts[front_idx]: + for j in dominated_solutions[i]: + domination_counts[j] -= 1 + if domination_counts[j] == 0: + next_front.append(j) + front_idx += 1 + if next_front: + fronts.append(next_front) + else: + break + + return fronts + + +def _calculate_crowding_distance(obj_vectors: np.ndarray) -> np.ndarray: + """ + Calculates the crowding distance for each solution in a set. + Used as a tie-breaker for selection and pruning. + + Args: + obj_vectors: A numpy array of shape (n_solutions, n_objectives). + + Returns: + A numpy array of shape (n_solutions,) with the crowding distance for each. + """ + num_solutions, num_obj = obj_vectors.shape + if num_solutions <= 2: + return np.full(num_solutions, float("inf")) + + distances = np.zeros(num_solutions) + + for i in range(num_obj): + sorted_indices = np.argsort(obj_vectors[:, i]) + + distances[sorted_indices[0]] = float("inf") + distances[sorted_indices[-1]] = float("inf") + + f_min = obj_vectors[sorted_indices[0], i] + f_max = obj_vectors[sorted_indices[-1], i] + + if f_max == f_min: + continue + + for j in range(1, num_solutions - 1): + distances[sorted_indices[j]] += ( + obj_vectors[sorted_indices[j + 1], i] - obj_vectors[sorted_indices[j - 1], i] + ) / (f_max - f_min) + + return distances + + +class CAPOEIRA(BaseOptimizer): + """ + Multi-Objective Cost-Aware Prompt Optimizer that evolves prompt instructions + using crossover, mutation, and intensification based on Pareto-optimal fronts + and dominance relationships. + """ + + def __init__( + self, + initial_prompts: List[str], + task: BaseTask, + df_few_shots: pd.DataFrame, + meta_llm: BaseLLM, + downstream_llm: BaseLLM, + crossovers_per_iter: int, + population_size: int, + upper_shots: int, + freeze_p_pop: bool = False, + w_in: float = 1.0, + w_out: float = 1.0, + crossover_meta_prompt: str = None, + mutation_meta_prompt: str = None, + callbacks: List[Callable] = [], + predictor: BasePredictor = None, + verbosity: int = 0, + logger=getLogger(__name__), + intensify_vs_all_incumbents: bool = False, + init_with_intensification: bool = False, + init_on_all_blocks: bool = False, + no_weak_dominance: bool = False, + random_parent_selection: bool = False, + random_pruning: bool = False, + random_seed: int = 42, + ): + """ + Initializes the MO-CAPO optimizer with parameters for multi-objective prompt evolution. + + Parameters: + initial_prompts (List[str]): Initial prompt instructions. + task (BaseTask): The task instance containing dataset and description. + df_few_shots (pd.DataFrame): DataFrame containing few-shot examples. + meta_llm (BaseLLM): The meta language model for crossover/mutation. + downstream_llm (BaseLLM): The downstream language model used for responses. + crossovers_per_iter (int): Number of crossover operations per iteration. + population_size (int): Maximum population size for pruning. + upper_shots (int): Maximum number of few-shot examples per prompt. + w_in (float): Weight for input tokens in cost calculation. + w_out (float): Weight for output tokens in cost calculation. + crossover_meta_prompt (str, optional): Template for crossover instructions. + mutation_meta_prompt (str, optional): Template for mutation instructions. + callbacks (List[Callable], optional): Callbacks for optimizer events. + predictor (BasePredictor, optional): Predictor to evaluate prompt performance. + verbosity (int, optional): Verbosity level for logging. Defaults to 0. + logger: Logger instance for debugging and information output. + """ + assert isinstance(task, MOCAPOClassificationTask), "MOCAPOptimizer requires a MO-CAPO task." + + super().__init__(initial_prompts, task, callbacks, predictor) + self.df_few_shots = df_few_shots + self.meta_llm = meta_llm + self.downstream_llm = downstream_llm + + self.crossover_meta_prompt = crossover_meta_prompt or CROSSOVER_TEMPLATE + self.mutation_meta_prompt = mutation_meta_prompt or MUTATION_TEMPLATE + + self.population_size = population_size + self.crossovers_per_iter = crossovers_per_iter + self.upper_shots = upper_shots + self.freeze_p_pop = freeze_p_pop + self.w_in = w_in + self.w_out = w_out + self.verbosity = verbosity + self.logger = logger + self.intensify_vs_all_incumbents = intensify_vs_all_incumbents + self.init_with_intensification = init_with_intensification + self.init_on_all_blocks = init_on_all_blocks + self.no_weak_dominance = no_weak_dominance + self.random_parent_selection = random_parent_selection + self.random_pruning = random_pruning + self.random_seed = random_seed + seed_everything(self.random_seed) + + self.P_inc: List[Prompt] = [] + self.P_pop: List[Prompt] = [] + self.runhistory = RunHistory() + + # Buffers for minimally invasive lineage tracking + self._selection_details_buffer: List[Dict[str, Any]] = [] + self._crossover_lineage_buffer: List[Dict[str, Any]] = [] + + initial_prompt_objects = self._create_initial_prompts(initial_prompts) + + if self.init_with_intensification: + # Shuffle and intensify each initial prompt + random.shuffle(initial_prompt_objects) + for prompt in initial_prompt_objects: + self._do_intensification(prompt) + else: + # Original initialization + self._initialize_population_and_fronts(initial_prompt_objects) + + def _create_initial_prompts(self, initial_prompts: List[str]) -> List[Prompt]: + """ + Initializes the population of Prompt objects from initial instructions. + + Parameters: + initial_prompts (List[str]): List of initial prompt instructions. + + Returns: + List[Prompt]: Initialized population of prompts with few-shot examples. + """ + population = [] + + for instruction_text in initial_prompts: + num_examples = random.randint(0, self.upper_shots) + few_shots = self._create_few_shot_examples(instruction_text, num_examples) + population.append(Prompt(instruction_text, few_shots)) + + if self.verbosity > 0: + self.logger.warning( + f"🍿Initialized population with {len(population)} prompts: \n {[p.construct_prompt() for p in population]}" + ) + return population + + def _initialize_population_and_fronts(self, initial_prompts: List[Prompt]): + if not initial_prompts: + return + + prompt_strings = [p.construct_prompt() for p in initial_prompts] + + # Determine which blocks to evaluate on during initialization + if self.init_on_all_blocks: + # Evaluate on all blocks + blocks_to_evaluate = list(range(len(self.task.blocks))) + else: + # Default: Evaluate only on a single random block + block_id = random.randrange(len(self.task.blocks)) + blocks_to_evaluate = [block_id] + + # Evaluate initial prompts on selected blocks + for block_id in blocks_to_evaluate: + self.task.evaluate_on_block( + prompt_strings, + block_id, + self.predictor, + self.runhistory, + ) + + # Compute current vectors (will aggregate across all evaluated blocks) + all_vectors = np.array( + [self.runhistory.compute_current_vector(prompt_str) for prompt_str in prompt_strings] + ) + all_fronts = fast_non_dominated_sort(all_vectors) + + inc_indices = all_fronts[0] + self.P_inc = [initial_prompts[i] for i in inc_indices] + + for front in all_fronts[1:]: + self.P_pop.extend([initial_prompts[i] for i in front]) + + if self.verbosity > 0: + self.logger.info( + f"🌱 Population Initialized. Incumbents: {len(self.P_inc)}, Population: {len(self.P_pop)} " + f"(evaluated on {len(blocks_to_evaluate)} block(s))" + ) + + def _create_few_shot_examples( + self, instruction: str, num_examples: int + ) -> List[Tuple[str, str]]: + if num_examples == 0: + return [] + few_shot_samples = self.df_few_shots.sample(num_examples, replace=False) + sample_inputs = few_shot_samples["input"].values + sample_targets = few_shot_samples["target"].values + few_shots = [ + FEWSHOT_TEMPLATE.replace("", i).replace( + "", + f"{self.predictor.begin_marker}{t}{self.predictor.end_marker}", + ) + for i, t in zip(sample_inputs, sample_targets) + ] + # Select partition of the examples to generate reasoning from downstream model + preds, seqs = self.predictor.predict( + instruction, + sample_inputs, + return_seq=True, + ) + preds, seqs = preds.reshape(num_examples), seqs.reshape(num_examples) + + # Check which predictions are correct and get a single one per example + for j in range(num_examples): + # Process and clean up the generated sequences + seqs[j] = seqs[j].replace(sample_inputs[j], "").strip() + # Check if the prediction is correct and add reasoning if so + if preds[j] == sample_targets[j]: + few_shots[j] = FEWSHOT_TEMPLATE.replace("", sample_inputs[j]).replace( + "", seqs[j] + ) + + if self.verbosity > 1: + self.logger.warning(f"🔫Few-shot examples: {few_shots}") + self.logger.warning(f"💆‍♂️Generated reasoning: {seqs}") + + return few_shots + + def _is_dominated(self, vec1: np.ndarray, vec2: np.ndarray) -> bool: + return np.all(vec2 <= vec1) and np.any(vec2 < vec1) + + def _is_weakly_dominated( + self, + prompt_a: Prompt, + prompt_b: Prompt, + str_a: str | None = None, + str_b: str | None = None, + ) -> Tuple[Prompt | None, str]: + """ + Check weak dominance relationship between two prompts according to ParamILS weak dominance. + + Args: + prompt_a: First prompt + prompt_b: Second prompt + str_a: Optional pre-constructed prompt string for prompt_a (optimization) + str_b: Optional pre-constructed prompt string for prompt_b (optimization) + + Returns: + Tuple of (winner, reason) where: + - winner: The selected prompt or None if no clear winner + - reason: One of: + * "better_rank": winner is on a better (lower) front than loser (same blocks) + * "crowding_distance": same blocks, same front, CD used + * "weak_dominance": winner weakly dominates loser (different blocks) + * "no_weak_dominance": no subset relationship exists + """ + # Use pre-constructed strings if provided, otherwise construct them + if str_a is None: + str_a = prompt_a.construct_prompt() + if str_b is None: + str_b = prompt_b.construct_prompt() + + # Get evaluated blocks for both prompts (returns Set[int] for single prompt) + blocks_a = self.runhistory.get_evaluated_blocks(str_a) + blocks_b = self.runhistory.get_evaluated_blocks(str_b) + + # Ensure we have Set[int] type (type narrowing) + assert isinstance(blocks_a, set), f"Expected set, got {type(blocks_a)}" + assert isinstance(blocks_b, set), f"Expected set, got {type(blocks_b)}" + + # CASE 1: Same blocks - use rank (front) then crowding distance + if blocks_a == blocks_b and len(blocks_a) > 0: + # Get all P_pop prompts with the same block set + same_block_prompts = [ + p + for p in self.P_pop + if self.runhistory.get_evaluated_blocks(p.construct_prompt()) == blocks_a + ] + + # Perform NDS to identify fronts + vectors = np.array( + [ + self.runhistory.compute_current_vector(p.construct_prompt(), blocks_a) + for p in same_block_prompts + ] + ) + fronts = fast_non_dominated_sort(vectors) + + # Find indices of prompt_a and prompt_b in same_block_prompts + idx_a = same_block_prompts.index(prompt_a) + idx_b = same_block_prompts.index(prompt_b) + + # Find which front each prompt is on + rank_a = None + rank_b = None + for front_idx, front in enumerate(fronts): + if idx_a in front: + rank_a = front_idx + if idx_b in front: + rank_b = front_idx + if rank_a is not None and rank_b is not None: + break + + # Both prompts must be found in some front + assert rank_a is not None, f"prompt_a (idx={idx_a}) missing from all fronts" + assert rank_b is not None, f"prompt_b (idx={idx_b}) missing from all fronts" + + # Select by rank first (lower rank = better front = better prompt) + if rank_a < rank_b: + return (prompt_a, "better_rank") + elif rank_b < rank_a: + return (prompt_b, "better_rank") + else: + # Same rank - use crowding distance + front = fronts[rank_a] + front_vectors = vectors[front] + distances = _calculate_crowding_distance(front_vectors) + + # Find positions of our prompts in the front + pos_a = front.index(idx_a) + pos_b = front.index(idx_b) + + # Select based on crowding distance (higher is better for diversity) + if distances[pos_a] > distances[pos_b]: + return (prompt_a, "crowding_distance") + elif distances[pos_b] > distances[pos_a]: + return (prompt_b, "crowding_distance") + else: + # Legitimate tie - both have same crowding distance + return (random.choice([prompt_a, prompt_b]), "crowding_distance") + + # CASE 2: Different blocks - check weak dominance (information + performance) + # Check Information Rule: blocks_b ⊆ blocks_a (a weakly dominates b) + a_weakly_dominates_b = blocks_b.issubset(blocks_a) + # Check Information Rule: blocks_a ⊆ blocks_b (b weakly dominates a) + b_weakly_dominates_a = blocks_a.issubset(blocks_b) + + # Early return if no subset relationship exists + if not a_weakly_dominates_b and not b_weakly_dominates_a: + return (None, "no_weak_dominance") + + if a_weakly_dominates_b and len(blocks_b) > 0: + # Check Performance Rule on shared blocks (which is blocks_b) + vec_a = self.runhistory.compute_current_vector(str_a, blocks_b) + vec_b = self.runhistory.compute_current_vector(str_b, blocks_b) + + if self._is_dominated(vec_b, vec_a): # a weakly dominates b + return (prompt_a, "weak_dominance") + + if b_weakly_dominates_a and len(blocks_a) > 0: + # Check Performance Rule on shared blocks (which is blocks_a) + vec_a = self.runhistory.compute_current_vector(str_a, blocks_a) + vec_b = self.runhistory.compute_current_vector(str_b, blocks_a) + + if self._is_dominated(vec_a, vec_b): # b weakly dominates a + return (prompt_b, "weak_dominance") + + return (None, "no_weak_dominance") + + def _select_parent_from_pool(self, selection_pool: List[Prompt]) -> Prompt: + """Select one parent from the selection pool using tournament rules.""" + p1, p2 = random.sample(selection_pool, 2) + details: Dict[str, Any] = { + "sampled_candidates": [p1.construct_prompt(), p2.construct_prompt()], + "reason": "", + } + + p1_is_inc = p1 in self.P_inc + p2_is_inc = p2 in self.P_inc + + winner = None + if p1_is_inc and p2_is_inc: + # Case 1: Both prompts are from incumbent set -> use crowding distance to break ties + # Use B_common (intersection of all incumbent blocks) for fair comparison + inc_strings = [inc.construct_prompt() for inc in self.P_inc] + all_inc_blocks = self.runhistory.get_evaluated_blocks(inc_strings) + + # Ensure we have List[Set[int]] type + assert isinstance(all_inc_blocks, list), f"Expected list, got {type(all_inc_blocks)}" + B_common = set.intersection(*all_inc_blocks) if all_inc_blocks else set() + + inc_vectors = np.array( + [ + self.runhistory.compute_current_vector(inc_str, B_common) + for inc_str in inc_strings + ] + ) + inc_distances = _calculate_crowding_distance(inc_vectors) + + # Find the indices of p1 and p2 within the incumbent set + p1_idx = self.P_inc.index(p1) + p2_idx = self.P_inc.index(p2) + + if inc_distances[p1_idx] > inc_distances[p2_idx]: + winner = p1 + elif inc_distances[p2_idx] > inc_distances[p1_idx]: + winner = p2 + else: + winner = random.choice([p1, p2]) + details["reason"] = "crowding_distance_tie_inc" + + elif p1_is_inc: + # Case 2: One is from pop and one from inc -> use incumbent + winner = p1 + details["reason"] = "incumbent" + elif p2_is_inc: + # Case 2: One is from pop and one from inc -> use incumbent + winner = p2 + details["reason"] = "incumbent" + else: + # Case 3: Both are from population + if self.no_weak_dominance: + # Ablation: Use random sampling directly + winner = random.choice([p1, p2]) + details["reason"] = "random" + else: + # Use weak dominance with crowding distance tie-breaking + weak_dom_winner, reason = self._is_weakly_dominated(p1, p2) + if weak_dom_winner is not None: + winner = weak_dom_winner + details["reason"] = reason + else: + # No relationship - use random sampling + winner = random.choice([p1, p2]) + details["reason"] = "random" + + details["winner"] = winner.construct_prompt() + self._selection_details_buffer.append(details) + return winner + + def _tournament_selection(self) -> Tuple[Prompt, Prompt]: + selection_pool = self.P_inc + self.P_pop + + parent1 = self._select_parent_from_pool(selection_pool) + parent2 = self._select_parent_from_pool(selection_pool) + + # Ensure we don't select the same parent twice + while parent1 == parent2: + # Remove the last selection details and try again + self._selection_details_buffer.pop() + parent2 = self._select_parent_from_pool(selection_pool) + + return parent1, parent2 + + def _crossover(self) -> List[Prompt]: + """ + Performs crossover among parent prompts to generate offsprings. + """ + self._selection_details_buffer = [] # Clear buffer for this generation + crossover_prompts = [] + offspring_few_shots = [] + parent_pairs = [] + for _ in range(self.crossovers_per_iter): + if self.random_parent_selection: + mother, father = random.sample(self.P_inc + self.P_pop, 2) + # For random selection, we don't have detailed tournament data + mother_details = { + "sampled_candidates": [mother.construct_prompt()], # List with single element + "reason": "random_selection", + "winner": mother.construct_prompt(), + } + father_details = { + "sampled_candidates": [father.construct_prompt()], + "reason": "random_selection", + "winner": father.construct_prompt(), + } + self._selection_details_buffer.extend([mother_details, father_details]) + else: + mother, father = self._tournament_selection() + parent_pairs.append((mother, father)) + crossover_prompt = ( + self.crossover_meta_prompt.replace("", mother.instruction_text) + .replace("", father.instruction_text) + .replace("", self.task.description) + .strip() + ) + crossover_prompts.append(crossover_prompt) + combined_few_shots = mother.few_shots + father.few_shots + if combined_few_shots: + num_few_shots = (len(mother.few_shots) + len(father.few_shots)) // 2 + offspring_few_shot = random.sample(combined_few_shots, num_few_shots) + else: + offspring_few_shot = [] + offspring_few_shots.append(offspring_few_shot) + + child_instructions = self.meta_llm.get_response(crossover_prompts) + if self.verbosity > 1: + self.logger.warning(f"🥐Generated crossover prompts: \n{child_instructions}") + + offsprings = [] + self._crossover_lineage_buffer = [] # Clear/prepare for mutation step + for i, (instruction, examples) in enumerate(zip(child_instructions, offspring_few_shots)): + instruction = instruction.split("")[-1].split("")[0].strip() + offspring = Prompt(instruction, examples) + offsprings.append(offspring) + + mother_details = self._selection_details_buffer[i * 2] + father_details = self._selection_details_buffer[i * 2 + 1] + + lineage = { + "step": self.runhistory.current_step, + "mother_selection": mother_details, + "father_selection": father_details, + } + self._crossover_lineage_buffer.append(lineage) + + return offsprings + + def _mutate(self, offsprings: List[Prompt]) -> List[Prompt]: + """ + Applies mutation to offsprings to generate new candidate prompts. + """ + mutation_prompts = [ + self.mutation_meta_prompt.replace("", prompt.instruction_text).replace( + "", self.task.description + ) + for prompt in offsprings + ] + new_instructions = self.meta_llm.get_response(mutation_prompts) + + mutated = [] + for i, (new_instruction, original_offspring) in enumerate( + zip(new_instructions, offsprings) + ): + new_instruction = new_instruction.split("")[-1].split("")[0].strip() + p = random.random() + + if ( + p < 1 / 3 and len(original_offspring.few_shots) < self.upper_shots + ): # add a random few shot + new_few_shot = self._create_few_shot_examples(new_instruction, 1) + new_few_shots = original_offspring.few_shots + new_few_shot + elif ( + 1 / 3 <= p < 2 / 3 and len(original_offspring.few_shots) > 0 + ): # remove a random few shot + new_few_shots = random.sample( + original_offspring.few_shots, len(original_offspring.few_shots) - 1 + ) + else: # do not change few shots, but shuffle + new_few_shots = original_offspring.few_shots + + random.shuffle(new_few_shots) + mutated_prompt = Prompt(new_instruction, new_few_shots) + mutated.append(mutated_prompt) + + # Log full lineage from the buffer + lineage_data = self._crossover_lineage_buffer[i] + lineage_data["offspring"] = original_offspring.construct_prompt() + self.runhistory.add_lineage(mutated_prompt.construct_prompt(), lineage_data.copy()) + + if self.verbosity > 0: + self.logger.warning(f"🧟Generated {len(mutated)} mutated prompts.") + self.logger.warning(f"😶Generated Prompts: {[p.construct_prompt() for p in mutated]}") + + return mutated + + def _get_closest_incumbent(self, challenger_vec: np.ndarray): + """ + Finds the incumbent prompt closest to a challenger in a normalized + multi-objective space using Euclidean distance. This implementation is + vectorized for efficiency. + + Normalization bounds (min and max for each objective) are calculated on-demand + from all evaluations stored in the runhistory to ensure they are up-to-date. + + Args: + challenger_vec: The challenger's objective vector. + challenger: The challenger prompt object. + + Returns: + The incumbent prompt object that is closest to the challenger, or None + if no incumbents exist. + """ + if not self.P_inc: + return None + + # Step 1: Calculate Global Objective Bounds (On-Demand) + all_obj_vectors = self.runhistory.get_all_objective_vectors() + + if all_obj_vectors.shape[0] < 2: + return random.choice(self.P_inc) + + min_bounds = np.min(all_obj_vectors, axis=0) + max_bounds = np.max(all_obj_vectors, axis=0) + + # Step 2: Normalize Vectors + range_val = max_bounds - min_bounds + range_val[range_val == 0] = 1.0 # Avoid division by zero + + norm_chal_vec = (challenger_vec - min_bounds) / range_val + + # Vectorized normalization of all incumbent vectors + inc_vectors = np.array( + [self.runhistory.compute_current_vector(p.construct_prompt()) for p in self.P_inc] + ) + norm_inc_vectors = (inc_vectors - min_bounds) / range_val + + # Step 3: Find Closest Incumbent (Vectorized) + # Calculate Euclidean distance for all incumbents at once + distances = np.linalg.norm(norm_inc_vectors - norm_chal_vec, axis=1) + + # Find the index of the incumbent with the minimum distance + closest_inc_idx = np.argmin(distances) + closest_incumbent = self.P_inc[closest_inc_idx] + + return closest_incumbent + + def _do_intensification(self, challenger: Prompt): + """ + Implements the MO-CAPO intensification algorithm as defined in Algorithm 3. + + Args: + challenger: The challenger prompt to intensify + """ + + # Handle the edge case for initialization with intensification when P_inc is empty + if not self.P_inc: + # Evaluate the first challenger on a single random block + challenger_str = challenger.construct_prompt() + random_block_id = random.choice(range(len(self.task.blocks))) + self.task.evaluate_on_block( + [challenger_str], random_block_id, self.predictor, self.runhistory + ) + + # This first prompt becomes the first incumbent + self.P_inc.append(challenger) + if self.verbosity > 0: + self.logger.info(f"🐣 Initializing with first incumbent: {challenger_str[:30]}...") + return + + # Step 1: Get intersection of all incumbent evaluated blocks (common blocks) + inc_strings = [inc.construct_prompt() for inc in self.P_inc] + all_inc_blocks = self.runhistory.get_evaluated_blocks(inc_strings) + + # Ensure we have List[Set[int]] type (since we passed a list of strings) + assert isinstance(all_inc_blocks, list), f"Expected list, got {type(all_inc_blocks)}" + + # Get intersection of all blocks that ALL incumbents have been evaluated on + B_common = set.intersection(*all_inc_blocks) + + # Step 2: Initialize challenger evaluation + challenger_str = challenger.construct_prompt() + B_eval_challenger = set() # Blocks challenger has been evaluated on + new_cost_vector = np.full(2, np.inf) + # Step 3: Main while loop + while True: + # Step 3.1: Set old_cost_vector = new_cost_vector + old_cost_vector = new_cost_vector.copy() + + # Step 3.2: Repeat-until loop (check conditions at END) + while True: + # Step 3.2.2: Sample random block from B_common \ B_eval_challenger + available_blocks = B_common - B_eval_challenger + + sampled_block = random.choice(list(available_blocks)) + B_eval_challenger.add(sampled_block) + + # Step 3.2.3: Evaluate challenger on sampled block + eval_results = self.task.evaluate_on_block( + [challenger_str], sampled_block, self.predictor, self.runhistory + ) + obj_vec, _, _ = eval_results[0] + + # Update cost vector incrementally (running average) + n_evals = len(B_eval_challenger) + if n_evals == 1: + new_cost_vector = obj_vec.copy() + else: + # Correct incremental average: new_avg = (old_avg * (n-1) + new_val) / n + new_cost_vector = (old_cost_vector * (n_evals - 1) + obj_vec) / n_evals + + # Step 3.2.5: Check exit conditions + condition_1 = self._is_dominated( + old_cost_vector, new_cost_vector + ) # new dominates old + condition_2 = B_eval_challenger == B_common # evaluated on all common blocks + + if condition_1 or condition_2: + break + + # Step 3.2: Check if challenger evaluated on ALL common blocks + if B_eval_challenger == B_common: + # Step 3.2.1: Create temporary list with all prompts (don't modify P_inc yet) + all_prompts = self.P_inc + [challenger] # Temporary combined list + all_inc_strings = inc_strings + [challenger_str] # Reuse inc_strings + + # Step 3.2.2: Perform NDS on all prompts + inc_vectors = np.array( + [ + self.runhistory.compute_current_vector(prompt_str, B_common) + for prompt_str in all_inc_strings + ] + ) + fronts = fast_non_dominated_sort(inc_vectors) + + # Step 3.2.3: Assign fronts correctly (indices now match all_prompts) + self.P_inc = [all_prompts[i] for i in fronts[0]] # First front becomes new P_inc + for front_idx in range(1, len(fronts)): + for i in fronts[front_idx]: + self._add_to_population(all_prompts[i]) # Dominated fronts to P_pop + + # Prune immediately after updating sets + self._prune_population() + + # Step 3.2.3: Stop the while loop + break + + if self.intensify_vs_all_incumbents: + any_inc_dominates = False + for inc in self.P_inc: + inc_vector = self.runhistory.compute_current_vector( + inc.construct_prompt(), B_eval_challenger + ) + if self._is_dominated(new_cost_vector, inc_vector): + any_inc_dominates = True + break + + if any_inc_dominates: + self._add_to_population(challenger) + self._prune_population() + break + else: + # Step 3.3: Get closest incumbent in normalized objective space + closest_incumbent = self._get_closest_incumbent(new_cost_vector) + assert closest_incumbent is not None, "There should always be incumbents" + + # Step 3.4: Get closest incumbent's cost vector on challenger's evaluated blocks + closest_inc_str = closest_incumbent.construct_prompt() + closest_inc_vector = self.runhistory.compute_current_vector( + closest_inc_str, B_eval_challenger + ) + + # Step 3.5: Check if inc domiantes the challenger (on same block subset) + inc_dominates = self._is_dominated(new_cost_vector, closest_inc_vector) + + # Step 3.6: Decision based on dominance + if inc_dominates: # Incumbent dominates or no dominance + # Add challenger to P_pop and stop + self._add_to_population(challenger) + self._prune_population() + break + + # Continue while loop if challenger is not dominated (update old_cost_vector at beginning of next iteration) + + # Step 5-6: Incumbent evaluation on additional block + # Step 5: Get incumbent with least evaluations + inc_strings = [inc.construct_prompt() for inc in self.P_inc] + least_evaluated_results = self.runhistory.get_least_evaluated_prompts(inc_strings) + + # Get union of all blocks that ANY incumbent has been evaluated on + all_inc_blocks = self.runhistory.get_evaluated_blocks(inc_strings) + assert isinstance(all_inc_blocks, list), f"Expected list, got {type(all_inc_blocks)}" + union_all_inc_blocks = set().union(*all_inc_blocks) + + # Randomly select from least evaluated incumbents + chosen_inc_str, chosen_inc_blocks = random.choice(least_evaluated_results) + + # Calculate gap: blocks that other incumbents have evaluated but this one hasn't + gap_blocks = union_all_inc_blocks - chosen_inc_blocks + + if gap_blocks: + # Case 1: Catch up - evaluate on a gap block (should be exactly one block) + new_block = random.choice(list(gap_blocks)) + + if self.verbosity > 1: + self.logger.info(f"📈 Catching up incumbent: evaluating on gap block {new_block}") + else: + # Case 2: All incumbents at same level - evaluate on a completely new block + all_available_blocks = set(range(len(self.task.blocks))) # All possible blocks + unevaluated_blocks = all_available_blocks - union_all_inc_blocks + + if not unevaluated_blocks: + # End case: All incumbents have been evaluated on all available blocks + if self.verbosity > 0: + self.logger.info( + "🏁 All incumbents evaluated on all blocks - no further incumbent evaluation needed" + ) + return + # Continue optimization - we can still find better solutions through evolution + + new_block = random.choice(list(unevaluated_blocks)) + + if self.verbosity > 0: + self.logger.info( + f"🆕 All incumbents at same level: evaluating on new block {new_block}" + ) + + # Evaluate the chosen incumbent on the selected block + _ = self.task.evaluate_on_block( + [chosen_inc_str], new_block, self.predictor, self.runhistory + ) + + def _add_to_population(self, prompt_to_add: Prompt): + """ + Adds a prompt to P_pop only if it's not already present. + """ + # Create a set of existing prompt strings in P_pop for efficient checking + population_strings = {p.construct_prompt() for p in self.P_pop} + prompt_str = prompt_to_add.construct_prompt() + if prompt_str not in population_strings: + self.P_pop.append(prompt_to_add) + + def _prune_population(self): + while len(self.P_inc) + len(self.P_pop) > self.population_size: + if self.P_pop: + if self.random_pruning: + worst_idx = random.randrange(len(self.P_pop)) + else: + # Compute prompts once per iteration to avoid redundancy + prompts = [p.construct_prompt() for p in self.P_pop] + + # Get eval counts and actual block sets for uniformity check + pop_eval_counts = [ + len(self.runhistory.get_evaluated_blocks(prompt)) for prompt in prompts + ] + pop_block_sets = [ + set( + self.runhistory.get_evaluated_blocks(prompt) + ) # Actual sets for comparison + for prompt in prompts + ] + + # Check if all have the SAME block set (not just count) + all_same_block_set = ( + all(bs == pop_block_sets[0] for bs in pop_block_sets) + if pop_block_sets + else False + ) + + if all_same_block_set: + # Case: Uniform block sets → Full NDS + CD on entire P_pop + all_pop_vectors = np.array( + [self.runhistory.compute_current_vector(prompt) for prompt in prompts] + ) + fronts = fast_non_dominated_sort(all_pop_vectors) + + if fronts: + # If only one front (all non-dominated), treat full set as "worst front" for CD pruning + if len(fronts) == 1: + worst_front = list(range(len(self.P_pop))) # All indices + else: + worst_front = fronts[-1] # Standard: most-dominated front + + # Now prune from worst_front via CD + worst_front_vectors = all_pop_vectors[worst_front] + distances = _calculate_crowding_distance(worst_front_vectors) + + min_distance = np.min(distances) + tied_indices = np.where(distances == min_distance)[0] + + if len(tied_indices) == 1: + local_worst = tied_indices[0] + else: + local_worst = random.choice(tied_indices) + + worst_idx = worst_front[local_worst] + else: + # Fail explicitly if no fronts at all (empty P_pop or NDS bug) + raise AssertionError( + f"No fronts found in P_pop during pruning (empty or NDS failure). " + f"P_pop size: {len(self.P_pop)}, vectors shape: {all_pop_vectors.shape if len(all_pop_vectors) > 0 else 'empty'}" + ) + + else: + # Case: Non-uniform block sets → Random from least evaluated (by count) + min_eval_count = min(pop_eval_counts) + least_evaluated_indices = [ + i for i, count in enumerate(pop_eval_counts) if count == min_eval_count + ] + worst_idx = random.choice(least_evaluated_indices) + + if self.verbosity > 1: + self.logger.info( + f"Pruning {self.P_pop[worst_idx].instruction_text[:30]}... from P_pop." + ) + + self.P_pop.pop(worst_idx) + + elif self.P_inc: + # Prune from P_inc based on worst crowding distance (unchanged) + prompts = [p.construct_prompt() for p in self.P_inc] + all_inc_blocks = self.runhistory.get_evaluated_blocks(prompts) + + # Ensure we have List[Set[int]] type + assert isinstance( + all_inc_blocks, list + ), f"Expected list, got {type(all_inc_blocks)}" + B_common = set.intersection(*all_inc_blocks) if all_inc_blocks else set() + + obj_vectors = np.array( + [self.runhistory.compute_current_vector(prompt, B_common) for prompt in prompts] + ) + distances = _calculate_crowding_distance(obj_vectors) + + min_distance = np.min(distances) + tied_indices = np.where(distances == min_distance)[0] + + if len(tied_indices) == 1: + worst_idx = tied_indices[0] + else: + worst_idx = random.choice(tied_indices) + + if self.verbosity > 0: + self.logger.info( + f"Pruning {self.P_inc[worst_idx].instruction_text[:30]}... from P_inc." + ) + + self.P_inc.pop(worst_idx) + else: + break + + def _intensify_challengers(self, challengers: List[Prompt]): + """ + Routes challengers to intensification based on the optimizer's strategy. + This version now uses the self.prompts_P_inc and self.prompts_P_pop sets + which are maintained by the main optimize loop. + """ + if self.freeze_p_pop: + # Original behavior: filter for only brand-new prompts + for challenger in challengers: + prompt_str = challenger.construct_prompt() + if not self.runhistory.get_evaluated_blocks(prompt_str): + self._do_intensification(challenger) + return + + # --- Default MO-SMAC-like Logic (freeze_population=False) --- + # Use the live P_inc to determine incumbency rather than the + # stale snapshot `self.prompts_P_inc` created at the start of the + # optimization step. Build a set of current incumbent strings once + # for efficient membership checks, and also allow direct object + # identity checks to short-circuit when the exact instance was + # already added to P_inc. + current_inc_strings = {p.construct_prompt() for p in self.P_inc} + + for challenger in challengers: + challenger_str = challenger.construct_prompt() + + # 1) If challenger is already an incumbent (live view), skip + if challenger_str in current_inc_strings: + if self.verbosity > 0: + self.logger.info( + f"Skipping challenger (already incumbent): {challenger_str[:30]}..." + ) + continue + + # 2) If challenger string already exists in P_pop, remove all old + # objects with the same constructed string and replace them with + # the challenger as the canonical representative. This prevents + # duplicates across P_inc and P_pop and ensures the challenger + # (freshly generated object) is the one that gets intensified. + if self.P_pop: + # Filter out any P_pop entries that match the challenger string. We remove + # stale duplicates so the challenger will not coexist with old objects + # that have the same constructed prompt. Do NOT append the challenger + # here — intensification (_do_intensification) will decide whether the + # challenger becomes an incumbent or a population member. This keeps the + # evaluation flow faithful to ParamILS/SMAC-like intensification. + new_pop = [p for p in self.P_pop if p.construct_prompt() != challenger_str] + if len(new_pop) != len(self.P_pop): + if self.verbosity > 1: + self.logger.info( + f"🔁 Removed {len(self.P_pop) - len(new_pop)} old P_pop object(s) matching challenger: {challenger_str[:30]}..." + ) + # Replace the population with the filtered list. Do NOT append challenger. + self.P_pop = new_pop + + # Intensify the challenger (now canonical in P_pop or new) + if self.verbosity > 0: + self.logger.info(f"Intensifying challenger: {challenger_str[:30]}...") + + self._do_intensification(challenger) + + def optimize(self, n_steps: int) -> List[str]: + """ + Main optimization loop that evolves the prompt population. + + Parameters: + n_steps (int): Number of optimization steps to perform. + + Returns: + List[str]: The final population of prompts after optimization. + """ + + self.prompts_P_inc = [p.construct_prompt() for p in self.P_inc] + self.prompts_P_pop = [p.construct_prompt() for p in self.P_pop] + + self._on_step_end() + self.runhistory.set_current_step() + + for step in range(n_steps): + if self.verbosity > 0: + self.logger.info(f"--- Starting Step {step + 1}/{n_steps} ---") + + offsprings = self._crossover() + mutated_challengers = self._mutate(offsprings) + + self._intensify_challengers(mutated_challengers) + + # Update attributes for callbacks + self.prompts_P_inc = [p.construct_prompt() for p in self.P_inc] + self.prompts_P_pop = [p.construct_prompt() for p in self.P_pop] + + continue_optimization = self._on_step_end() + self.runhistory.set_current_step() + if not continue_optimization: + break + + self._on_train_end() + + # Return final Pareto front with detailed information + final_pareto_front = [] + prompt_strings = [] + for p in self.P_inc: + prompt_str = p.construct_prompt() + prompt_strings.append(prompt_str) + obj_vector = self.runhistory.compute_current_vector(prompt_str) + total_in, total_out = self.runhistory.get_total_token_counts(prompt_str) + final_pareto_front.append( + { + "prompt": prompt_str, + "objectives": obj_vector.tolist(), + "total_input_tokens": total_in, + "total_output_tokens": total_out, + } + ) + + # Return just the prompt strings for compatibility with base class + return prompt_strings + + def __getstate__(self): + """Return state values to be pickled.""" + state = self.__dict__.copy() + state.pop("predictor", None) + state.pop("logger", None) + state.pop("meta_llm", None) + state.pop("downstream_llm", None) + + return state + + def __setstate__(self, state): + """Restore state from the unpickled state values.""" + self.__dict__.update(state) + self.predictor = None + self.logger = getLogger(__name__) + + def _on_step_end(self) -> bool: + """ + Override base _on_step_end to only call MO-CAPO compatible callbacks. + The base implementation expects 'scores' attribute which doesn't exist in MO-CAPO. + """ + continue_optimization = True + for callback in self.callbacks: + # Only call callbacks that are MO-CAPO aware (have on_step_end method) + if hasattr(callback, "on_step_end"): + continue_optimization &= callback.on_step_end(self) + return continue_optimization \ No newline at end of file From 732d5c038b120623564b0a77b99a058f7adb36f0 Mon Sep 17 00:00:00 2001 From: Tom Zehle Date: Sat, 27 Dec 2025 14:40:38 +0100 Subject: [PATCH 03/53] Simplify CAPO helper usage and adjust tests --- docs/examples/getting_started.md | 4 +- docs/examples/reward_task_tutorial.md | 2 +- promptolution/helpers.py | 13 +- promptolution/optimizers/__init__.py | 2 + promptolution/optimizers/base_optimizer.py | 2 +- promptolution/optimizers/capo.py | 156 +++--------- promptolution/optimizers/capo_utils.py | 111 +++++++++ promptolution/optimizers/capoeira.py | 261 +++++++++++++++++++++ promptolution/tasks/base_task.py | 54 ++++- promptolution/utils/capo_utils.py | 112 +++++++++ tests/optimizers/test_capo.py | 115 ++++++--- tests/optimizers/test_capoeira.py | 107 +++++++++ tutorials/getting_started.ipynb | 4 +- tutorials/reward_task_tutorial.ipynb | 4 +- 14 files changed, 785 insertions(+), 162 deletions(-) create mode 100644 promptolution/optimizers/capo_utils.py create mode 100644 promptolution/utils/capo_utils.py create mode 100644 tests/optimizers/test_capoeira.py diff --git a/docs/examples/getting_started.md b/docs/examples/getting_started.md index 81e1f57a..2dfc1f14 100644 --- a/docs/examples/getting_started.md +++ b/docs/examples/getting_started.md @@ -83,7 +83,7 @@ api_key = "YOUR_API_KEY" # Replace with your Promptolution API key ``` Here's an explanation of each configuration parameter in the ExperimentConfig: -- `optimizer`: The algorithm used for prompt optimization. Currently we support "capo", "evopromptga", "evopromptde", and "opro". For this example, we use "capo" as it is capable of leveraging few-shot examples. +- `optimizer`: The algorithm used for prompt optimization. Currently we support "capo", "capoeira", "evopromptga", "evopromptde", and "opro". For this example, we use "capo" as it is capable of leveraging few-shot examples. - `task_description`: A string describing the task you're optimizing prompts for. This is used to provide the meta-llm with context about your task. - `prompts`: A list of initial prompt strings that will be used as the starting point for optimization. - `n_steps`: The number of optimization steps to run. Higher values allow more exploration and refinement but require more API calls and computational resources. @@ -114,7 +114,7 @@ With everything configured, you're ready to optimize your prompts! The `run_expe prompts = run_experiment(df, config) ``` - 📌 CAPO requires block evaluation strategy. Setting it to 'sequential_block'. + 📌 CAPO-style optimizers require block evaluation strategy. Setting it to 'sequential_block'. ⚠️ The LLM does not have a tokenizer. Using simple token count. 🔥 Starting optimization... 📊 Starting evaluation... diff --git a/docs/examples/reward_task_tutorial.md b/docs/examples/reward_task_tutorial.md index 82d0e973..da51cfdd 100644 --- a/docs/examples/reward_task_tutorial.md +++ b/docs/examples/reward_task_tutorial.md @@ -102,7 +102,7 @@ api_key = "YOUR_API_KEY" # Replace with your Promptolution API key ``` Here's an explanation of each configuration parameter in the ExperimentConfig: -- `optimizer`: The algorithm used for prompt optimization. Currently we support "capo", "evopromptga", "evopromptde", and "opro". For this example, we use "capo" as it is capable of leveraging few-shot examples. +- `optimizer`: The algorithm used for prompt optimization. Currently we support "capo", "capoeira", "evopromptga", "evopromptde", and "opro". For this example, we use "capo" as it is capable of leveraging few-shot examples. - `task_description`: A string describing the task you're optimizing prompts for. This is used to provide the meta-llm with context about your task. - `prompts`: A list of initial prompt strings that will be used as the starting point for optimization. - `n_steps`: The number of optimization steps to run. Higher values allow more exploration and refinement but require more API calls and computational resources. diff --git a/promptolution/helpers.py b/promptolution/helpers.py index a25c0080..bb0624e5 100644 --- a/promptolution/helpers.py +++ b/promptolution/helpers.py @@ -26,6 +26,7 @@ from promptolution.llms.local_llm import LocalLLM from promptolution.llms.vllm import VLLM from promptolution.optimizers.capo import CAPO +from promptolution.optimizers.capoeira import Capoeira from promptolution.optimizers.evoprompt_de import EvoPromptDE from promptolution.optimizers.evoprompt_ga import EvoPromptGA from promptolution.optimizers.opro import OPRO @@ -79,8 +80,8 @@ def run_optimization(df: pd.DataFrame, config: "ExperimentConfig") -> List[Promp ) config.prompts = [Prompt(p) for p in initial_prompts] - if config.optimizer == "capo" and (config.eval_strategy is None or "block" not in config.eval_strategy): - logger.warning("📌 CAPO requires block evaluation strategy. Setting it to 'sequential_block'.") + if config.optimizer in {"capo", "capoeira"} and (config.eval_strategy is None or "block" not in config.eval_strategy): + logger.warning("📌 CAPO-style optimizers require block evaluation strategy. Setting it to 'sequential_block'.") config.eval_strategy = "sequential_block" task = get_task(df, config, judge_llm=llm) @@ -238,6 +239,14 @@ def get_optimizer( config=config, ) + if final_optimizer == "capoeira": + return Capoeira( + predictor=predictor, + meta_llm=meta_llm, + task=task, + config=config, + ) + if final_optimizer == "evopromptde": return EvoPromptDE(predictor=predictor, meta_llm=meta_llm, task=task, config=config) diff --git a/promptolution/optimizers/__init__.py b/promptolution/optimizers/__init__.py index 4b7a7dbb..96e9a484 100644 --- a/promptolution/optimizers/__init__.py +++ b/promptolution/optimizers/__init__.py @@ -1,12 +1,14 @@ """Module for prompt optimizers.""" from promptolution.optimizers.capo import CAPO +from promptolution.optimizers.capoeira import Capoeira from promptolution.optimizers.evoprompt_de import EvoPromptDE from promptolution.optimizers.evoprompt_ga import EvoPromptGA from promptolution.optimizers.opro import OPRO __all__ = [ "CAPO", + "Capoeira", "EvoPromptDE", "EvoPromptGA", "OPRO", diff --git a/promptolution/optimizers/base_optimizer.py b/promptolution/optimizers/base_optimizer.py index 7264f6fb..68717b28 100644 --- a/promptolution/optimizers/base_optimizer.py +++ b/promptolution/optimizers/base_optimizer.py @@ -15,7 +15,7 @@ logger = get_logger(__name__) -OptimizerType = Literal["evopromptde", "evopromptga", "opro", "capo"] +OptimizerType = Literal["evopromptde", "evopromptga", "opro", "capo", "capoeira"] class BaseOptimizer(ABC): diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index 3c5955a6..41174112 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -17,6 +17,7 @@ from promptolution.optimizers.base_optimizer import BaseOptimizer from promptolution.utils.formatting import extract_from_tag +from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation from promptolution.utils.logging import get_logger from promptolution.utils.prompt import Prompt, sort_prompts_by_scores from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE @@ -121,128 +122,27 @@ def __init__( self.target_end_marker = "" def _initialize_population(self, initial_prompts: List[Prompt]) -> List[Prompt]: - """Initializes the population of Prompt objects from initial instructions. - - Args: - initial_prompts (List[str]): List of initial prompt instructions. - - Returns: - List[Prompt]: Initialized population of prompts with few-shot examples. - """ + """Initializes the population of Prompt objects from initial instructions.""" population = [] for prompt in initial_prompts: num_examples = random.randint(0, self.upper_shots) - few_shots = self._create_few_shot_examples(prompt.instruction, num_examples) + few_shots = build_few_shot_examples( + instruction=prompt.instruction, + num_examples=num_examples, + df_few_shots=self.df_few_shots, + x_column=self.task.x_column, + y_column=self.task.y_column, + predictor=self.predictor, + fewshot_template=CAPO_FEWSHOT_TEMPLATE, + target_begin_marker=self.target_begin_marker, + target_end_marker=self.target_end_marker, + check_fs_accuracy=self.check_fs_accuracy, + create_fs_reasoning=self.create_fs_reasoning, + ) population.append(Prompt(prompt.instruction, few_shots)) return population - def _create_few_shot_examples(self, instruction: str, num_examples: int) -> List[str]: - if num_examples == 0: - return [] - - few_shot_samples = self.df_few_shots.sample(num_examples, replace=False) - sample_inputs = few_shot_samples[self.task.x_column].values.astype(str) - sample_targets = few_shot_samples[self.task.y_column].values - few_shots = [ - CAPO_FEWSHOT_TEMPLATE.replace("", i).replace( - "", f"{self.target_begin_marker}{t}{self.target_end_marker}" - ) - for i, t in zip(sample_inputs, sample_targets) - ] - - if not self.create_fs_reasoning: - # If we do not create reasoning, return the few-shot examples directly - return few_shots - - preds, seqs = self.predictor.predict( - [instruction] * num_examples, - list(sample_inputs), - return_seq=True, - ) - if isinstance(seqs, str): - seqs = [seqs] - if isinstance(preds, str): - preds = [preds] - - # Check which predictions are correct and get a single one per example - for j in range(num_examples): - # Process and clean up the generated sequences - seqs[j] = seqs[j].replace(sample_inputs[j], "", 1).strip() - # Check if the prediction is correct and add reasoning if so - if preds[j] == sample_targets[j] or not self.check_fs_accuracy: - few_shots[j] = CAPO_FEWSHOT_TEMPLATE.replace("", sample_inputs[j]).replace("", seqs[j]) - - return few_shots - - def _crossover(self, parents: List[Prompt]) -> List[Prompt]: - """Performs crossover among parent prompts to generate offsprings. - - Args: - parents (List[Prompt]): List of parent prompts. - - Returns: - List[Prompt]: List of new offsprings after crossover. - """ - crossover_prompts = [] - offspring_few_shots = [] - for _ in range(self.crossovers_per_iter): - mother, father = random.sample(parents, 2) - crossover_prompt = ( - self.crossover_template.replace("", mother.instruction) - .replace("", father.instruction) - .strip() - ) - # collect all crossover prompts then pass them bundled to the meta llm (speedup) - crossover_prompts.append(crossover_prompt) - combined_few_shots = mother.few_shots + father.few_shots - num_few_shots = (len(mother.few_shots) + len(father.few_shots)) // 2 - offspring_few_shot = random.sample(combined_few_shots, num_few_shots) if combined_few_shots else [] - offspring_few_shots.append(offspring_few_shot) - - child_instructions = self.meta_llm.get_response(crossover_prompts) - - offsprings = [] - for instruction, examples in zip(child_instructions, offspring_few_shots): - instruction = extract_from_tag(instruction, "", "") - offsprings.append(Prompt(instruction, examples)) - - return offsprings - - def _mutate(self, offsprings: List[Prompt]) -> List[Prompt]: - """Apply mutation to offsprings to generate new candidate prompts. - - Args: - offsprings (List[Prompt]): List of offsprings to mutate. - - Returns: - List[Prompt]: List of mutated prompts. - """ - # collect all mutation prompts then pass them bundled to the meta llm (speedup) - mutation_prompts = [ - self.mutation_template.replace("", prompt.instruction) for prompt in offsprings - ] - new_instructions = self.meta_llm.get_response(mutation_prompts) - - mutated = [] - for new_instruction, prompt in zip(new_instructions, offsprings): - new_instruction = extract_from_tag(new_instruction, "", "") - p = random.random() - - new_few_shots: List[str] - if p < 1 / 3 and len(prompt.few_shots) < self.upper_shots: # add a random few shot - new_few_shot = self._create_few_shot_examples(new_instruction, 1) - new_few_shots = prompt.few_shots + new_few_shot - elif 1 / 3 <= p < 2 / 3 and len(prompt.few_shots) > 0: # remove a random few shot - new_few_shots = random.sample(prompt.few_shots, len(prompt.few_shots) - 1) - else: # do not change few shots, but shuffle - new_few_shots = prompt.few_shots - - random.shuffle(new_few_shots) - mutated.append(Prompt(new_instruction, new_few_shots)) - - return mutated - def _do_racing(self, candidates: List[Prompt], k: int) -> Tuple[List[Prompt], List[float]]: """Perform the racing (selection) phase by comparing candidates based on their evaluation scores using the provided test statistic. @@ -297,13 +197,25 @@ def _pre_optimization_loop(self) -> None: self.task.reset_block_idx() def _step(self) -> List[Prompt]: - """Perform a single optimization step. - - Returns: - List[Prompt]: The optimized list of prompts after the step. - """ - offsprings = self._crossover(self.prompts) - mutated = self._mutate(offsprings) + """Perform a single optimization step.""" + offsprings = perform_crossover(self.prompts, self.crossovers_per_iter, self.crossover_template, self.meta_llm) + mutated = perform_mutation( + offsprings=offsprings, + mutation_template=self.mutation_template, + upper_shots=self.upper_shots, + meta_llm=self.meta_llm, + few_shot_kwargs=dict( + df_few_shots=self.df_few_shots, + x_column=self.task.x_column, + y_column=self.task.y_column, + predictor=self.predictor, + fewshot_template=CAPO_FEWSHOT_TEMPLATE, + target_begin_marker=self.target_begin_marker, + target_end_marker=self.target_end_marker, + check_fs_accuracy=self.check_fs_accuracy, + create_fs_reasoning=self.create_fs_reasoning, + ), + ) combined = self.prompts + mutated self.prompts, self.scores = self._do_racing(combined, self.population_size) diff --git a/promptolution/optimizers/capo_utils.py b/promptolution/optimizers/capo_utils.py new file mode 100644 index 00000000..45bc2ff9 --- /dev/null +++ b/promptolution/optimizers/capo_utils.py @@ -0,0 +1,111 @@ +"""Shared utilities for CAPO-style optimizers.""" + +from __future__ import annotations + +import random +from typing import Callable, List, Optional + +import pandas as pd + +from promptolution.utils.formatting import extract_from_tag +from promptolution.utils.prompt import Prompt + + +def build_few_shot_examples( + instruction: str, + num_examples: int, + df_few_shots: pd.DataFrame, + task, + predictor, + fewshot_template: str, + target_begin_marker: str, + target_end_marker: str, + check_fs_accuracy: bool, + create_fs_reasoning: bool, +) -> List[str]: + """Create few-shot examples with optional reasoning replacement.""" + if num_examples == 0: + return [] + + few_shot_samples = df_few_shots.sample(num_examples, replace=False) + sample_inputs = few_shot_samples[task.x_column].values.astype(str) + sample_targets = few_shot_samples[task.y_column].values + few_shots = [ + fewshot_template.replace("", i).replace("", f"{target_begin_marker}{t}{target_end_marker}") + for i, t in zip(sample_inputs, sample_targets) + ] + + if not create_fs_reasoning: + return few_shots + + preds, seqs = predictor.predict( + [instruction] * num_examples, + list(sample_inputs), + return_seq=True, + ) + if isinstance(seqs, str): + seqs = [seqs] + if isinstance(preds, str): + preds = [preds] + + for j in range(num_examples): + seqs[j] = seqs[j].replace(sample_inputs[j], "", 1).strip() + if preds[j] == sample_targets[j] or not check_fs_accuracy: + few_shots[j] = fewshot_template.replace("", sample_inputs[j]).replace("", seqs[j]) + + return few_shots + + +def perform_crossover( + parents: List[Prompt], + crossovers_per_iter: int, + template: str, + meta_llm, +) -> List[Prompt]: + """Generate crossover offspring prompts.""" + crossover_prompts: List[str] = [] + offspring_few_shots: List[List[str]] = [] + for _ in range(crossovers_per_iter): + mother, father = (parents if len(parents) == 2 else random.sample(parents, 2)) + crossover_prompt = template.replace("", mother.instruction).replace("", father.instruction).strip() + crossover_prompts.append(crossover_prompt) + combined_few_shots = mother.few_shots + father.few_shots + num_few_shots = (len(mother.few_shots) + len(father.few_shots)) // 2 + offspring_few_shot = random.sample(combined_few_shots, num_few_shots) if combined_few_shots else [] + offspring_few_shots.append(offspring_few_shot) + + child_instructions = meta_llm.get_response(crossover_prompts) + return [ + Prompt(extract_from_tag(instr, "", ""), examples) + for instr, examples in zip(child_instructions, offspring_few_shots) + ] + + +def perform_mutation( + offsprings: List[Prompt], + mutation_template: str, + create_few_shots: Callable[[str, int], List[str]], + upper_shots: int, + meta_llm, +) -> List[Prompt]: + """Mutate offspring prompts.""" + mutation_prompts = [mutation_template.replace("", prompt.instruction) for prompt in offsprings] + new_instructions = meta_llm.get_response(mutation_prompts) + + mutated: List[Prompt] = [] + for new_instruction, prompt in zip(new_instructions, offsprings): + new_instruction = extract_from_tag(new_instruction, "", "") + p = random.random() + + if p < 1 / 3 and len(prompt.few_shots) < upper_shots: + new_few_shot = create_few_shots(new_instruction, 1) + new_few_shots = prompt.few_shots + new_few_shot + elif 1 / 3 <= p < 2 / 3 and len(prompt.few_shots) > 0: + new_few_shots = random.sample(prompt.few_shots, len(prompt.few_shots) - 1) + else: + new_few_shots = prompt.few_shots + + random.shuffle(new_few_shots) + mutated.append(Prompt(new_instruction, new_few_shots)) + + return mutated diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index e69de29b..d35887a7 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -0,0 +1,261 @@ +"""Implementation of the Capoeira (Multi-Objective CAPO) optimizer.""" + +from __future__ import annotations + +import random +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import numpy as np +import pandas as pd + +if TYPE_CHECKING: # pragma: no cover + from promptolution.utils.callbacks import BaseCallback + from promptolution.llms.base_llm import BaseLLM + from promptolution.predictors.base_predictor import BasePredictor + from promptolution.tasks.base_task import BaseTask + from promptolution.utils.config import ExperimentConfig + +from promptolution.optimizers.base_optimizer import BaseOptimizer +from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation +from promptolution.utils.formatting import extract_from_tag +from promptolution.utils.logging import get_logger +from promptolution.utils.prompt import Prompt +from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE +from promptolution.utils.token_counter import get_token_counter + +logger = get_logger(__name__) + + +class Capoeira(BaseOptimizer): + """Multi-objective variant of CAPO with Pareto-based selection.""" + + def __init__( + self, + predictor: "BasePredictor", + task: "BaseTask", + meta_llm: "BaseLLM", + initial_prompts: Optional[List[str]] = None, + crossover_template: Optional[str] = None, + mutation_template: Optional[str] = None, + crossovers_per_iter: int = 4, + upper_shots: int = 5, + population_size: Optional[int] = None, + cost_per_input_token: float = 1.0, + cost_per_output_token: float = 0.0, + check_fs_accuracy: bool = True, + create_fs_reasoning: bool = True, + df_few_shots: Optional[pd.DataFrame] = None, + callbacks: Optional[List["BaseCallback"]] = None, + config: Optional["ExperimentConfig"] = None, + ) -> None: + """Initialize the Capoeira optimizer. + + Args: + predictor: The predictor used to evaluate prompt performance. + task: The task instance containing data and evaluation settings. + meta_llm: Meta language model for crossover and mutation generation. + initial_prompts: Starting prompt strings to seed the population. + crossover_template: Optional meta-prompt template for crossover. + mutation_template: Optional meta-prompt template for mutation. + crossovers_per_iter: Number of crossover operations per iteration. + upper_shots: Maximum number of few-shot examples to attach. + population_size: Target population size used when pruning fronts. + cost_per_input_token: Weight applied to input token cost for the cost objective. + cost_per_output_token: Weight applied to output token cost for the cost objective. + check_fs_accuracy: Whether to verify few-shot correctness before use. + create_fs_reasoning: Whether to replace few-shots with model reasoning. + df_few_shots: Optional dataframe providing few-shot examples. If None, will pop 10% of datapoints from task. + callbacks: Optional list of optimization callbacks. + config: Optional experiment configuration object. + """ + self.meta_llm = meta_llm + self.downstream_llm = predictor.llm + self.crossovers_per_iter = crossovers_per_iter + self.upper_shots = upper_shots + self.population_size = population_size + self.cost_per_input_token = cost_per_input_token + self.cost_per_output_token = cost_per_output_token + self.check_fs_accuracy = check_fs_accuracy + self.create_fs_reasoning = create_fs_reasoning + + super().__init__(predictor, task, initial_prompts, callbacks, config) + + self.crossover_template = self._initialize_meta_template(crossover_template or CAPO_CROSSOVER_TEMPLATE) + self.mutation_template = self._initialize_meta_template(mutation_template or CAPO_MUTATION_TEMPLATE) + self.token_counter = get_token_counter(self.downstream_llm) + self.df_few_shots = df_few_shots if df_few_shots is not None else task.pop_datapoints(frac=0.1) + self.population_size = self.population_size or len(self.prompts) or 1 + + if hasattr(self.predictor, "begin_marker") and hasattr(self.predictor, "end_marker"): + self.target_begin_marker = self.predictor.begin_marker + self.target_end_marker = self.predictor.end_marker + else: + self.target_begin_marker = "" + self.target_end_marker = "" + + def _pre_optimization_loop(self) -> None: + population: List[Prompt] = [] + for prompt in self.prompts: + num_examples = random.randint(0, self.upper_shots) + few_shots = build_few_shot_examples( + instruction=prompt.instruction, + num_examples=num_examples, + df_few_shots=self.df_few_shots, + x_column=self.task.x_column, + y_column=self.task.y_column, + predictor=self.predictor, + fewshot_template=CAPO_FEWSHOT_TEMPLATE, + target_begin_marker=self.target_begin_marker, + target_end_marker=self.target_end_marker, + check_fs_accuracy=self.check_fs_accuracy, + create_fs_reasoning=self.create_fs_reasoning, + ) + population.append(Prompt(prompt.instruction, few_shots)) + + self.prompts = population + self.max_prompt_length = max(self.token_counter(p.construct_prompt()) for p in self.prompts) if self.prompts else 1 + initial_vectors = self._evaluate_candidates(self.prompts) + self.prompts, selected_vectors = self._select_population(self.prompts, initial_vectors) + self.scores = (-selected_vectors[:, 0]).tolist() + + def _evaluate_candidates(self, candidates: List[Prompt]) -> np.ndarray: + evaluation = self.task.evaluate( + candidates, + self.predictor, + eval_strategy=self.task.eval_strategy, + return_costs=True, + return_seq=False, + return_agg_scores=True, + ) + + if isinstance(evaluation, tuple) and len(evaluation) == 3: + scores, input_tokens, output_tokens = evaluation + else: + scores = evaluation # type: ignore[assignment] + input_tokens = [self.token_counter(c.construct_prompt()) for c in candidates] + output_tokens = [0.0 for _ in candidates] + + input_tokens_arr = np.array(input_tokens, dtype=float) + output_tokens_arr = np.array(output_tokens, dtype=float) + + if not input_tokens_arr.any() and not output_tokens_arr.any(): + input_tokens_arr = np.array([self.token_counter(c.construct_prompt()) for c in candidates], dtype=float) + output_tokens_arr = np.zeros_like(input_tokens_arr, dtype=float) + + score_vectors = np.column_stack( + [ + -np.array(scores, dtype=float), + self.cost_per_input_token * input_tokens_arr + self.cost_per_output_token * output_tokens_arr, + ] + ) + return score_vectors + + def _select_population(self, candidates: List[Prompt], score_vectors: np.ndarray) -> Tuple[List[Prompt], np.ndarray]: + selected_indices: List[int] = [] + fronts = self.fast_non_dominated_sort(score_vectors) + for front in fronts: + if len(selected_indices) + len(front) <= self.population_size: + selected_indices.extend(front) + else: + remaining = self.population_size - len(selected_indices) + front_vectors = score_vectors[front] + distances = self.calculate_crowding_distance(front_vectors) + sorted_front = [i for _, i in sorted(zip(distances, front), reverse=True)] + selected_indices.extend(sorted_front[:remaining]) + break + + selected_prompts = [candidates[i] for i in selected_indices] + selected_vectors = score_vectors[selected_indices] + return selected_prompts, selected_vectors + + def _step(self) -> List[Prompt]: + offsprings = perform_crossover(self.prompts, self.crossovers_per_iter, self.crossover_template, self.meta_llm) + mutated = perform_mutation( + offsprings=offsprings, + mutation_template=self.mutation_template, + upper_shots=self.upper_shots, + meta_llm=self.meta_llm, + few_shot_kwargs=dict( + df_few_shots=self.df_few_shots, + x_column=self.task.x_column, + y_column=self.task.y_column, + predictor=self.predictor, + fewshot_template=CAPO_FEWSHOT_TEMPLATE, + target_begin_marker=self.target_begin_marker, + target_end_marker=self.target_end_marker, + check_fs_accuracy=self.check_fs_accuracy, + create_fs_reasoning=self.create_fs_reasoning, + ), + ) + combined = self.prompts + mutated + + score_vectors = self._evaluate_candidates(combined) + self.prompts, selected_vectors = self._select_population(combined, score_vectors) + self.scores = (-selected_vectors[:, 0]).tolist() + return self.prompts + + def get_pareto_front(self) -> List[Dict[str, Any]]: + """Return the current Pareto front with objective values.""" + score_vectors = self._evaluate_candidates(self.prompts) + return [ + { + "prompt": prompt.construct_prompt(), + "score": float(score_vectors[i][0] * -1), + "cost": float(score_vectors[i][1]), + } + for i, prompt in enumerate(self.prompts) + ] + + @staticmethod + def fast_non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: + """Perform fast non-dominated sorting (NSGA-II) in a vectorized manner.""" + num_solutions = obj_vectors.shape[0] + if num_solutions == 0: + return [] + + less = obj_vectors[:, None, :] < obj_vectors[None, :, :] + less_equal = obj_vectors[:, None, :] <= obj_vectors[None, :, :] + dominates = np.all(less_equal, axis=2) & np.any(less, axis=2) + + domination_counts = dominates.sum(axis=0) + dominated_solutions = [list(np.where(dominates[i])[0]) for i in range(num_solutions)] + + fronts: List[List[int]] = [list(np.where(domination_counts == 0)[0])] + current_front = 0 + + while current_front < len(fronts) and fronts[current_front]: + next_front: List[int] = [] + for i in fronts[current_front]: + for dominated in dominated_solutions[i]: + domination_counts[dominated] -= 1 + if domination_counts[dominated] == 0: + next_front.append(dominated) + if next_front: + fronts.append(next_front) + current_front += 1 + + return fronts + + @staticmethod + def calculate_crowding_distance(obj_vectors: np.ndarray) -> np.ndarray: + """Calculate crowding distance for a set of solutions.""" + num_solutions, num_obj = obj_vectors.shape + if num_solutions <= 2: + return np.full(num_solutions, float("inf")) + + distances = np.zeros(num_solutions) + for i in range(num_obj): + sorted_indices = np.argsort(obj_vectors[:, i]) + distances[sorted_indices[0]] = float("inf") + distances[sorted_indices[-1]] = float("inf") + + f_min = obj_vectors[sorted_indices[0], i] + f_max = obj_vectors[sorted_indices[-1], i] + if f_max == f_min: + continue + + slice_indices = sorted_indices[1:-1] + next_vals = obj_vectors[sorted_indices[2:], i] + prev_vals = obj_vectors[sorted_indices[:-2], i] + distances[slice_indices] += (next_vals - prev_vals) / (f_max - f_min) + return distances diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 2f1c164b..6d897162 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -9,6 +9,12 @@ from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union, overload from promptolution.utils.prompt import Prompt +from promptolution.utils.token_counter import get_token_counter +from promptolution.utils.logging import get_logger + +logger = get_logger(__name__) + +logger = get_logger(__name__) if TYPE_CHECKING: # pragma: no cover from promptolution.predictors.base_predictor import BasePredictor @@ -247,7 +253,13 @@ def evaluate( return_agg_scores: bool = True, return_seq: bool = False, eval_strategy: Optional["EvalStrategy"] = None, - ) -> Union[List[float], List[List[float]], Tuple[List[List[float]], List[List[str]]]]: + return_costs: bool = False, + ) -> Union[ + List[float], + List[List[float]], + Tuple[List[List[float]], List[List[str]]], + Tuple[List[float], List[float], List[float]], + ]: """Evaluate a set of prompts using a given predictor. This method orchestrates subsampling, prediction, caching, and result collection. @@ -255,8 +267,10 @@ def evaluate( Note: Cannot return both aggregated scores and sequences (assertion will fail). """ assert not (return_agg_scores and return_seq), "Cannot return both aggregated scores and sequences" + assert not return_seq or not return_costs, "Token cost reporting is not supported together with sequences." seqs: List[str] = [] + token_counter = get_token_counter(predictor.llm) if return_costs else None prompts = [prompts] if isinstance(prompts, Prompt) else prompts eval_strategy = eval_strategy or self.eval_strategy @@ -285,7 +299,7 @@ def evaluate( if return_seq: self.seq_cache[cache_key] = seqs[i] - return self._collect_results_from_cache( + agg_scores = self._collect_results_from_cache( prompts, xs, ys, @@ -293,6 +307,42 @@ def evaluate( return_seq, ) + if not return_costs: + return agg_scores + + per_prompt_inputs: List[float] = [] + per_prompt_outputs: List[float] = [] + + if token_counter is None: + logger.warning("⚠️ Token counting unavailable; returning zero costs.") + per_prompt_inputs = [0.0 for _ in prompts] + per_prompt_outputs = [0.0 for _ in prompts] + return agg_scores, per_prompt_inputs, per_prompt_outputs + + preds_by_prompt: List[List[str]] = [] + if isinstance(preds, list): + if preds and isinstance(preds[0], list): + preds_by_prompt = preds # type: ignore[assignment] + elif preds and isinstance(preds[0], str): + preds_by_prompt = [preds for _ in prompts] + + xs_token_mean = float(np.mean([token_counter(x) for x in xs])) if xs else 0.0 + + for idx, prompt in enumerate(prompts): + prompt_tokens = token_counter(prompt.construct_prompt()) + input_tokens = prompt_tokens + xs_token_mean + + if preds_by_prompt and idx < len(preds_by_prompt) and preds_by_prompt[idx]: + avg_output = float(np.mean([token_counter(p) for p in preds_by_prompt[idx]])) + else: + avg_output = 0.0 + logger.warning("⚠️ Unable to estimate output tokens; defaulting to 0.") + + per_prompt_inputs.append(input_tokens) + per_prompt_outputs.append(avg_output) + + return agg_scores, per_prompt_inputs, per_prompt_outputs + def pop_datapoints(self, n: Optional[int] = None, frac: Optional[float] = None) -> pd.DataFrame: """Pop a number of datapoints from the dataset. diff --git a/promptolution/utils/capo_utils.py b/promptolution/utils/capo_utils.py new file mode 100644 index 00000000..d451476d --- /dev/null +++ b/promptolution/utils/capo_utils.py @@ -0,0 +1,112 @@ +"""Shared utilities for CAPO-style optimizers.""" + +from __future__ import annotations + +import random +from typing import Callable, List + +import pandas as pd + +from promptolution.utils.formatting import extract_from_tag +from promptolution.utils.prompt import Prompt + + +def build_few_shot_examples( + instruction: str, + num_examples: int, + df_few_shots: pd.DataFrame, + x_column: str, + y_column: str, + predictor, + fewshot_template: str, + target_begin_marker: str, + target_end_marker: str, + check_fs_accuracy: bool, + create_fs_reasoning: bool, +) -> List[str]: + """Create few-shot examples with optional reasoning replacement.""" + if num_examples == 0: + return [] + + few_shot_samples = df_few_shots.sample(num_examples, replace=False) + sample_inputs = few_shot_samples[x_column].values.astype(str) + sample_targets = few_shot_samples[y_column].values + few_shots = [ + fewshot_template.replace("", i).replace("", f"{target_begin_marker}{t}{target_end_marker}") + for i, t in zip(sample_inputs, sample_targets) + ] + + if not create_fs_reasoning: + return few_shots + + preds, seqs = predictor.predict( + [instruction] * num_examples, + list(sample_inputs), + return_seq=True, + ) + if isinstance(seqs, str): + seqs = [seqs] + if isinstance(preds, str): + preds = [preds] + + for j in range(num_examples): + seqs[j] = seqs[j].replace(sample_inputs[j], "", 1).strip() + if preds[j] == sample_targets[j] or not check_fs_accuracy: + few_shots[j] = fewshot_template.replace("", sample_inputs[j]).replace("", seqs[j]) + + return few_shots + + +def perform_crossover( + parents: List[Prompt], + crossovers_per_iter: int, + template: str, + meta_llm, +) -> List[Prompt]: + """Generate crossover offspring prompts.""" + crossover_prompts: List[str] = [] + offspring_few_shots: List[List[str]] = [] + for _ in range(crossovers_per_iter): + mother, father = (parents if len(parents) == 2 else random.sample(parents, 2)) + crossover_prompt = template.replace("", mother.instruction).replace("", father.instruction).strip() + crossover_prompts.append(crossover_prompt) + combined_few_shots = mother.few_shots + father.few_shots + num_few_shots = (len(mother.few_shots) + len(father.few_shots)) // 2 + offspring_few_shot = random.sample(combined_few_shots, num_few_shots) if combined_few_shots else [] + offspring_few_shots.append(offspring_few_shot) + + child_instructions = meta_llm.get_response(crossover_prompts) + return [ + Prompt(extract_from_tag(instr, "", ""), examples) + for instr, examples in zip(child_instructions, offspring_few_shots) + ] + + +def perform_mutation( + offsprings: List[Prompt], + mutation_template: str, + upper_shots: int, + meta_llm, + few_shot_kwargs: dict, +) -> List[Prompt]: + """Mutate offspring prompts.""" + mutation_prompts = [mutation_template.replace("", prompt.instruction) for prompt in offsprings] + new_instructions = meta_llm.get_response(mutation_prompts) + + mutated: List[Prompt] = [] + for new_instruction, prompt in zip(new_instructions, offsprings): + new_instruction = extract_from_tag(new_instruction, "", "") + p = random.random() + + if p < 1 / 3 and len(prompt.few_shots) < upper_shots: + new_few_shot = build_few_shot_examples(new_instruction, 1, **few_shot_kwargs) + new_few_shots = prompt.few_shots + new_few_shot + elif 1 / 3 <= p < 2 / 3 and len(prompt.few_shots) > 0: + new_few_shots = random.sample(prompt.few_shots, len(prompt.few_shots) - 1) + else: + new_few_shots = prompt.few_shots + + random.shuffle(new_few_shots) + mutated.append(Prompt(new_instruction, new_few_shots)) + + return mutated diff --git a/tests/optimizers/test_capo.py b/tests/optimizers/test_capo.py index 305f290a..5fff4cb0 100644 --- a/tests/optimizers/test_capo.py +++ b/tests/optimizers/test_capo.py @@ -5,8 +5,9 @@ from tests.mocks.mock_task import MockTask from promptolution.optimizers.capo import CAPO +from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation from promptolution.utils.prompt import Prompt -from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_MUTATION_TEMPLATE +from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE def test_capo_initialization(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): @@ -28,7 +29,7 @@ def test_capo_initialization(mock_meta_llm, mock_predictor, initial_prompts, moc def test_capo_initialize_population(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): - """Test the _initialize_population method.""" + """Test initializing the population using pre-optimization loop.""" optimizer = CAPO( predictor=mock_predictor, task=mock_task, @@ -37,15 +38,9 @@ def test_capo_initialize_population(mock_meta_llm, mock_predictor, initial_promp df_few_shots=mock_df, ) - # Mock the _create_few_shot_examples method to simplify - def mock_create_few_shot_examples(instruction, num_examples): - return [f"Example {i}" for i in range(num_examples)] - - optimizer._create_few_shot_examples = mock_create_few_shot_examples - - # Control randomness with patch("random.randint", return_value=2): - population = optimizer._initialize_population([Prompt(p) for p in initial_prompts]) + optimizer._pre_optimization_loop() + population = optimizer.prompts # Verify population was created assert len(population) == len(initial_prompts) @@ -69,17 +64,16 @@ def test_capo_step(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mo # Mock the internal methods to avoid complexity mock_offspring = [Prompt("Offspring", ["Example"])] - optimizer._crossover = lambda x: mock_offspring - mock_mutated = [Prompt("Mutated", ["Example"])] - optimizer._mutate = lambda x: mock_mutated - - mock_survivors = [Prompt("Survivor 1", ["Example"]), Prompt("Survivor 2", ["Example"])] - mock_scores = [0.9, 0.8] - optimizer._do_racing = lambda x, k: (mock_survivors, mock_scores) + with patch("promptolution.optimizers.capo.perform_crossover", return_value=mock_offspring), patch( + "promptolution.optimizers.capo.perform_mutation", return_value=mock_mutated + ): + mock_survivors = [Prompt("Survivor 1", ["Example"]), Prompt("Survivor 2", ["Example"])] + mock_scores = [0.9, 0.8] + optimizer._do_racing = lambda x, k: (mock_survivors, mock_scores) - # Call _step - result = optimizer._step() + # Call _step + result = optimizer._step() # Verify results assert len(result) == 2 # Should match population_size @@ -117,7 +111,7 @@ def mock_step(): def test_create_few_shots(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): - """Test the _create_few_shot_examples method.""" + """Test the few-shot example builder.""" optimizer = CAPO( predictor=mock_predictor, task=mock_task, @@ -126,14 +120,37 @@ def test_create_few_shots(mock_meta_llm, mock_predictor, initial_prompts, mock_t df_few_shots=mock_df, ) - # Call the method - few_shot_examples = optimizer._create_few_shot_examples("Classify the sentiment of the text.", 2) + few_shot_examples = build_few_shot_examples( + instruction="Classify the sentiment of the text.", + num_examples=2, + df_few_shots=mock_df, + x_column=mock_task.x_column, + y_column=mock_task.y_column, + predictor=mock_predictor, + fewshot_template=CAPO_FEWSHOT_TEMPLATE, + target_begin_marker=optimizer.target_begin_marker, + target_end_marker=optimizer.target_end_marker, + check_fs_accuracy=True, + create_fs_reasoning=True, + ) # Verify results assert len(few_shot_examples) == 2 assert all(isinstance(example, str) for example in few_shot_examples) - few_shot_examples = optimizer._create_few_shot_examples("Classify the sentiment of the text.", 0) + few_shot_examples = build_few_shot_examples( + instruction="Classify the sentiment of the text.", + num_examples=0, + df_few_shots=mock_df, + x_column=mock_task.x_column, + y_column=mock_task.y_column, + predictor=mock_predictor, + fewshot_template=CAPO_FEWSHOT_TEMPLATE, + target_begin_marker=optimizer.target_begin_marker, + target_end_marker=optimizer.target_end_marker, + check_fs_accuracy=True, + create_fs_reasoning=True, + ) assert len(few_shot_examples) == 0 @@ -148,7 +165,12 @@ def test_crossover(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mo crossovers_per_iter=5, ) - offsprings = optimizer._crossover([Prompt("Instruction 1", ["Example 1"]), Prompt("Instruction 2", ["Example 2"])]) + offsprings = perform_crossover( + [Prompt("Instruction 1", ["Example 1"]), Prompt("Instruction 2", ["Example 2"])], + optimizer.crossovers_per_iter, + optimizer.crossover_template, + optimizer.meta_llm, + ) assert len(offsprings) == 5 @@ -161,7 +183,23 @@ def test_mutate(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_ df_few_shots=mock_df, ) - mutated = optimizer._mutate([Prompt("Instruction 1", ["Example 1"]), Prompt("Instruction 2", ["Example 2"])]) + mutated = perform_mutation( + offsprings=[Prompt("Instruction 1", ["Example 1"]), Prompt("Instruction 2", ["Example 2"])], + mutation_template=optimizer.mutation_template, + upper_shots=optimizer.upper_shots, + meta_llm=optimizer.meta_llm, + few_shot_kwargs=dict( + df_few_shots=mock_df, + x_column=mock_task.x_column, + y_column=mock_task.y_column, + predictor=mock_predictor, + fewshot_template=CAPO_FEWSHOT_TEMPLATE, + target_begin_marker=optimizer.target_begin_marker, + target_end_marker=optimizer.target_end_marker, + check_fs_accuracy=True, + create_fs_reasoning=True, + ), + ) assert len(mutated) == 2 @@ -199,7 +237,7 @@ def test_capo_crossover_prompt(mock_meta_llm, mock_predictor, initial_prompts, m mother = Prompt("Classify the sentiment of the text.", ["Input: I love this! Output: Positive"]) father = Prompt("Determine if the review is positive or negative.", ["Input: This is terrible. Output: Negative"]) - optimizer._crossover([mother, father]) + perform_crossover([mother, father], optimizer.crossovers_per_iter, optimizer.crossover_template, optimizer.meta_llm) full_task_desc = mock_task.task_description + "\n" + optimizer.predictor.extraction_description @@ -208,8 +246,13 @@ def test_capo_crossover_prompt(mock_meta_llm, mock_predictor, initial_prompts, m .replace("", father.instruction) .replace("", full_task_desc) ) + alt_meta_prompt = ( + CAPO_CROSSOVER_TEMPLATE.replace("", father.instruction) + .replace("", mother.instruction) + .replace("", full_task_desc) + ) - assert str(mock_meta_llm.call_history[0]["prompts"][0]) == expected_meta_prompt + assert str(mock_meta_llm.call_history[0]["prompts"][0]) in {expected_meta_prompt, alt_meta_prompt} def test_capo_mutate_prompt(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): @@ -224,7 +267,23 @@ def test_capo_mutate_prompt(mock_meta_llm, mock_predictor, initial_prompts, mock full_task_desc = mock_task.task_description + "\n" + optimizer.predictor.extraction_description parent = Prompt("Classify the sentiment of the text.", ["Input: I love this! Output: Positive"]) - optimizer._mutate([parent]) + perform_mutation( + offsprings=[parent], + mutation_template=optimizer.mutation_template, + upper_shots=optimizer.upper_shots, + meta_llm=optimizer.meta_llm, + few_shot_kwargs=dict( + df_few_shots=mock_df, + x_column=mock_task.x_column, + y_column=mock_task.y_column, + predictor=mock_predictor, + fewshot_template=CAPO_FEWSHOT_TEMPLATE, + target_begin_marker=optimizer.target_begin_marker, + target_end_marker=optimizer.target_end_marker, + check_fs_accuracy=True, + create_fs_reasoning=True, + ), + ) expected_meta_prompt = CAPO_MUTATION_TEMPLATE.replace("", parent.instruction).replace( "", full_task_desc diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py new file mode 100644 index 00000000..580a0c78 --- /dev/null +++ b/tests/optimizers/test_capoeira.py @@ -0,0 +1,107 @@ +from unittest.mock import MagicMock, patch + +import pandas as pd + +from promptolution.optimizers.capoeira import Capoeira +from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation +from promptolution.utils.prompt import Prompt +from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE + + +def test_capoeira_initialization(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): + optimizer = Capoeira( + predictor=mock_predictor, + task=mock_task, + meta_llm=mock_meta_llm, + initial_prompts=initial_prompts, + df_few_shots=mock_df, + population_size=None, + ) + + assert optimizer.crossovers_per_iter == 4 + assert optimizer.population_size == len(initial_prompts) + assert isinstance(optimizer.df_few_shots, pd.DataFrame) + + +def test_capoeira_initialize_population(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): + optimizer = Capoeira( + predictor=mock_predictor, + task=mock_task, + meta_llm=mock_meta_llm, + initial_prompts=initial_prompts, + df_few_shots=mock_df, + ) + + with patch("random.randint", return_value=1): + optimizer._pre_optimization_loop() + population = optimizer.prompts + + assert len(population) == len(initial_prompts) + assert all(isinstance(p, Prompt) for p in population) + + +def test_capoeira_selection_prefers_better_score(mock_meta_llm, mock_predictor, mock_task, mock_df): + optimizer = Capoeira( + predictor=mock_predictor, + task=mock_task, + meta_llm=mock_meta_llm, + initial_prompts=["short", "longer prompt"], + df_few_shots=mock_df, + population_size=1, + ) + optimizer.token_counter = lambda _: 1 + candidates = [Prompt("short"), Prompt("longer prompt")] + optimizer.task.evaluate = MagicMock(return_value=[0.4, 0.9]) + + objectives = optimizer._evaluate_candidates(candidates) + selected, _ = optimizer._select_population(candidates, objectives) + + assert len(selected) == 1 + assert selected[0].instruction == "longer prompt" + + +def test_capoeira_meta_prompts(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): + optimizer = Capoeira( + predictor=mock_predictor, + task=mock_task, + meta_llm=mock_meta_llm, + initial_prompts=initial_prompts, + df_few_shots=mock_df, + crossovers_per_iter=2, + ) + + mother = Prompt("Instruction 1", ["Example 1"]) + father = Prompt("Instruction 2", ["Example 2"]) + perform_crossover([mother, father], optimizer.crossovers_per_iter, optimizer.crossover_template, optimizer.meta_llm) + + full_task_desc = mock_task.task_description + "\n" + optimizer.predictor.extraction_description + expected_crossover = ( + CAPO_CROSSOVER_TEMPLATE.replace("", mother.instruction) + .replace("", father.instruction) + .replace("", full_task_desc) + ) + assert expected_crossover in mock_meta_llm.call_history[0]["prompts"] + + mock_meta_llm.reset() + parent = Prompt("Instruction 3", ["Example 3"]) + perform_mutation( + offsprings=[parent], + mutation_template=optimizer.mutation_template, + upper_shots=optimizer.upper_shots, + meta_llm=optimizer.meta_llm, + few_shot_kwargs=dict( + df_few_shots=mock_df, + x_column=mock_task.x_column, + y_column=mock_task.y_column, + predictor=mock_predictor, + fewshot_template=CAPO_FEWSHOT_TEMPLATE, + target_begin_marker=optimizer.target_begin_marker, + target_end_marker=optimizer.target_end_marker, + check_fs_accuracy=True, + create_fs_reasoning=True, + ), + ) + expected_mutation = CAPO_MUTATION_TEMPLATE.replace("", parent.instruction).replace( + "", full_task_desc + ) + assert expected_mutation in mock_meta_llm.call_history[0]["prompts"] diff --git a/tutorials/getting_started.ipynb b/tutorials/getting_started.ipynb index 2c140f6d..b049196d 100644 --- a/tutorials/getting_started.ipynb +++ b/tutorials/getting_started.ipynb @@ -163,7 +163,7 @@ "metadata": {}, "source": [ "Here's an explanation of each configuration parameter in the ExperimentConfig:\n", - "- `optimizer`: The algorithm used for prompt optimization. Currently we support \"capo\", \"evopromptga\", \"evopromptde\", and \"opro\". For this example, we use \"capo\" as it is capable of leveraging few-shot examples.\n", + "- `optimizer`: The algorithm used for prompt optimization. Currently we support \"capo\", \"capoeira\", \"evopromptga\", \"evopromptde\", and \"opro\". For this example, we use \"capo\" as it is capable of leveraging few-shot examples.\n", "- `task_description`: A string describing the task you're optimizing prompts for. This is used to provide the meta-llm with context about your task.\n", "- `prompts`: A list of initial prompt strings that will be used as the starting point for optimization.\n", "- `n_steps`: The number of optimization steps to run. Higher values allow more exploration and refinement but require more API calls and computational resources.\n", @@ -399,4 +399,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file diff --git a/tutorials/reward_task_tutorial.ipynb b/tutorials/reward_task_tutorial.ipynb index e0922408..56d17b2a 100644 --- a/tutorials/reward_task_tutorial.ipynb +++ b/tutorials/reward_task_tutorial.ipynb @@ -201,7 +201,7 @@ "metadata": {}, "source": [ "Here's an explanation of each configuration parameter in the ExperimentConfig:\n", - "- `optimizer`: The algorithm used for prompt optimization. Currently we support \"capo\", \"evopromptga\", \"evopromptde\", and \"opro\". For this example, we use \"capo\" as it is capable of leveraging few-shot examples.\n", + "- `optimizer`: The algorithm used for prompt optimization. Currently we support \"capo\", \"capoeira\", \"evopromptga\", \"evopromptde\", and \"opro\". For this example, we use \"capo\" as it is capable of leveraging few-shot examples.\n", "- `task_description`: A string describing the task you're optimizing prompts for. This is used to provide the meta-llm with context about your task.\n", "- `prompts`: A list of initial prompt strings that will be used as the starting point for optimization.\n", "- `n_steps`: The number of optimization steps to run. Higher values allow more exploration and refinement but require more API calls and computational resources.\n", @@ -466,4 +466,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} +} \ No newline at end of file From 3961739a412d9da4cc6923b4ca312c8dfa68c0b3 Mon Sep 17 00:00:00 2001 From: finitearth Date: Sat, 27 Dec 2025 15:07:07 +0100 Subject: [PATCH 04/53] debugging --- promptolution/optimizers/capo.py | 4 ++-- promptolution/optimizers/capo_utils.py | 2 +- promptolution/optimizers/capoeira.py | 10 ++++------ promptolution/utils/capo_utils.py | 4 ++-- 4 files changed, 9 insertions(+), 11 deletions(-) diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index 41174112..73811f2c 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -115,8 +115,8 @@ def __init__( self.population_size = len(self.prompts) if hasattr(self.predictor, "begin_marker") and hasattr(self.predictor, "end_marker"): - self.target_begin_marker = self.predictor.begin_marker - self.target_end_marker = self.predictor.end_marker + self.target_begin_marker = self.predictor.begin_marker # type: ignore + self.target_end_marker = self.predictor.end_marker # type: ignore else: self.target_begin_marker = "" self.target_end_marker = "" diff --git a/promptolution/optimizers/capo_utils.py b/promptolution/optimizers/capo_utils.py index 45bc2ff9..38635aa4 100644 --- a/promptolution/optimizers/capo_utils.py +++ b/promptolution/optimizers/capo_utils.py @@ -3,7 +3,7 @@ from __future__ import annotations import random -from typing import Callable, List, Optional +from typing import Callable, List import pandas as pd diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index d35887a7..5bf4c283 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -39,7 +39,6 @@ def __init__( mutation_template: Optional[str] = None, crossovers_per_iter: int = 4, upper_shots: int = 5, - population_size: Optional[int] = None, cost_per_input_token: float = 1.0, cost_per_output_token: float = 0.0, check_fs_accuracy: bool = True, @@ -59,7 +58,6 @@ def __init__( mutation_template: Optional meta-prompt template for mutation. crossovers_per_iter: Number of crossover operations per iteration. upper_shots: Maximum number of few-shot examples to attach. - population_size: Target population size used when pruning fronts. cost_per_input_token: Weight applied to input token cost for the cost objective. cost_per_output_token: Weight applied to output token cost for the cost objective. check_fs_accuracy: Whether to verify few-shot correctness before use. @@ -72,7 +70,7 @@ def __init__( self.downstream_llm = predictor.llm self.crossovers_per_iter = crossovers_per_iter self.upper_shots = upper_shots - self.population_size = population_size + self.cost_per_input_token = cost_per_input_token self.cost_per_output_token = cost_per_output_token self.check_fs_accuracy = check_fs_accuracy @@ -84,11 +82,11 @@ def __init__( self.mutation_template = self._initialize_meta_template(mutation_template or CAPO_MUTATION_TEMPLATE) self.token_counter = get_token_counter(self.downstream_llm) self.df_few_shots = df_few_shots if df_few_shots is not None else task.pop_datapoints(frac=0.1) - self.population_size = self.population_size or len(self.prompts) or 1 + self.population_size = len(self.prompts) if hasattr(self.predictor, "begin_marker") and hasattr(self.predictor, "end_marker"): - self.target_begin_marker = self.predictor.begin_marker - self.target_end_marker = self.predictor.end_marker + self.target_begin_marker = self.predictor.begin_marker # type: ignore + self.target_end_marker = self.predictor.end_marker # type: ignore else: self.target_begin_marker = "" self.target_end_marker = "" diff --git a/promptolution/utils/capo_utils.py b/promptolution/utils/capo_utils.py index d451476d..e5e720b3 100644 --- a/promptolution/utils/capo_utils.py +++ b/promptolution/utils/capo_utils.py @@ -3,7 +3,7 @@ from __future__ import annotations import random -from typing import Callable, List +from typing import List, Optional import pandas as pd @@ -16,7 +16,7 @@ def build_few_shot_examples( num_examples: int, df_few_shots: pd.DataFrame, x_column: str, - y_column: str, + y_column: Optional[str], predictor, fewshot_template: str, target_begin_marker: str, From b64b53de54bea7f2cd35f252da6a54f83ded6243 Mon Sep 17 00:00:00 2001 From: finitearth Date: Sat, 27 Dec 2025 15:10:46 +0100 Subject: [PATCH 05/53] delete redundant util file --- promptolution/optimizers/capo.py | 1 - promptolution/optimizers/capo_utils.py | 111 ------------------------- promptolution/optimizers/capoeira.py | 1 + 3 files changed, 1 insertion(+), 112 deletions(-) delete mode 100644 promptolution/optimizers/capo_utils.py diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index 73811f2c..f259774a 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -16,7 +16,6 @@ from promptolution.utils.test_statistics import TestStatistics from promptolution.optimizers.base_optimizer import BaseOptimizer -from promptolution.utils.formatting import extract_from_tag from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation from promptolution.utils.logging import get_logger from promptolution.utils.prompt import Prompt, sort_prompts_by_scores diff --git a/promptolution/optimizers/capo_utils.py b/promptolution/optimizers/capo_utils.py deleted file mode 100644 index 38635aa4..00000000 --- a/promptolution/optimizers/capo_utils.py +++ /dev/null @@ -1,111 +0,0 @@ -"""Shared utilities for CAPO-style optimizers.""" - -from __future__ import annotations - -import random -from typing import Callable, List - -import pandas as pd - -from promptolution.utils.formatting import extract_from_tag -from promptolution.utils.prompt import Prompt - - -def build_few_shot_examples( - instruction: str, - num_examples: int, - df_few_shots: pd.DataFrame, - task, - predictor, - fewshot_template: str, - target_begin_marker: str, - target_end_marker: str, - check_fs_accuracy: bool, - create_fs_reasoning: bool, -) -> List[str]: - """Create few-shot examples with optional reasoning replacement.""" - if num_examples == 0: - return [] - - few_shot_samples = df_few_shots.sample(num_examples, replace=False) - sample_inputs = few_shot_samples[task.x_column].values.astype(str) - sample_targets = few_shot_samples[task.y_column].values - few_shots = [ - fewshot_template.replace("", i).replace("", f"{target_begin_marker}{t}{target_end_marker}") - for i, t in zip(sample_inputs, sample_targets) - ] - - if not create_fs_reasoning: - return few_shots - - preds, seqs = predictor.predict( - [instruction] * num_examples, - list(sample_inputs), - return_seq=True, - ) - if isinstance(seqs, str): - seqs = [seqs] - if isinstance(preds, str): - preds = [preds] - - for j in range(num_examples): - seqs[j] = seqs[j].replace(sample_inputs[j], "", 1).strip() - if preds[j] == sample_targets[j] or not check_fs_accuracy: - few_shots[j] = fewshot_template.replace("", sample_inputs[j]).replace("", seqs[j]) - - return few_shots - - -def perform_crossover( - parents: List[Prompt], - crossovers_per_iter: int, - template: str, - meta_llm, -) -> List[Prompt]: - """Generate crossover offspring prompts.""" - crossover_prompts: List[str] = [] - offspring_few_shots: List[List[str]] = [] - for _ in range(crossovers_per_iter): - mother, father = (parents if len(parents) == 2 else random.sample(parents, 2)) - crossover_prompt = template.replace("", mother.instruction).replace("", father.instruction).strip() - crossover_prompts.append(crossover_prompt) - combined_few_shots = mother.few_shots + father.few_shots - num_few_shots = (len(mother.few_shots) + len(father.few_shots)) // 2 - offspring_few_shot = random.sample(combined_few_shots, num_few_shots) if combined_few_shots else [] - offspring_few_shots.append(offspring_few_shot) - - child_instructions = meta_llm.get_response(crossover_prompts) - return [ - Prompt(extract_from_tag(instr, "", ""), examples) - for instr, examples in zip(child_instructions, offspring_few_shots) - ] - - -def perform_mutation( - offsprings: List[Prompt], - mutation_template: str, - create_few_shots: Callable[[str, int], List[str]], - upper_shots: int, - meta_llm, -) -> List[Prompt]: - """Mutate offspring prompts.""" - mutation_prompts = [mutation_template.replace("", prompt.instruction) for prompt in offsprings] - new_instructions = meta_llm.get_response(mutation_prompts) - - mutated: List[Prompt] = [] - for new_instruction, prompt in zip(new_instructions, offsprings): - new_instruction = extract_from_tag(new_instruction, "", "") - p = random.random() - - if p < 1 / 3 and len(prompt.few_shots) < upper_shots: - new_few_shot = create_few_shots(new_instruction, 1) - new_few_shots = prompt.few_shots + new_few_shot - elif 1 / 3 <= p < 2 / 3 and len(prompt.few_shots) > 0: - new_few_shots = random.sample(prompt.few_shots, len(prompt.few_shots) - 1) - else: - new_few_shots = prompt.few_shots - - random.shuffle(new_few_shots) - mutated.append(Prompt(new_instruction, new_few_shots)) - - return mutated diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 5bf4c283..6ed49a76 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -111,6 +111,7 @@ def _pre_optimization_loop(self) -> None: population.append(Prompt(prompt.instruction, few_shots)) self.prompts = population + # TODO: align placement of the logic with capo self.max_prompt_length = max(self.token_counter(p.construct_prompt()) for p in self.prompts) if self.prompts else 1 initial_vectors = self._evaluate_candidates(self.prompts) self.prompts, selected_vectors = self._select_population(self.prompts, initial_vectors) From f65deefdaf17128885811de2d2d6929af56c438a Mon Sep 17 00:00:00 2001 From: finitearth Date: Sat, 27 Dec 2025 15:19:05 +0100 Subject: [PATCH 06/53] reduce code complexity --- promptolution/optimizers/capo.py | 38 ++++++++++++++-------------- promptolution/optimizers/capoeira.py | 18 ++----------- promptolution/tasks/base_task.py | 32 ++++++++--------------- 3 files changed, 31 insertions(+), 57 deletions(-) diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index f259774a..ea367ee4 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -221,26 +221,26 @@ def _step(self) -> List[Prompt]: return self.prompts + @staticmethod + def filter_survivors( + candidates: List[Prompt], scores: List[List[float]], mask: Any + ) -> Tuple[List[Prompt], List[List[float]]]: + """Filter candidates and scores based on a boolean mask. -def filter_survivors( - candidates: List[Prompt], scores: List[List[float]], mask: Any -) -> Tuple[List[Prompt], List[List[float]]]: - """Filter candidates and scores based on a boolean mask. - - Args: - candidates (List[Prompt]): List of candidate prompts. - scores (List[List[float]]): Corresponding scores for the candidates. - mask (Any): Boolean mask indicating which candidates to keep. + Args: + candidates (List[Prompt]): List of candidate prompts. + scores (List[List[float]]): Corresponding scores for the candidates. + mask (Any): Boolean mask indicating which candidates to keep. - Returns: - Tuple[List[Prompt], List[List[float]]]: Filtered candidates and their scores. - """ - assert len(candidates) == len(mask), "Length of candidates, and mask must be the same." - assert all( - len(candidates) == len(score) for score in scores - ), "Each score list must have the same length as candidates." + Returns: + Tuple[List[Prompt], List[List[float]]]: Filtered candidates and their scores. + """ + assert len(candidates) == len(mask), "Length of candidates, and mask must be the same." + assert all( + len(candidates) == len(score) for score in scores + ), "Each score list must have the same length as candidates." - filtered_candidates = [c for c, m in zip(candidates, mask) if m] - filtered_scores = [[s for s, m in zip(score, mask) if m] for score in scores] + filtered_candidates = [c for c, m in zip(candidates, mask) if m] + filtered_scores = [[s for s, m in zip(score, mask) if m] for score in scores] - return filtered_candidates, filtered_scores + return filtered_candidates, filtered_scores diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 6ed49a76..0ceb0124 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -118,7 +118,7 @@ def _pre_optimization_loop(self) -> None: self.scores = (-selected_vectors[:, 0]).tolist() def _evaluate_candidates(self, candidates: List[Prompt]) -> np.ndarray: - evaluation = self.task.evaluate( + scores, input_tokens, output_tokens = self.task.evaluate( candidates, self.predictor, eval_strategy=self.task.eval_strategy, @@ -127,24 +127,10 @@ def _evaluate_candidates(self, candidates: List[Prompt]) -> np.ndarray: return_agg_scores=True, ) - if isinstance(evaluation, tuple) and len(evaluation) == 3: - scores, input_tokens, output_tokens = evaluation - else: - scores = evaluation # type: ignore[assignment] - input_tokens = [self.token_counter(c.construct_prompt()) for c in candidates] - output_tokens = [0.0 for _ in candidates] - - input_tokens_arr = np.array(input_tokens, dtype=float) - output_tokens_arr = np.array(output_tokens, dtype=float) - - if not input_tokens_arr.any() and not output_tokens_arr.any(): - input_tokens_arr = np.array([self.token_counter(c.construct_prompt()) for c in candidates], dtype=float) - output_tokens_arr = np.zeros_like(input_tokens_arr, dtype=float) - score_vectors = np.column_stack( [ -np.array(scores, dtype=float), - self.cost_per_input_token * input_tokens_arr + self.cost_per_output_token * output_tokens_arr, + self.cost_per_input_token * input_tokens + self.cost_per_output_token * output_tokens, ] ) return score_vectors diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 6d897162..620e2bc8 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -173,7 +173,8 @@ def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> List[floa """ raise NotImplementedError - @overload + # TODO: create overload for return_costs=True + @overload def evaluate( self, prompts: List[Prompt], @@ -270,7 +271,6 @@ def evaluate( assert not return_seq or not return_costs, "Token cost reporting is not supported together with sequences." seqs: List[str] = [] - token_counter = get_token_counter(predictor.llm) if return_costs else None prompts = [prompts] if isinstance(prompts, Prompt) else prompts eval_strategy = eval_strategy or self.eval_strategy @@ -310,37 +310,25 @@ def evaluate( if not return_costs: return agg_scores + token_counter = get_token_counter(predictor.llm) + per_prompt_inputs: List[float] = [] per_prompt_outputs: List[float] = [] - if token_counter is None: - logger.warning("⚠️ Token counting unavailable; returning zero costs.") - per_prompt_inputs = [0.0 for _ in prompts] - per_prompt_outputs = [0.0 for _ in prompts] - return agg_scores, per_prompt_inputs, per_prompt_outputs - - preds_by_prompt: List[List[str]] = [] - if isinstance(preds, list): - if preds and isinstance(preds[0], list): - preds_by_prompt = preds # type: ignore[assignment] - elif preds and isinstance(preds[0], str): - preds_by_prompt = [preds for _ in prompts] - - xs_token_mean = float(np.mean([token_counter(x) for x in xs])) if xs else 0.0 + xs_token_mean = np.mean([token_counter(x) for x in xs]) for idx, prompt in enumerate(prompts): prompt_tokens = token_counter(prompt.construct_prompt()) input_tokens = prompt_tokens + xs_token_mean - - if preds_by_prompt and idx < len(preds_by_prompt) and preds_by_prompt[idx]: - avg_output = float(np.mean([token_counter(p) for p in preds_by_prompt[idx]])) - else: - avg_output = 0.0 - logger.warning("⚠️ Unable to estimate output tokens; defaulting to 0.") + avg_output = np.mean([token_counter(p) for p in preds[idx]]) per_prompt_inputs.append(input_tokens) per_prompt_outputs.append(avg_output) + # convert to numpy + per_prompt_inputs = np.array(per_prompt_inputs, dtype=float) + per_prompt_outputs = np.array(per_prompt_outputs, dtype=float) + return agg_scores, per_prompt_inputs, per_prompt_outputs def pop_datapoints(self, n: Optional[int] = None, frac: Optional[float] = None) -> pd.DataFrame: From c532b48d8b2aac50d6c5cdad6fba1703d910db84 Mon Sep 17 00:00:00 2001 From: finitearth Date: Sat, 27 Dec 2025 17:46:21 +0100 Subject: [PATCH 07/53] clean up for pre commit --- promptolution/helpers.py | 12 +-- promptolution/llms/api_llm.py | 2 +- promptolution/optimizers/capo.py | 10 +-- promptolution/optimizers/capoeira.py | 23 ++++-- promptolution/tasks/base_task.py | 108 +++++++++++++++++++++------ promptolution/tasks/judge_tasks.py | 2 +- promptolution/utils/callbacks.py | 6 +- promptolution/utils/capo_utils.py | 9 ++- promptolution/utils/formatting.py | 2 +- promptolution/utils/prompt.py | 8 +- tests/conftest.py | 2 +- tests/optimizers/test_capoeira.py | 2 +- tutorials/getting_started.ipynb | 2 +- tutorials/reward_task_tutorial.ipynb | 2 +- 14 files changed, 134 insertions(+), 56 deletions(-) diff --git a/promptolution/helpers.py b/promptolution/helpers.py index bb0624e5..94879df3 100644 --- a/promptolution/helpers.py +++ b/promptolution/helpers.py @@ -80,7 +80,9 @@ def run_optimization(df: pd.DataFrame, config: "ExperimentConfig") -> List[Promp ) config.prompts = [Prompt(p) for p in initial_prompts] - if config.optimizer in {"capo", "capoeira"} and (config.eval_strategy is None or "block" not in config.eval_strategy): + if config.optimizer in {"capo", "capoeira"} and ( + config.eval_strategy is None or "block" not in config.eval_strategy + ): logger.warning("📌 CAPO-style optimizers require block evaluation strategy. Setting it to 'sequential_block'.") config.eval_strategy = "sequential_block" @@ -133,7 +135,7 @@ def run_evaluation( def get_llm(model_id: Optional[str] = None, config: Optional["ExperimentConfig"] = None) -> "BaseLLM": - """Factory function to create and return a language model instance based on the provided model_id. + """Create and return a language model instance based on the provided model_id. This function supports three types of language models: 1. LocalLLM: For running models locally. @@ -208,7 +210,7 @@ def get_optimizer( task_description: Optional[str] = None, config: Optional["ExperimentConfig"] = None, ) -> "BaseOptimizer": - """Creates and returns an optimizer instance based on provided parameters. + """Create and return an optimizer instance based on provided parameters. Args: predictor: The predictor used for prompt evaluation @@ -262,7 +264,7 @@ def get_optimizer( def get_exemplar_selector( name: Literal["random", "random_search"], task: "BaseTask", predictor: "BasePredictor" ) -> "BaseExemplarSelector": - """Factory function to get an exemplar selector based on the given name. + """Get an exemplar selector based on the given name. Args: name (str): The name of the exemplar selector to instantiate. @@ -284,7 +286,7 @@ def get_exemplar_selector( def get_predictor(downstream_llm=None, type: "PredictorType" = "marker", *args, **kwargs) -> "BasePredictor": - """Factory function to create and return a predictor instance. + """Create and return a predictor instance. This function supports three types of predictors: 1. FirstOccurrencePredictor: A predictor that classifies based on first occurrence of the label. diff --git a/promptolution/llms/api_llm.py b/promptolution/llms/api_llm.py index c6971a63..1e34e532 100644 --- a/promptolution/llms/api_llm.py +++ b/promptolution/llms/api_llm.py @@ -210,7 +210,7 @@ def _submit(self, coro): return asyncio.run_coroutine_threadsafe(coro, self._loop) def _get_response(self, prompts: List[str], system_prompts: List[str]) -> List[str]: - """Synchronously obtain responses for a batch of prompts. + """Obtain responses synchronously for a batch of prompts. This is the main entrypoint used by external callers. It handles system prompt broadcasting and delegates the actual work to the async batch diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index ea367ee4..2f06b6a3 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -55,7 +55,7 @@ def __init__( callbacks: Optional[List["BaseCallback"]] = None, config: Optional["ExperimentConfig"] = None, ) -> None: - """Initializes the CAPOptimizer with various parameters for prompt evolution. + """Initialize the CAPOptimizer with various parameters for prompt evolution. Args: predictor (BasePredictor): The predictor for evaluating prompt performance. @@ -114,14 +114,14 @@ def __init__( self.population_size = len(self.prompts) if hasattr(self.predictor, "begin_marker") and hasattr(self.predictor, "end_marker"): - self.target_begin_marker = self.predictor.begin_marker # type: ignore - self.target_end_marker = self.predictor.end_marker # type: ignore + self.target_begin_marker = self.predictor.begin_marker # type: ignore + self.target_end_marker = self.predictor.end_marker # type: ignore else: self.target_begin_marker = "" self.target_end_marker = "" def _initialize_population(self, initial_prompts: List[Prompt]) -> List[Prompt]: - """Initializes the population of Prompt objects from initial instructions.""" + """Initialize the population of Prompt objects from initial instructions.""" population = [] for prompt in initial_prompts: num_examples = random.randint(0, self.upper_shots) @@ -178,7 +178,7 @@ def _do_racing(self, candidates: List[Prompt], k: int) -> Tuple[List[Prompt], Li # Sum along rows to get number of better scores for each candidate n_better = np.sum(comparison_matrix, axis=1) - candidates, block_scores = filter_survivors(candidates, block_scores, mask=n_better < k) + candidates, block_scores = self.filter_survivors(candidates, block_scores, mask=n_better < k) i += 1 self.task.increment_block_idx() diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 0ceb0124..60ebcc90 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -3,11 +3,12 @@ from __future__ import annotations import random -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple import numpy as np import pandas as pd +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + if TYPE_CHECKING: # pragma: no cover from promptolution.utils.callbacks import BaseCallback from promptolution.llms.base_llm import BaseLLM @@ -17,7 +18,6 @@ from promptolution.optimizers.base_optimizer import BaseOptimizer from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation -from promptolution.utils.formatting import extract_from_tag from promptolution.utils.logging import get_logger from promptolution.utils.prompt import Prompt from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE @@ -85,7 +85,7 @@ def __init__( self.population_size = len(self.prompts) if hasattr(self.predictor, "begin_marker") and hasattr(self.predictor, "end_marker"): - self.target_begin_marker = self.predictor.begin_marker # type: ignore + self.target_begin_marker = self.predictor.begin_marker # type: ignore self.target_end_marker = self.predictor.end_marker # type: ignore else: self.target_begin_marker = "" @@ -112,7 +112,9 @@ def _pre_optimization_loop(self) -> None: self.prompts = population # TODO: align placement of the logic with capo - self.max_prompt_length = max(self.token_counter(p.construct_prompt()) for p in self.prompts) if self.prompts else 1 + self.max_prompt_length = ( + max(self.token_counter(p.construct_prompt()) for p in self.prompts) if self.prompts else 1 + ) initial_vectors = self._evaluate_candidates(self.prompts) self.prompts, selected_vectors = self._select_population(self.prompts, initial_vectors) self.scores = (-selected_vectors[:, 0]).tolist() @@ -127,15 +129,22 @@ def _evaluate_candidates(self, candidates: List[Prompt]) -> np.ndarray: return_agg_scores=True, ) + # TODO move to evaluate method! + input_tokens_array = np.array(input_tokens, dtype=float) + output_tokens_array = np.array(output_tokens, dtype=float) + scores_array = np.array(scores, dtype=float) + score_vectors = np.column_stack( [ - -np.array(scores, dtype=float), - self.cost_per_input_token * input_tokens + self.cost_per_output_token * output_tokens, + -scores_array, + self.cost_per_input_token * input_tokens_array + self.cost_per_output_token * output_tokens_array, ] ) return score_vectors - def _select_population(self, candidates: List[Prompt], score_vectors: np.ndarray) -> Tuple[List[Prompt], np.ndarray]: + def _select_population( + self, candidates: List[Prompt], score_vectors: np.ndarray + ) -> Tuple[List[Prompt], np.ndarray]: selected_indices: List[int] = [] fronts = self.fast_non_dominated_sort(score_vectors) for front in fronts: diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 620e2bc8..a1e458cd 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -6,11 +6,11 @@ import numpy as np import pandas as pd -from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union, overload +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union, cast, overload +from promptolution.utils.logging import get_logger from promptolution.utils.prompt import Prompt from promptolution.utils.token_counter import get_token_counter -from promptolution.utils.logging import get_logger logger = get_logger(__name__) @@ -116,7 +116,7 @@ def _prepare_batch( ys: List[str], eval_strategy: Literal["full", "subsample", "sequential_block", "random_block", "evaluated"] = "full", ) -> List[Tuple[str, str, str]]: - """Generates (prompt, x, y) keys that require prediction. + """Generate (prompt, x, y) keys that require prediction. Returns keys not found in eval_cache. """ @@ -138,7 +138,7 @@ def _collect_results_from_cache( return_agg_scores: bool, return_seq: bool, ) -> Union[List[float], List[List[float]], Tuple[List[List[float]], List[List[str]]]]: - """Collects all results for the current batch from the cache and formats them.""" + """Collect all results for the current batch from the cache and format them.""" assert not (return_agg_scores and return_seq), "Cannot return both aggregated scores and sequences" scores = [] @@ -173,8 +173,7 @@ def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> List[floa """ raise NotImplementedError - # TODO: create overload for return_costs=True - @overload + @overload def evaluate( self, prompts: List[Prompt], @@ -183,6 +182,7 @@ def evaluate( return_agg_scores: Literal[True] = True, return_seq: Literal[False] = False, eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[False] = False, ) -> List[float]: ... @@ -195,6 +195,7 @@ def evaluate( return_agg_scores: Literal[False] = False, return_seq: Literal[False] = False, eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[False] = False, ) -> List[List[float]]: ... @@ -207,6 +208,7 @@ def evaluate( return_agg_scores: Literal[False] = False, return_seq: Literal[True] = True, eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[False] = False, ) -> Tuple[List[List[float]], List[List[str]]]: ... @@ -219,6 +221,7 @@ def evaluate( return_agg_scores: Literal[True] = True, return_seq: Literal[False] = False, eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[False] = False, ) -> List[float]: ... @@ -231,6 +234,7 @@ def evaluate( return_agg_scores: Literal[False] = False, return_seq: Literal[False] = False, eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[False] = False, ) -> List[List[float]]: ... @@ -243,9 +247,62 @@ def evaluate( return_agg_scores: Literal[False] = False, return_seq: Literal[True] = True, eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[False] = False, ) -> Tuple[List[List[float]], List[List[str]]]: ... + @overload + def evaluate( + self, + prompts: List[Prompt], + predictor: "BasePredictor", + system_prompts: Optional[Union[str, List[str]]] = None, + return_agg_scores: Literal[True] = True, + return_seq: Literal[False] = False, + eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[True] = True, + ) -> Tuple[List[float], List[float], List[float]]: + ... + + @overload + def evaluate( + self, + prompts: List[Prompt], + predictor: "BasePredictor", + system_prompts: Optional[Union[str, List[str]]] = None, + return_agg_scores: Literal[False] = False, + return_seq: Literal[False] = False, + eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[True] = True, + ) -> Tuple[List[List[float]], List[List[float]], List[List[float]]]: + ... + + @overload + def evaluate( + self, + prompts: Prompt, + predictor: "BasePredictor", + system_prompts: Optional[Union[str, List[str]]] = None, + return_agg_scores: Literal[True] = True, + return_seq: Literal[False] = False, + eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[True] = True, + ) -> Tuple[List[float], List[float], List[float]]: + ... + + @overload + def evaluate( + self, + prompts: Prompt, + predictor: "BasePredictor", + system_prompts: Optional[Union[str, List[str]]] = None, + return_agg_scores: Literal[False] = False, + return_seq: Literal[False] = False, + eval_strategy: Optional["EvalStrategy"] = None, + return_costs: Literal[True] = True, + ) -> Tuple[List[List[float]], List[List[float]], List[List[float]]]: + ... + def evaluate( self, prompts: Union[Prompt, List[Prompt]], @@ -260,6 +317,7 @@ def evaluate( List[List[float]], Tuple[List[List[float]], List[List[str]]], Tuple[List[float], List[float], List[float]], + Tuple[List[List[float]], List[List[float]], List[List[float]]], ]: """Evaluate a set of prompts using a given predictor. @@ -270,8 +328,6 @@ def evaluate( assert not (return_agg_scores and return_seq), "Cannot return both aggregated scores and sequences" assert not return_seq or not return_costs, "Token cost reporting is not supported together with sequences." - seqs: List[str] = [] - prompts = [prompts] if isinstance(prompts, Prompt) else prompts eval_strategy = eval_strategy or self.eval_strategy xs, ys = self.subsample(eval_strategy=eval_strategy) @@ -288,10 +344,11 @@ def evaluate( else: preds_seqs = ([], []) if return_seq else [] + seqs: List[str] = [] if return_seq: preds, seqs = preds_seqs if isinstance(preds_seqs, tuple) else (preds_seqs, []) else: - preds = preds_seqs + preds = cast(List[str], preds_seqs) scores: List[float] = self._evaluate(list(xs_to_evaluate), list(ys_to_evaluate), preds) for i, cache_key in enumerate(batches): @@ -311,25 +368,32 @@ def evaluate( return agg_scores token_counter = get_token_counter(predictor.llm) - - per_prompt_inputs: List[float] = [] - per_prompt_outputs: List[float] = [] - xs_token_mean = np.mean([token_counter(x) for x in xs]) + per_prompt_inputs: List[List[float]] = [] + per_prompt_outputs: List[List[float]] = [] + + input_token_counts = [float(token_counter(x)) for x in xs] for idx, prompt in enumerate(prompts): - prompt_tokens = token_counter(prompt.construct_prompt()) - input_tokens = prompt_tokens + xs_token_mean - avg_output = np.mean([token_counter(p) for p in preds[idx]]) + prompt_tokens = float(token_counter(prompt.construct_prompt())) + start = idx * len(xs) + end = (idx + 1) * len(xs) + preds_for_prompt = preds[start:end] + output_token_counts = [float(token_counter(p)) for p in preds_for_prompt] - per_prompt_inputs.append(input_tokens) - per_prompt_outputs.append(avg_output) + # Per-datapoint input tokens: prompt tokens + tokens of each x + prompt_input_tokens = [prompt_tokens + input_toks for input_toks in input_token_counts] + per_prompt_inputs.append(prompt_input_tokens) + per_prompt_outputs.append(output_token_counts) - # convert to numpy - per_prompt_inputs = np.array(per_prompt_inputs, dtype=float) - per_prompt_outputs = np.array(per_prompt_outputs, dtype=float) + if return_agg_scores: + agg_scores_list = cast(List[float], agg_scores) + per_prompt_inputs_mean = [float(np.mean(tokens)) for tokens in per_prompt_inputs] + per_prompt_outputs_mean = [float(np.mean(tokens)) for tokens in per_prompt_outputs] + return agg_scores_list, per_prompt_inputs_mean, per_prompt_outputs_mean - return agg_scores, per_prompt_inputs, per_prompt_outputs + score_matrix = cast(List[List[float]], agg_scores) + return score_matrix, per_prompt_inputs, per_prompt_outputs def pop_datapoints(self, n: Optional[int] = None, frac: Optional[float] = None) -> pd.DataFrame: """Pop a number of datapoints from the dataset. diff --git a/promptolution/tasks/judge_tasks.py b/promptolution/tasks/judge_tasks.py index 25704582..0f2fd4dc 100644 --- a/promptolution/tasks/judge_tasks.py +++ b/promptolution/tasks/judge_tasks.py @@ -112,7 +112,7 @@ def __init__( self.judge_llm = judge_llm def _construct_judge_prompt(self, x: str, pred: str, y: Optional[str] = None) -> str: - """Constructs the judge prompt based on whether ground truth is available.""" + """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: diff --git a/promptolution/utils/callbacks.py b/promptolution/utils/callbacks.py index 98129e2d..abad62f5 100644 --- a/promptolution/utils/callbacks.py +++ b/promptolution/utils/callbacks.py @@ -34,7 +34,7 @@ def __init__(self, **kwargs: Any) -> None: pass def on_step_end(self, optimizer: "BaseOptimizer") -> bool: - """Called at the end of each optimization step. + """Call at the end of each optimization step. Args: optimizer: The optimizer object that called the callback. @@ -45,7 +45,7 @@ def on_step_end(self, optimizer: "BaseOptimizer") -> bool: return True def on_epoch_end(self, optimizer: "BaseOptimizer") -> bool: - """Called at the end of each optimization epoch. + """Call at the end of each optimization epoch. Args: optimizer: The optimizer object that called the callback. @@ -56,7 +56,7 @@ def on_epoch_end(self, optimizer: "BaseOptimizer") -> bool: return True def on_train_end(self, optimizer: "BaseOptimizer") -> bool: - """Called at the end of the entire optimization process. + """Call at the end of the entire optimization process. Args: optimizer: The optimizer object that called the callback. diff --git a/promptolution/utils/capo_utils.py b/promptolution/utils/capo_utils.py index e5e720b3..1e18c915 100644 --- a/promptolution/utils/capo_utils.py +++ b/promptolution/utils/capo_utils.py @@ -3,10 +3,11 @@ from __future__ import annotations import random -from typing import List, Optional import pandas as pd +from typing import List, Optional + from promptolution.utils.formatting import extract_from_tag from promptolution.utils.prompt import Prompt @@ -67,8 +68,10 @@ def perform_crossover( crossover_prompts: List[str] = [] offspring_few_shots: List[List[str]] = [] for _ in range(crossovers_per_iter): - mother, father = (parents if len(parents) == 2 else random.sample(parents, 2)) - crossover_prompt = template.replace("", mother.instruction).replace("", father.instruction).strip() + mother, father = parents if len(parents) == 2 else random.sample(parents, 2) + crossover_prompt = ( + template.replace("", mother.instruction).replace("", father.instruction).strip() + ) crossover_prompts.append(crossover_prompt) combined_few_shots = mother.few_shots + father.few_shots num_few_shots = (len(mother.few_shots) + len(father.few_shots)) // 2 diff --git a/promptolution/utils/formatting.py b/promptolution/utils/formatting.py index 3d2ad198..c342baff 100644 --- a/promptolution/utils/formatting.py +++ b/promptolution/utils/formatting.py @@ -13,7 +13,7 @@ def extract_from_tag(text: List[str], start_tag: str, end_tag: str) -> List[str] def extract_from_tag(text: Union[str, List[str]], start_tag: str, end_tag: str) -> Union[List[str], str]: - """Extracts content from a string between specified start and end tags. + """Extract content from a string between specified start and end tags. Args: text (str): The input text to extract from. diff --git a/promptolution/utils/prompt.py b/promptolution/utils/prompt.py index d660e49d..e9067981 100644 --- a/promptolution/utils/prompt.py +++ b/promptolution/utils/prompt.py @@ -6,12 +6,12 @@ class Prompt: - """Represents a prompt consisting of an instruction and few-shot examples.""" + """Represent a prompt consisting of an instruction and few-shot examples.""" def __init__( self, instruction: str, few_shots: Optional[List[str]] = None, downstream_template: Optional[str] = None ) -> None: - """Initializes the Prompt with an instruction and associated examples. + """Initialize the Prompt with an instruction and associated examples. Args: instruction (str): The instruction or prompt text. @@ -28,7 +28,7 @@ def __init__( self.downstream_template = downstream_template def construct_prompt(self) -> str: - """Constructs the full prompt string by replacing placeholders in the template with the instruction and formatted examples. + """Construct the full prompt string by replacing placeholders in the template with the instruction and formatted examples. Returns: str: The constructed prompt string. @@ -43,7 +43,7 @@ def construct_prompt(self) -> str: return prompt def __str__(self) -> str: - """Returns the string representation of the prompt.""" + """Return the string representation of the prompt.""" return self.construct_prompt() diff --git a/tests/conftest.py b/tests/conftest.py index 2ba60f8e..d4499c54 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -104,7 +104,7 @@ def mock_classification_task_with_subsampling(mock_df): @pytest.fixture def simple_reward_function(): - """A simple reward function for testing RewardTask.""" + """Define a simple reward function for testing RewardTask.""" def reward_func(prediction: str) -> float: if "great" in prediction.lower() or "perfect" in prediction.lower(): diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py index 580a0c78..aa1e716e 100644 --- a/tests/optimizers/test_capoeira.py +++ b/tests/optimizers/test_capoeira.py @@ -3,7 +3,7 @@ import pandas as pd from promptolution.optimizers.capoeira import Capoeira -from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation +from promptolution.utils.capo_utils import perform_crossover, perform_mutation from promptolution.utils.prompt import Prompt from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE diff --git a/tutorials/getting_started.ipynb b/tutorials/getting_started.ipynb index b049196d..7761168c 100644 --- a/tutorials/getting_started.ipynb +++ b/tutorials/getting_started.ipynb @@ -399,4 +399,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} \ No newline at end of file +} diff --git a/tutorials/reward_task_tutorial.ipynb b/tutorials/reward_task_tutorial.ipynb index 56d17b2a..91f4af72 100644 --- a/tutorials/reward_task_tutorial.ipynb +++ b/tutorials/reward_task_tutorial.ipynb @@ -466,4 +466,4 @@ }, "nbformat": 4, "nbformat_minor": 2 -} \ No newline at end of file +} From 79395a0fc232eb325173c0f162cb60866bc2358d Mon Sep 17 00:00:00 2001 From: Tom Zehle Date: Sun, 28 Dec 2025 13:28:33 +0100 Subject: [PATCH 08/53] Potential fix for pull request finding 'Wrong name for an argument in a class instantiation' Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com> --- tests/optimizers/test_capoeira.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py index aa1e716e..66b79f2f 100644 --- a/tests/optimizers/test_capoeira.py +++ b/tests/optimizers/test_capoeira.py @@ -15,7 +15,6 @@ def test_capoeira_initialization(mock_meta_llm, mock_predictor, initial_prompts, meta_llm=mock_meta_llm, initial_prompts=initial_prompts, df_few_shots=mock_df, - population_size=None, ) assert optimizer.crossovers_per_iter == 4 From 0bb5e7d868af8034fe006f0a80e729b910d1dd15 Mon Sep 17 00:00:00 2001 From: finitearth Date: Sun, 28 Dec 2025 13:30:47 +0100 Subject: [PATCH 09/53] minor fixes --- promptolution/tasks/base_task.py | 4 +--- tests/optimizers/test_capoeira.py | 2 -- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index a1e458cd..47bccb90 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -12,14 +12,12 @@ from promptolution.utils.prompt import Prompt from promptolution.utils.token_counter import get_token_counter -logger = get_logger(__name__) - -logger = get_logger(__name__) if TYPE_CHECKING: # pragma: no cover from promptolution.predictors.base_predictor import BasePredictor from promptolution.utils.config import ExperimentConfig +logger = get_logger(__name__) TaskType = Literal["classification", "reward", "judge"] EvalStrategy = Literal["full", "subsample", "sequential_block", "random_block"] diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py index aa1e716e..a12e2816 100644 --- a/tests/optimizers/test_capoeira.py +++ b/tests/optimizers/test_capoeira.py @@ -15,7 +15,6 @@ def test_capoeira_initialization(mock_meta_llm, mock_predictor, initial_prompts, meta_llm=mock_meta_llm, initial_prompts=initial_prompts, df_few_shots=mock_df, - population_size=None, ) assert optimizer.crossovers_per_iter == 4 @@ -47,7 +46,6 @@ def test_capoeira_selection_prefers_better_score(mock_meta_llm, mock_predictor, meta_llm=mock_meta_llm, initial_prompts=["short", "longer prompt"], df_few_shots=mock_df, - population_size=1, ) optimizer.token_counter = lambda _: 1 candidates = [Prompt("short"), Prompt("longer prompt")] From f152077bff5b85d48aabead00f2c1681cf0907d5 Mon Sep 17 00:00:00 2001 From: finitearth Date: Sun, 28 Dec 2025 14:00:21 +0100 Subject: [PATCH 10/53] fix tests --- .coverage | Bin 69632 -> 69632 bytes .vscode/settings.json | 7 +++++++ promptolution/optimizers/capo.py | 2 +- promptolution/tasks/base_task.py | 31 +++++++++++++++--------------- tests/optimizers/test_capoeira.py | 5 ++--- 5 files changed, 25 insertions(+), 20 deletions(-) create mode 100644 .vscode/settings.json diff --git a/.coverage b/.coverage index 44c7c72d750d8c0618335cd34d0a10e261ad14b2..4f86ff21c29ceade31ad542efcb6270a1e4fdbd2 100644 GIT binary patch literal 69632 zcmeHQ3v?6LnI1`_w=`qpH#W8azhCkLV~k0Fq?C|QNU*U31mv+amIac=NU}}H22f~| zJ?(bWE!hy*6xwduw9c!$Z8yLYa9;<~CQU*&o0NveZSzjLNz;TAOtAOAccjtCkDM?X zrzh#1W9!bnGxPob{qBGNd;dEcP1_BteG%T#BZLB6#L-BlQ5p@k%;BIY$^!rC@E_ld zutUc)kZNT6)9qTQ+OJ-)vmKPp@KuW4VqeSF*zUELTA#56%wM)1GIC}&93UHEfG|K9 zAPoF{GZ4MOX3WXU)7*AzgzNV6kr3zRL*l3X>UArYx36%tFTZN_3Wr$cD7QFZtEq7; zcZ7sNM;{+@^!WU|!x!}U++4&L^g1HFeEjVGFz-PF9hAXPhfb`;*P}g`&jVLQcrQTo zg?s@nbhBd{e{(fHAto)t?}&&60K)scL3D#D(1OXHBXN6eOy$st17FLdb)(?+$T6cP}5Vc5qS+-2&*jBO=x|Bm|rh z!QuA>aW#WsU&JQ_9sCa7-5=pSQ&zAVkrv;8Jb7&Zs|!`AOT02BOU6#(;wi3-ogTZm z2b2p6aat1>BGGU$);ZW8=myHS_6Pb<Z7;D|8YmD61t5n`47*ljaKTPEt8)SPk3 zQHzRK^#^_1`+2GHs~w3et8q1gLbU@`0xY1!YRGBH&`@BN{EqN;Kg_pWTz^Esc^5Rk zs~$dVUraX{bBc;IcWlRog6MH|`yyd!$DCv|Qt>pVgN{#cH>j5+8wu(~e!oNL!CeMj z?GPIWtXbY?oN+R^5b9_SA$B5Y`Zu*Uq%7#*==KZU$+m_qCd@+z0bdY5t*1ZeMqMS0 z>fjaz_z>siYhy=6xA3<%Bru(Nvms|i(R8LGnSiTa+Oak z@-X`WaQf(9m!vW(NvM1se{;)=;7dGu3PAjVEaW`Y@TTB zN$4vAB=T|Y;j9XLI&|;rg>cvh!zvoK!i(djj@~E^lff2qLVR#41uuk#81-bcl=ydV^1v`NN7xcu&D>Mv?sNnli=R|TuCju8b z%Aiw*p^HwyjDyCGB99fvuk8a)5Nd$)h}|7EiVShEm!AYDjz*v3B7DRb;2r1`tPD1+ zK`!iYhj=&$V`)X1?xV~AwU(k99`T&$z-cs0X;X+Oe!Q z$w(#~OxVj>aJVqrk=qI4PDTgHiV~1>6I-xbWFdZ5Y_4Oeu!EKw_96v;WFrg^1_%R$ z0m1-bfG|K9APf)&2m^!x!T@35v&Vo=qtzHu{I6w4DE9y0gKUHW!T@1_FhCd}3=jqg z1B3y>0AYYIKo}qld@32RX?1hh_%}T2(lomKyx4aDnrrJDYwOUL{I%?B6#E+c{-?5n zL}|hRVSq3|7$6J~1_%R$0m1-bfG|K9APf)&BnE7{Ihyz~fYzeRH^){1^l%3H{y%K# z>_z4=#>=GHkJ#OIyX}ar%ciw{-@4B7wq?Ji)%*wZ?dEdRGo~I>rtzfFZFt|X&rqp9 ztB>jn=u`9{ouT`d&Z+&mwpW{`d04ZaxoqR7m4(j}z z;O-Cb!ARKI7ZSW7E)al+B+-7LFCzH+(eqprazYU1MJQUz-~+x8|vc^I}xl4?*edFIRLLn6<9_}Hy7qzQXV99l!1hL6(u13 z5#6OK)9rHkg1(5$1=zk4!0t>{XLuYk(C6nut}yTC-4Ovy2A;n42mz@8MD-PeD1HV+ z;f9!`0+Kj?h9pXrgJ-{b<#p_y2NGA!KpkaGwq`EKn47AMUfvJSe20OhHV5EUT|Qof z4}&*};W$ywl>l`Y0qAu~LEW6+-_5zVg+*8?3w8_7nTO!0>Vs!YaWM$$ z$_GIyJ2?q4?8JRb%7X;|Y>=>7#j5ncgPReT;CAPFIKz2=D}bK=Aa# zW4A#e2#-&Oe0}%#zcP$_*LsH)w(ey$7^zvjDCF$2S!n zVi|z8n*p>YRqGIy^$T9FNN}4;JyqeS5qf;6)jEvo@hVLs0C&9sa8tHkkph|Zw(Hek z77r=7JBv=ajaGwMjLP--VxGQTr<$rV0=}RH)S}Bbtx2;nvF0l^0GFW@4p~uzhnPIA zN(-MF_C6G!v+RA^U(h?)4CVyWYdT}f()u)yXw#`t>Y&kMxXWC|DDN;V3lqzI=WnQ`87=L;8JKi zDgLk8L`PfH#FxeYl{eDS#+0$ePLve?S8#N6MasZ3QgCOI^3eawJLzb>ni691f0^2J z%lcO7Ryx|5vd)u&@2TSdlI?Vqp9xX8AttE+4N^Q)5~q#-=ez0Xs+p*xtjXrBrK59G zmXQe6=eE<)Dg}6=dz~i!cX;UNb&5fyQFTISJ|X@uQs^s0&8HIrhEu4J;z=V;Z0bCx zmyWh6U?bTfOK~PFE=&{u7pUUoWP)(tlJY=Jet?c{RUT|LrPTwkXUci~nt^G-SNrB;tRoLdQ219by@b3zqeCv?gWi5S2~D z|7I0cmEwPsio8mb2%x~IiuH;V$gJ0(3bS}f!QEMmg7vB}i%}EtKdqjsG6HmKX!tR} z%MWMKYz!i3MwRBC!XJ{~|5Lg?L9QbV5C#YXgaN_;VSq3|7$6J~1_%R$0m8tif&raI zN7>Qy{~Go~3jWAO7$6J~1_%R$0m1-bfG|K9APf)&2m^!x!oX*c0iDinM)&`hSe~KS zKd?`+r`YeXcd~og4eT}S6>K%jFn?iQX3jBBFkfbNF&<_;^LeI$DPi>X^Y%a6&)J`} zAGbef|B`*P-D%g^Ms0twy=;5l_LS}Wwy)WC+Sc2a*~)FKO=tb9^|#iCtoK-Nw}z~n ztZS`J)*P$O@`2?|5KK100AYYIKo}ql5C#YXgaN`pDh%WBb?%&e* zSkLHa8`ky6SZJoOx!+w? zdf3yj&;LukWg#5NUifw2Kd*izs|oh2o3@VLxpwzm_wHJ~X6d!xUHWVDtOZb9y5Qj0 z_R&*!ts1{))#$0>v_>eK*LZMj(WTSxox1D3qJ0?+P*UFT#)a_M3m*pl`qZtL#=V-# zdMNCuAG&|{y}N$>(y8P7USGL?_2`A&Z;d~`)HU|u_~R!x{pP9n554#1cu_`f9h}=x zH}u;}hfbf~RowFUD_Qs}Jw@0Z%XK7LO_i=`Iw*|paO8xIxa)WCjS z%{ph-o-vmC(DLxG#k2E7Hm}XDh6AvC>1P}Jz#&z`HK(z*)!ax$u*yu1o7ps3eQ zKBpf!d*$Xw=2t@DHI+kIAMlsPb=1gXC!QZadG^wu4vk%^rKmmB^Dmq~IgnWaCskGq z9ehz(dh^RKzkjh%xTVONQx2td|Nd^MxC{#;v-YY4sr5G~V#q0EY zsV?JubjJLH%Z>BU&b+gl%DJ%9G4~DqH^*JtcS@-9mR+|n)YumEX!IETj}(daXoNeZ z)ebm!qvKd>e)+{aFBKJCx%QfOUp9`s8a`U=y4(8-z3R|k@BOz`3-`U$y8qReR+SdP z*=vgaOZ(QV{QXZnzWQMJ-SC>meTTpM#1FksyN1cKTHn8i!$7nqFg$ zQfX(7{ouiI>Ye7bk5U)LR~C)E`P|s9=Z43AyY!-W{H?>6ez^PmQ;@l|d-%lRhu`)- z^y00ik&&KK%ayn3SQgM6Y$$8*n2{PD?_doK6fn#&iyeuv_xuxed}+Q7_7~ZPvUMTq z>0f~=BO8rW-=iODp1*UmroajZ7h5Z#!8H_RqS7b_rCsK5IF>CpnM|w&jyf!rts|S> zd7xm=fW~Zwf=qLzrpwHl&@t1H?buDb^uz!9Q%$)M3f37zRk^0LJrq4mJ>-g3=}h-^ z{$%&PRk|OISQ;GV-}~l+x(Dge>P$WCSL)aOhiwmalrm|4ui3L(Z*$O4 zdKG=Ff5UhYHFEgLv;RPYhH0^XAT`aPgEJW2u{15D^R$(Fsgfhq)3%{_{Ur+Ea1y)ba6q_Pn+9l_GTie+jdNVlT4qv43V?VP9mQXP;qz$^MKz z&igJ7}|1_%R$0m1-b zfG|K9APf)&2m^!x!oX*M0TkJ4mMp+oaE=@WcF;F0AYYIKo}ql5C#YXgaN_;VSq3|7?_#? zbpM~k|5I~AN(lpm0m1-bfG|K9APf)&2m^!x!T@1_Fz{Jn06qUt;{VS|Ym#~p1_%R$ z0m1-bfG|K9APf)&2m^!x!T@1lY6dLyX)2pKsG&1;b-HhBe?jkHGnf-hujz~_OY74- zqD`kpse}4oX*!G^!(E2+hAL~D`H$v1m~{J5O^4lM9kMgFqc)w5Ggnxa=x@`{VUN(K zS)XxO|2va+YMCFe)I~>^D|lgI@(VttSKEi3{D2U9J*kI>7hP7V^7dL8K6=qT?+bB& zU8&038fDl>bHJ^sQ03LwGGP2NO(_o|%GFjK(E_+-YSS%yw`QsMcIowDJ`{HDIIPhr&{=%heGJEma;^lZ!$U`}Qw%D-Rab%@S{{q)`rz^w<1Tgl1;OY1CuAmTfh20@vUxe%Sqj5c3 zl^0lJ<^2M5FPB(@EbEZAGG$KU?=c5s2}i&Kgd;vTa$lJn>1cDx=;G_0L;N5Y^0;sY zpffhm(I%CkxAuFyyo5Xpek3a=W#kwo!ZtqWatr<78$+OVx+>FO!30U+1E_aBJ z4{tX3buyz$=n*_Y1J?ywTUB}WyGVzqHGo=FdG)&p>KD9Tkzlims=9?>j}Nt)NtF&F zswzz)0N1FB^@_MM>oustEcVKHG#8z+UKM6BD%aRg4^~q974(Bi1GQLl z?MgbDVRl}2$P=42f)$}XWb&X_iLQxazrkM3yvFxt-|_@HEhkX zoUm-Lq??bMtj0%;z2<;9*L2#LVK`>kY@qZH=vUK!p}$5i*8NfUHC>Zw^T&(s$ytPf zPb~uqXKP~T6MZAvhe`2v&B&ymGNZH#rG6F{uP? zTe|4z)r#GuOxxtiSy-B#ElR2FWI-6_B#{T*o43%>rHXY=K#A)GlY-WT3I{upaA}P% zEW!(5d!TdEO?0$XvCb1g6HBak#bChK?_=xBusaPkP(nUd>CsF;{Z)r#w~pGefNa6M7%9I|x~7|WxeLRFURVk7;;r41G1ZR``0SCSMC_u^C} zerF^;M?p74KnB>ZRYk2QPkU?%O?fJ<#WGNNhN5|eO$^m#qs=zW*aFikCp&qBk%uj%mZmDk2JFS$of%RE?_IR zds!n(@L)OD=KyF)s`TI$f!KP5FB=d$Kau$57JPI($45`w#b%!zR<;I34M zjpQX}X@z+`10qs3Kv8vBPw%t?tb%bO+5L)aC&isN*#KB2%NB*nZ)kK_0X^l0jQ5+w zdPJKAa1}Vdspt^Pz@ppD09uo(b%@F)?*F%$)KgUo>^qF=@hVLsfP(b~z)jhDMG9or z+pbrGSv;iR?kui=w$W-Zi%}D+*6ljgRFx6%1udW!UA}2envIDyU#S7O45e_$iV!fg zH;Vt$Ee;Ca@^=k$p7}P@W*@VE$9|3NP21hJX6p;qTdXCPCoQ+YDP$uI5C#YXgaN_; zVSq3|82CG2U<16BTFJ_UIB$)|ZBu^N?fO=LR z1n8wobx%M|^If;LCV*EE-HVz__p_3}>$YhDfVL{tc_L`yb|U_++or~u0mt_$I~xF8 z!P?K%bjWU0w$-c2bL>{%7wQ01aj1)=6Q;utni->|FMq=Hkeoj zQd+S+-lmmm=vn08x+%NyUn%Y?$1<+ozSmJ^_=<|u;px_dN^0#Ad$NV{+^ZCv>mo1sk%lr5$D-WOgdVUTcls5B& zVvz6m`b6FvaC==s*ctac;(mei7V@HF#=K@j|u@6P;N8k&d8)Guu5@DXrm9tTTWppEaAKpD&N@*U-loR zn@zdJ#qAbqzC7DpHCdX(- zV~6~r5E=@`8j`=(NXD*6w30C&i5`LC)SHaCa|*^QPV_uyb9BQTF)^mx>})zBdjx8i z(7z$65NJ=r1x9oP-(3AcAw0m3)*bBz#VZZQ+*#S0@(uZzeLpn% z=wF|tGTM_+d9iqP{U*T&t;wL^4T8%;8!Yz#JUY=Wz{x_NG!(}6i?Y}}QQwo$R|ZJr zW8A}8)%bL9?<=KH$P3*n>b62NV|)3^0U?B1p-?5?CrLh02*`Z{T(K|&GQA-l>oB@k zYENK=J-W3$# zAoQixmHNMC2B@~2&FC9N$!0W^ixktQ+}c_?G8A=2xt(!_pi}n?F}0&vU6PTEIGC`P zWzgWlXh&%$h&u%xC@V`qjhkG9eKHHNv!Y`iYqbNkq`C7H{E-h~fG|K9APf)&2m^!x z!T@1_FhCd}3=jqg1K&Re^t6sPq4;0Ny-jf+!x#Aw1_%R$0m1-bfG|K9APf)&2m^!x z!T@1_F!0@Ez^>Cz=3+nbXyRynL2mR%0BsHb)X;=}0AYYI5M#iupG?PQ0d!VO6J~8bwwHrS&?ljIYoHT4UR4^|v z5|gbzq+g{s>+aGu(NELMsSDK4W70K9SJX08q;~Gi-mVZl7SbEuDh~L1=Si+1zZeLI zdIy7&Cn)&+@YE#Q_78?7-w=Aht2Zo$;JDA{hl^m->62W7&k1>WT4Z?*fL3Y%Dhzs~ zkn5`fa+M0mWUb-hJpZ6i2s%TePjrPPC^9??>X!WR9iaZkDv;%w2w8nXNOX?U0g0@f zD3S5v!LvvM%2K*2K;+_yD5cZs4S2&&CzNYhIS46DS;&CsgQt)~0N7aufb}W>WqXYR z$tIgXb(aF_6{=Ce@_l`RYeUo(;#uGUo_!9&Bhyjaiu|crqR0DH72AnQw2--hn)oZf=Jawmuf}`O^C=9J? zDC~8iPC>{7_|)BAQLWhR<9QI#lK~<$wH%i)h%9GKx+c8Gbq)0@9KI(_6RqWb0ds;j zLF+yafET80?=o<#0YvW>kpl7Fp_>I6sW$?$3JI-j5%Xmd6i0Hh@pv24h}?2Ex#hg5I5009Mm4jKzq&4^X=-fZCX{jmYZy zB#%cXxzntnvhYrXelM!J9+QUL#s`1^+-(Hl)U8-%L1D#R22I$-T_pBwx##U`kfI9RF!*b)^MHfsDR7yfb@@Um6oA!7=Q5R|9Y>foE^cB25D<95>>dJC7v z9${Uk*~Zh>T8m^kVeT+}YT9j_Y1m@;y=}hrM(f+mab}RUIUaHJ+U{@|?1yZBv3Fb2 z%#SiT`iHo~`qk7hb&&HIcNr=XtX7{+CQ?hQE;6oX#AZ_bU)#Y%DzyMjSWwLdCbCKm zWO5TVD)t^L{;v))5zmCk!j3jd2WTa#CQRhG@qeY8i7cLwQYt*O;z}k`n!1ogq+Z^| zMC#Q5$}T@n{9h(Akt@`r#@l;&ULl?x5&xI2Vj|PDmjdQ9WepRVp$0kGJLI9#i1@!` zH4|ym3OhOe=ha%rsP;QLEG_P1A`4R2H6GoDda=jG^FVdcRZOH^QxUP@FOHNZ_cD=L zn!+o*u`o5~87KZPP}|xk8z1gUqVfNv9wyS9Itv(58UN?6WFoDp<0^u!X#AhIj)~0G z6nI4ZpF4p9aQvU6jp0RuixKgEwl+=^i(E#=|5@6&y`oyN+sE_Z_L-VnPBi|{(1!Q8 zu1WEKx;9$N{Q~AB+R&!Cn8?D^?Og_rHGm`He@>f5K$f9s1zEM-aB_Pk8)H)Z@6giB zWob@D{BPH$He+6ti2rSB`-QO>k@tbiSi6}>W9l{{tDA`bEm|rYkN?eDavL820t+T> ztXO72VZ}ym*u`BW_H5V!bq(6Ei(wPd8>6AJ3J&zkm`KUrWML3Bv)@M3Mc)6X)%*m( zAq)@(2m^!x!T@1_FhCd}3=jqg1B3y>z&|7ddRk98(DVN^_Z0 z0AYYIKo}ql5C#YXgaN|9_mBa--eE!a|J!UUDDG|UW$qYvH}`Mc_1sEs0e1;k2TuZg z$)0B4Wq-%s$?j$c*k$Z&b}C!UnjC*}oOB#<+~>H0AYYIKo}ql5C#YX z-xmh*+8El>(zf)qpS`xYXV+D&PxqfY*ScrJiPv2M@(`iVVDPt@xz(;$;I?ZmSE z>ocapX2sM++qQ0PfAhEhecJnfj};ZR!oHcUfe-I_^Y)q7?)%*(LhJtbedk_3boP+p z@cFzJINH*((|zVQ$KKiVj_b_Aiv8}ETYb+PY|U^auX)jd87+JC*-enCYg&Kq=4IP% z+k4%TrR|Fjw7+D@XoTGrjkjOec<$J33x{_vJa?=ty#e->HJm(q?2ii$v>zz@dN)&A z54$d}Klt1BrR_WSe6{U%&M_zb>;5PO7RqxC&Pvtb7g& zpQo|4u(PPP=E-+f+%mkoWrno|@;Nn&0EpVO z-WTsV{+lJYht7nSw%)q`!0|^tFaNRZ{j%297mrVE)t66!^Oj5*|GZmQRlHEAca*?s zg(V$7o-<6JYIo%EkZs{tH^QA&hT8twFhxBxyngNQ=~}iJjuaIiY%%}h&9`)~nAjqe zFX||td+IFxlP8TilOf+Qxx=QX7+=;`wfgPXQHNJN+hZ(*{cK@|qX5#vf`j>C>M3ew zP0=LCUOK7cevgs*a2Rue_@K}e-B1ZkMm`*`$Uky#9z~~JIB%d1KVZtugMCeT9gXX@ zU*M>(tPdWwy0<=)BkFQ;;Xp%fNA5}^MV_ zLk1kUD&u6$sa^EfNZW^`4LkbiXa2Nn%lrqZ;U7M&TlwGoNyB@#l*`MkMjp=Yc zn|`Jx4buL!6E(-ran!ayhiyL@-gfrwUtV|mVP!I(f=>hQ|;My*sim8Ja?3y-8f8B+c!+6j=rgz zWP{z)Z8g9l*qND1qj*X;o9AcGF`HQ{94oTc%sRE^lluy1KWVW*KFd-=uY>pE!6t7$ zXg|_>o#E(De??cBU`K~3Ses`~+fFe@sRx{qTD|%1o@cl1t<^tv%G$zL-FxT#`uhx4 zBb=FI%=m)ZPG=clyT;J*uzfrA1ZAe*qPJgTu$M5fcOG+OXw`5rb!z{MCw~ZT2Aat_ zQlKGi)Wa#PKA5J1bc(KK2UY$H>XmEUbLh!5W3gO%Zl9rU#U5(UcHid@wf9$_Klaq{ z?(HA7zgLXz|F^LbiaXDpBh$^SMj8Hm-@Q=B983To&g5ria-t*iYFH+5ctVU|(ZjVvn&$;jMy? zu=lZdvv;sN*_+uLK^*xI1_%R$0m1-bfG|K9APf)&2m^!x!obDG0E&L;wlS~-+RpF$v5+@ZE zI4Lj3Nm&_AN=tDvWeQG8N^rvSI4Lg1Nl_6_CQrsmVIfWm3UD%M5>E2-agvvZliXaK z8Gd;>6*=iQSG9n++#cJbyr2EI4O2`=chk@4n&}R`1>X91gmsx_8&6wn zEt2Jgxx@6SX}58vVTzx2itRPvl zfr+eAqi6CGrQ#M5EqfYthD4v}3QNIIZ%_!hC4YPetbth_WFnpkk%br99Hj$@teP;9 z@!|%kHe0T0?1R$k5Y%#=$}lpyi{^f+$;p$ zPMiVklGRM4O)Kp6LvD{4hvwBZ#OA5M5kc z+8_p;E@=qbKV;R}+Bi)V91S-@VQ5uDVXq5y3R&8?y@Gx0_VGN3$kg0&T*4r-oD6Lo zKCWx1SK;vK+Gs8J3z(B=Lz`y1gL0>CEi!Pd0d(PxGTu9I+B5>P3blFu z*e?bBLOARVc%UvFTAF#(`&~gXHh!>cQ=2g_LIYuF*gCfNeYnzT~ZB?bDusOrpGavLA10kBCME0$SMSg}zXc5xSpJzMU14P8v6K6Q36 zY$AGNG*ni>fgS?kJ!+I4HJ*_RkKXaP&@y*sZ&wJSjo$E9alqF*PjU_U#XvaJ3ziLo zdA}d*5N$)>EF}fQ$e0i^vVA}*VQ!_kpK%M}y?wjc%N!p#{@pRl{)T;rz07tD-qvrj zK4TrQW?PP02H_Ql&zr9{mzth4c}(fXhmD%aHP*4{bX~OlM3}+ZBxhPRwP1zEcGzGG9Ht z3VoAD!=ds@E@pr<{D6h`S=G)y{ac;g>U!Y)yv2A``hnJ&o}&YQ%*1w>Pb> zw#iM74x;_rw>&4NwgVkgcYN0jJOq5#Ho=n|YE619z2$wN^hC`?sJ){zma@rHNm|nI z9iHV-J8DfPSRP)D7zs|~jV<*1|61zj6!$pS&vEdk{yx^~c-XPV!PpO6J~8bwwHrS&?ljIYoHT4UR4^|v5|gbzq+g{s>+aGu(NELM zsS5zB-DgEDLq%#;PkD^p(=D$7&`J$J6EjHbs{wMA3dm$b7{*k&7pyl!_U;W#u5GG-V-)JN2Dq09daAP!1%<8S%JF0rd*i zsNZm)vT_PYnWniEFrU>WfILG5a_p6rh7!w)0a}f9$d#!G1CZtb zC=>zUf|PZQhbp1oQIYK0$snR#OA(1k97?dJ5a4HN39sq{=;|zhy%O?S< zn$|wq_;6Pejla6{0kSz|7BHkT1nbEI+}4zF6%z;1nFwbt0MFGDc*JCcGiTxi;Q5F( z*&rZw!y7x+S(^pmsXL8gqA=WzRbL(vMZbO5fJZaBHUl8rGb zXj^LsWUb6xmgZF4|L?H@eCjqB^PL(|G(3$ zp|bJFxW}X+xAE~Hu+VJ;;MA>HWHxMxg@uh8Kgxx_oCdsX)p&~g|0w>q+n%Smhq%@3SL{Caa=8D0lcU;x!hW^A0uGQ5 zVSq3|7$6J~1_%R$0m1-bfH3e6$iV70hzZn8t;;b+EZ83NcXP|91Gbu(({ZrTJk>OS zR5M2)L&}pSXnW-E=DMZ=u$pLo6ma4m0O-4}6=3J94!IQiCeI=zemA$K1th4M*pwBI z&vcLZySb~H0ep^X^+$q_y4%p%1o&!ZW@P0Rvx2xzqqEh$jez Tuple[List[Prompt], Li i = 0 while len(candidates) > k and i < self.max_n_blocks_eval: # new_scores shape: (n_candidates, n_samples) - new_scores: List[float] = self.task.evaluate(candidates, self.predictor, return_agg_scores=False) + new_scores = self.task.evaluate(candidates, self.predictor, return_agg_scores=False) # subtract length penalty prompt_lengths = np.array([self.token_counter(c.construct_prompt()) for c in candidates]) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 47bccb90..4bd05272 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -12,15 +12,15 @@ from promptolution.utils.prompt import Prompt from promptolution.utils.token_counter import get_token_counter - if TYPE_CHECKING: # pragma: no cover from promptolution.predictors.base_predictor import BasePredictor from promptolution.utils.config import ExperimentConfig -logger = get_logger(__name__) TaskType = Literal["classification", "reward", "judge"] -EvalStrategy = Literal["full", "subsample", "sequential_block", "random_block"] +EvalStrategy = Literal["full", "subsample", "sequential_block", "random_block", "evaluated"] + +logger = get_logger(__name__) class BaseTask(ABC): @@ -49,34 +49,34 @@ def __init__( seed (int): Random seed for reproducibility. config (ExperimentConfig, optional): Configuration for the task, overriding defaults. """ - self.df = df - self.x_column = x_column - self.y_column = y_column - self.task_description = task_description - self.n_subsamples = n_subsamples - self.eval_strategy = eval_strategy - self.seed = seed + self.df: pd.DataFrame = df + self.x_column: str = x_column + self.y_column: Optional[str] = y_column + 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] = df[self.x_column].values.astype(str).tolist() - self.has_y = y_column is not None + self.has_y: bool = y_column is not None if self.has_y and y_column is not None: self.ys: List[str] = 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 = 0 - self.n_blocks = len(self.xs) // self.n_subsamples if self.n_subsamples > 0 else 1 + 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): generating sequence per datapoint - def subsample(self, eval_strategy: "EvalStrategy" = None) -> Tuple[List[str], List[str]]: + def subsample(self, eval_strategy: Optional["EvalStrategy"] = None) -> Tuple[List[str], List[str]]: """Subsample the dataset based on the specified parameters. Args: @@ -123,7 +123,7 @@ def _prepare_batch( keys_to_predict = [] for prompt in prompts: for x, y in zip(xs, ys): - cache_key = (prompt.construct_prompt(), x, str(y)) + cache_key = (str(prompt), x, str(y)) if cache_key not in self.eval_cache: keys_to_predict.append(cache_key) return keys_to_predict @@ -379,7 +379,6 @@ def evaluate( preds_for_prompt = preds[start:end] output_token_counts = [float(token_counter(p)) for p in preds_for_prompt] - # Per-datapoint input tokens: prompt tokens + tokens of each x prompt_input_tokens = [prompt_tokens + input_toks for input_toks in input_token_counts] per_prompt_inputs.append(prompt_input_tokens) per_prompt_outputs.append(output_token_counts) diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py index a12e2816..8650d2b8 100644 --- a/tests/optimizers/test_capoeira.py +++ b/tests/optimizers/test_capoeira.py @@ -44,12 +44,11 @@ def test_capoeira_selection_prefers_better_score(mock_meta_llm, mock_predictor, predictor=mock_predictor, task=mock_task, meta_llm=mock_meta_llm, - initial_prompts=["short", "longer prompt"], + initial_prompts=["short"], df_few_shots=mock_df, ) - optimizer.token_counter = lambda _: 1 candidates = [Prompt("short"), Prompt("longer prompt")] - optimizer.task.evaluate = MagicMock(return_value=[0.4, 0.9]) + optimizer.task.evaluate = MagicMock(return_value=([0.1, 0.9], [len("short"), len("longer prompt")], [5, 5])) # second candidate is better objectives = optimizer._evaluate_candidates(candidates) selected, _ = optimizer._select_population(candidates, objectives) From cffa7893a19d8ef0e9d9b2920726b8947898da4e Mon Sep 17 00:00:00 2001 From: finitearth Date: Mon, 29 Dec 2025 14:19:20 +0100 Subject: [PATCH 11/53] implemented evalresults class --- .coverage | Bin 69632 -> 69632 bytes .vscode/settings.json | 2 +- .../random_search_selector.py | 8 +- .../exemplar_selectors/random_selector.py | 8 +- promptolution/helpers.py | 6 +- promptolution/optimizers/capo.py | 20 +- promptolution/optimizers/capoeira.py | 16 +- promptolution/optimizers/evoprompt_de.py | 6 +- promptolution/optimizers/evoprompt_ga.py | 10 +- promptolution/optimizers/opro.py | 6 +- promptolution/predictors/base_predictor.py | 10 +- promptolution/tasks/base_task.py | 359 ++++++------------ promptolution/tasks/classification_tasks.py | 8 +- promptolution/tasks/judge_tasks.py | 5 +- promptolution/tasks/reward_tasks.py | 5 +- promptolution/utils/capo_utils.py | 1 - promptolution/utils/prompt.py | 41 +- promptolution/utils/token_counter.py | 6 +- tests/helpers/test_helpers.py | 27 +- tests/optimizers/test_capoeira.py | 16 +- tests/predictors/test_base_predictor.py | 4 +- tests/predictors/test_predictors.py | 18 +- tests/tasks/test_classifications_tasks.py | 47 +-- tests/tasks/test_judge_task.py | 24 +- tests/tasks/test_reward_tasks.py | 8 +- tests/utils/test_prompt.py | 13 + 26 files changed, 317 insertions(+), 357 deletions(-) diff --git a/.coverage b/.coverage index 4f86ff21c29ceade31ad542efcb6270a1e4fdbd2..35124674714a2ecbc725f89aa85e571a7946542f 100644 GIT binary patch delta 676 zcmZozz|ydQWy1;&4zmiw?26QqvdtSje()plBEr<`Z!+*dPVr6W>*g!uOX3UWv*$D76XIj${l)u=_ZII--krRYc(ZtIc?Een zc>eQz=6S|*muDN#T%LBGBpx3g6CPm~yTY!(v;;hKE0OUjqWnuU>5%GxX4%5%F=TE z;2X0ZU%>j$^)??+8}+}LHY>{3@`Ry=RUBtcGq4eWi6nnht?|5AE|%; zGc!EcyYI&T{kLl0y}bXPiGhKm;cV{zY<3Y%kXFrAv+n2H@BH-XS8f0Hz5T4>8bDD4 zja3P$3}2W(tYs+R%U{OES9#NjnvjTU={rw+o%aG1E`FxMJ6CVpO z)c!CsFdVQ4g&zY09|HrEwgON?gu*UsX^r3gf92)v{5?KBWq`oZ;@ z>mmaPY!(yn;hKE2OUjGenuU>5!urxoze~EzoK`@Nh}9+k`B_4iK&G0dSAA8L{oTvg z-e$huA}3=36tl9(d$s)TEUVgiS8NiE=ReQ+b$$Jp^{m@}Pj=|`)%N^#Tt&bbs7%Rt z)<28CC%zn?vP6Eeun|y1!|2t|;(vGkT-M**lmN@16^w8roLzw+{S{vMy6vOjoUyi+B9 zTKY2<&v~Ec%=R>!{KR_t^CzB@vwCY-m=(%-C%@`l>*OH7%E;NsAaH5dhBf~gKA8J& zXZZQwUH;$e>;LMmANY6E{&)KSH}n5)uK#u($oX4;VBh?8@6*@bujP2~AWe3--FSycgp)T!C~ZFR;sHWr^M)S*;s9EU B6vO}k diff --git a/.vscode/settings.json b/.vscode/settings.json index 9b388533..a3a18383 100644 --- a/.vscode/settings.json +++ b/.vscode/settings.json @@ -4,4 +4,4 @@ ], "python.testing.unittestEnabled": false, "python.testing.pytestEnabled": true -} \ No newline at end of file +} diff --git a/promptolution/exemplar_selectors/random_search_selector.py b/promptolution/exemplar_selectors/random_search_selector.py index b8cb6ee2..1eef61b4 100644 --- a/promptolution/exemplar_selectors/random_search_selector.py +++ b/promptolution/exemplar_selectors/random_search_selector.py @@ -28,12 +28,12 @@ def select_exemplars(self, prompt: Prompt, n_trials: int = 5) -> Prompt: best_prompt = prompt for _ in range(n_trials): - _, seq = self.task.evaluate( - prompt, self.predictor, eval_strategy="subsample", return_seq=True, return_agg_scores=False - ) + result = self.task.evaluate(prompt, self.predictor, eval_strategy="subsample") + seq = result.sequences prompt_with_examples = Prompt(prompt.instruction, [seq[0][0]]) # evaluate prompts as few shot prompt - score = self.task.evaluate(prompt_with_examples, self.predictor, eval_strategy="subsample")[0] + result = self.task.evaluate(prompt_with_examples, self.predictor, eval_strategy="subsample") + score = float(result.agg_scores[0]) if score > best_score: best_score = score best_prompt = prompt_with_examples diff --git a/promptolution/exemplar_selectors/random_selector.py b/promptolution/exemplar_selectors/random_selector.py index 7b0ae0ff..63b02955 100644 --- a/promptolution/exemplar_selectors/random_selector.py +++ b/promptolution/exemplar_selectors/random_selector.py @@ -53,10 +53,10 @@ def select_exemplars(self, prompt: Prompt, n_examples: int = 5) -> Prompt: """ examples: List[str] = [] while len(examples) < n_examples: - scores, seqs = self.task.evaluate( - prompt, self.predictor, eval_strategy="subsample", return_seq=True, return_agg_scores=False - ) - score = np.mean(scores) + result = self.task.evaluate(prompt, self.predictor, eval_strategy="subsample") + scores = result.scores + seqs = result.sequences + score = float(np.mean(scores)) seq = seqs[0][0] if score == self.desired_score: examples.append(seq) diff --git a/promptolution/helpers.py b/promptolution/helpers.py index 94879df3..c7a3278c 100644 --- a/promptolution/helpers.py +++ b/promptolution/helpers.py @@ -124,10 +124,12 @@ def run_evaluation( logger.warning("📊 Starting evaluation...") if isinstance(prompts[0], str): str_prompts = cast(List[str], prompts) - prompts = [Prompt(p) for p in str_prompts] + prompt_objs = [Prompt(p) for p in str_prompts] else: str_prompts = [p.construct_prompt() for p in cast(List[Prompt], prompts)] - scores = task.evaluate(prompts, predictor, eval_strategy="full") + prompt_objs = cast(List[Prompt], prompts) + results = task.evaluate(prompt_objs, predictor, eval_strategy="full") + scores = results.agg_scores.tolist() df = pd.DataFrame(dict(prompt=str_prompts, score=scores)) df = df.sort_values("score", ascending=False, ignore_index=True) diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index 47713a06..97f1f659 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -153,21 +153,20 @@ def _do_racing(self, candidates: List[Prompt], k: int) -> Tuple[List[Prompt], Li List[Prompt]: List of surviving prompts after racing. """ self.task.reset_block_idx() - block_scores: List[List[float]] = [] + block_scores: List[np.ndarray] = [] i = 0 while len(candidates) > k and i < self.max_n_blocks_eval: # new_scores shape: (n_candidates, n_samples) - new_scores = self.task.evaluate(candidates, self.predictor, return_agg_scores=False) + results = self.task.evaluate(candidates, self.predictor) + new_scores = results.scores # subtract length penalty prompt_lengths = np.array([self.token_counter(c.construct_prompt()) for c in candidates]) rel_prompt_lengths = prompt_lengths / self.max_prompt_length - penalized_new_scores = np.array(new_scores) - self.length_penalty * rel_prompt_lengths[:, None] + penalized_new_scores = new_scores - self.length_penalty * rel_prompt_lengths[:, None] - new_scores = penalized_new_scores.tolist() - - block_scores.append(new_scores) + block_scores.append(penalized_new_scores) scores = np.concatenate(block_scores, axis=1) # boolean matrix C_ij indicating if candidate j is better than candidate i @@ -183,7 +182,8 @@ def _do_racing(self, candidates: List[Prompt], k: int) -> Tuple[List[Prompt], Li i += 1 self.task.increment_block_idx() - avg_scores = self.task.evaluate(candidates, self.predictor, eval_strategy="evaluated") + final_result = self.task.evaluate(candidates, self.predictor, eval_strategy="evaluated") + avg_scores = final_result.scores.tolist() prompts, avg_scores = sort_prompts_by_scores(candidates, avg_scores, top_k=k) return prompts, avg_scores @@ -223,8 +223,8 @@ def _step(self) -> List[Prompt]: @staticmethod def filter_survivors( - candidates: List[Prompt], scores: List[List[float]], mask: Any - ) -> Tuple[List[Prompt], List[List[float]]]: + candidates: List[Prompt], scores: List[np.ndarray], mask: Any + ) -> Tuple[List[Prompt], List[np.ndarray]]: """Filter candidates and scores based on a boolean mask. Args: @@ -241,6 +241,6 @@ def filter_survivors( ), "Each score list must have the same length as candidates." filtered_candidates = [c for c, m in zip(candidates, mask) if m] - filtered_scores = [[s for s, m in zip(score, mask) if m] for score in scores] + filtered_scores = [np.asarray(score)[mask] for score in scores] return filtered_candidates, filtered_scores diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 60ebcc90..9ea6b3e7 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -120,24 +120,20 @@ def _pre_optimization_loop(self) -> None: self.scores = (-selected_vectors[:, 0]).tolist() def _evaluate_candidates(self, candidates: List[Prompt]) -> np.ndarray: - scores, input_tokens, output_tokens = self.task.evaluate( + result = self.task.evaluate( candidates, self.predictor, eval_strategy=self.task.eval_strategy, - return_costs=True, - return_seq=False, - return_agg_scores=True, ) - # TODO move to evaluate method! - input_tokens_array = np.array(input_tokens, dtype=float) - output_tokens_array = np.array(output_tokens, dtype=float) - scores_array = np.array(scores, dtype=float) + scores = result.scores + input_tokens = result.costs.input_tokens + output_tokens = result.costs.output_tokens score_vectors = np.column_stack( [ - -scores_array, - self.cost_per_input_token * input_tokens_array + self.cost_per_output_token * output_tokens_array, + -scores, + self.cost_per_input_token * input_tokens + self.cost_per_output_token * output_tokens, ] ) return score_vectors diff --git a/promptolution/optimizers/evoprompt_de.py b/promptolution/optimizers/evoprompt_de.py index f6e701a8..0ae339da 100644 --- a/promptolution/optimizers/evoprompt_de.py +++ b/promptolution/optimizers/evoprompt_de.py @@ -60,7 +60,8 @@ def __init__( self.prompt_template = self._initialize_meta_template(prompt_template or EVOPROMPT_DE_TEMPLATE_TD) def _pre_optimization_loop(self) -> None: - self.scores = self.task.evaluate(self.prompts, self.predictor, return_agg_scores=True) + result = self.task.evaluate(self.prompts, self.predictor) + self.scores = result.agg_scores.tolist() self.prompts, self.scores = sort_prompts_by_scores(self.prompts, self.scores) def _step(self) -> List[Prompt]: @@ -99,7 +100,8 @@ def _step(self) -> List[Prompt]: child_instructions = extract_from_tag(child_instructions, "", "") child_prompts = [Prompt(p) for p in child_instructions] - child_scores = self.task.evaluate(child_prompts, self.predictor, return_agg_scores=True) + child_result = self.task.evaluate(child_prompts, self.predictor) + child_scores = child_result.agg_scores.tolist() for i in range(len(self.prompts)): if child_scores[i] > self.scores[i]: diff --git a/promptolution/optimizers/evoprompt_ga.py b/promptolution/optimizers/evoprompt_ga.py index 9a0b4e39..ae8dfb4f 100644 --- a/promptolution/optimizers/evoprompt_ga.py +++ b/promptolution/optimizers/evoprompt_ga.py @@ -68,17 +68,19 @@ def __init__( assert self.selection_mode in ["random", "wheel", "tour"], "Invalid selection mode." def _pre_optimization_loop(self) -> None: - self.scores = self.task.evaluate(self.prompts, self.predictor, return_agg_scores=True) + result = self.task.evaluate(self.prompts, self.predictor) + self.scores = result.agg_scores self.prompts, self.scores = sort_prompts_by_scores(self.prompts, self.scores) def _step(self) -> List[Prompt]: new_prompts = self._crossover(self.prompts, self.scores) - new_scores = self.task.evaluate(new_prompts, self.predictor, return_agg_scores=True) + new_result = self.task.evaluate(new_prompts, self.predictor) + new_scores = new_result.agg_scores prompts = self.prompts + new_prompts - scores = self.scores + new_scores + combined_scores = np.concatenate([np.asarray(self.scores), np.asarray(new_scores)], axis=0) - self.prompts, self.scores = sort_prompts_by_scores(prompts, scores, top_k=len(self.prompts)) + self.prompts, self.scores = sort_prompts_by_scores(prompts, combined_scores, top_k=len(self.prompts)) return self.prompts diff --git a/promptolution/optimizers/opro.py b/promptolution/optimizers/opro.py index e7b9048f..2e613ace 100644 --- a/promptolution/optimizers/opro.py +++ b/promptolution/optimizers/opro.py @@ -105,7 +105,8 @@ def _add_prompt_and_score(self, prompt: Prompt, score: float) -> None: self.prompts, self.scores = sort_prompts_by_scores(self.prompts, self.scores, top_k=self.max_num_instructions) def _pre_optimization_loop(self): - self.scores = self.task.evaluate(self.prompts, self.predictor) + result = self.task.evaluate(self.prompts, self.predictor) + self.scores = result.agg_scores.tolist() self.meta_prompt = self.meta_prompt_template.replace("", self._format_instructions()).replace( "", self._sample_examples() ) @@ -125,7 +126,8 @@ def _step(self) -> List[Prompt]: duplicate_prompts += 1 continue - score = self.task.evaluate(prompt, self.predictor)[0] + prompt_result = self.task.evaluate([prompt], self.predictor) + score = prompt_result.agg_scores.tolist()[0] self._add_prompt_and_score(prompt, score) diff --git a/promptolution/predictors/base_predictor.py b/promptolution/predictors/base_predictor.py index 292d56d8..f0060a16 100644 --- a/promptolution/predictors/base_predictor.py +++ b/promptolution/predictors/base_predictor.py @@ -37,8 +37,7 @@ def predict( prompts: Union[str, List[str]], xs: List[str], system_prompts: Optional[Union[str, List[str]]] = None, - return_seq: bool = False, - ) -> Union[List[str], Tuple[List[str], List[str]]]: + ) -> Tuple[List[str], List[str]]: """Abstract method to make predictions based on prompts and input data. Args: @@ -57,11 +56,8 @@ def predict( outputs = self.llm.get_response(inputs, system_prompts=system_prompts) preds = self._extract_preds(outputs) - if return_seq: - seqs = [f"{x}\n{out}" for x, out in zip(xs, outputs)] - return preds, seqs - - return preds + seqs = [f"{x}\n{out}" for x, out in zip(xs, outputs)] + return preds, seqs @abstractmethod def _extract_preds(self, preds: List[str]) -> List[str]: diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 4bd05272..11aa6364 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -2,11 +2,12 @@ from abc import ABC, abstractmethod +from dataclasses import dataclass import numpy as np import pandas as pd -from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union, cast, overload +from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union from promptolution.utils.logging import get_logger from promptolution.utils.prompt import Prompt @@ -23,6 +24,26 @@ logger = get_logger(__name__) +@dataclass +class Costs: + """Token costs accounting for model inputs and outputs.""" + + 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 + + +@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) + costs: Costs + + class BaseTask(ABC): """Abstract base class for tasks in the promptolution library.""" @@ -74,7 +95,7 @@ def __init__( 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): generating sequence per datapoint + self.seq_cache: Dict[Tuple[str, str, str], str] = {} # (prompt, x, y): raw model output per datapoint def subsample(self, eval_strategy: Optional["EvalStrategy"] = None) -> Tuple[List[str], List[str]]: """Subsample the dataset based on the specified parameters. @@ -113,284 +134,152 @@ def _prepare_batch( xs: List[str], ys: List[str], eval_strategy: Literal["full", "subsample", "sequential_block", "random_block", "evaluated"] = "full", - ) -> List[Tuple[str, str, str]]: - """Generate (prompt, x, y) keys that require prediction. - - Returns keys not found in eval_cache. - """ + ) -> 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 [] - keys_to_predict = [] + 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 not in self.eval_cache: - keys_to_predict.append(cache_key) - return keys_to_predict + 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) - def _collect_results_from_cache( - self, - prompts: List[Prompt], - xs: List[str], - ys: List[str], - return_agg_scores: bool, - return_seq: bool, - ) -> Union[List[float], List[List[float]], Tuple[List[List[float]], List[List[str]]]]: - """Collect all results for the current batch from the cache and format them.""" - assert not (return_agg_scores and return_seq), "Cannot return both aggregated scores and sequences" + return prompts_to_predict, xs_to_predict, ys_to_predict, keys_to_predict - scores = [] - seqs = [] + @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 = [] - datapoint_seqs = [] + datapoint_scores: List[float] = [] + datapoint_seqs: List[str] = [] for x, y in zip(xs, ys): - cache_key = (prompt.construct_prompt(), x, y) - datapoint_score = self.eval_cache.get(cache_key) - if datapoint_score is None: - continue + cache_key = self._cache_key(prompt, x, str(y)) + datapoint_score = self.eval_cache.get(cache_key, np.nan) datapoint_scores.append(datapoint_score) - if return_seq: - datapoint_seqs.append(self.seq_cache.get(cache_key, "")) - scores.append(datapoint_scores) - if return_seq: - seqs.append(datapoint_seqs) - - if return_agg_scores: - agg_scores = [np.nanmean(s).item() for s in scores] - return agg_scores - - return scores if not return_seq else (scores, seqs) + datapoint_seqs.append(self.seq_cache.get(cache_key, "")) + score_rows.append(datapoint_scores) + seq_rows.append(datapoint_seqs) - @abstractmethod - def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> List[float]: - """Abstract method to calculate the score for a predictions. - - This method should be implemented by subclasses based on their specific evaluation logic. - """ - raise NotImplementedError + 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 - @overload - def evaluate( + def _compute_costs( self, prompts: List[Prompt], + xs: List[str], + ys: List[str], + seq_cache: Dict[Tuple[str, str, str], str], predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[True] = True, - return_seq: Literal[False] = False, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[False] = False, - ) -> List[float]: - ... + ) -> Costs: + token_counter = get_token_counter(predictor.llm) - @overload - def evaluate( - self, - prompts: List[Prompt], - predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[False] = False, - return_seq: Literal[False] = False, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[False] = False, - ) -> List[List[float]]: - ... + per_prompt_inputs: List[np.ndarray] = [] + per_prompt_outputs: List[np.ndarray] = [] - @overload - def evaluate( - self, - prompts: List[Prompt], - predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[False] = False, - return_seq: Literal[True] = True, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[False] = False, - ) -> Tuple[List[List[float]], List[List[str]]]: - ... + input_token_counts = np.array([token_counter(x) for x in xs], dtype=float) - @overload - def evaluate( - self, - prompts: Prompt, - predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[True] = True, - return_seq: Literal[False] = False, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[False] = False, - ) -> List[float]: - ... + for prompt in prompts: + prompt_tokens = token_counter(prompt.construct_prompt()) + seq_token_counts: List[float] = [] + for x, y in zip(xs, ys): + cache_key = self._cache_key(prompt, x, str(y)) + seq_text = seq_cache.get(cache_key, "") + seq_token_counts.append(token_counter(seq_text)) - @overload - def evaluate( - self, - prompts: Prompt, - predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[False] = False, - return_seq: Literal[False] = False, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[False] = False, - ) -> List[List[float]]: - ... + prompt_input_tokens = prompt_tokens + input_token_counts + output_token_counts = np.maximum(np.array(seq_token_counts, dtype=float) - input_token_counts, 0.0) - @overload - def evaluate( - self, - prompts: Prompt, - predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[False] = False, - return_seq: Literal[True] = True, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[False] = False, - ) -> Tuple[List[List[float]], List[List[str]]]: - ... + per_prompt_inputs.append(np.asarray(prompt_input_tokens, dtype=float)) + per_prompt_outputs.append(output_token_counts) - @overload - def evaluate( - self, - prompts: List[Prompt], - predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[True] = True, - return_seq: Literal[False] = False, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[True] = True, - ) -> Tuple[List[float], List[float], List[float]]: - ... + inputs_array = np.vstack(per_prompt_inputs) + outputs_array = np.vstack(per_prompt_outputs) - @overload - def evaluate( - self, - prompts: List[Prompt], - predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[False] = False, - return_seq: Literal[False] = False, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[True] = True, - ) -> Tuple[List[List[float]], List[List[float]], List[List[float]]]: - ... + agg_input_tokens = inputs_array.mean(axis=1) if inputs_array.size else np.array([]) + agg_output_tokens = outputs_array.mean(axis=1) if outputs_array.size else np.array([]) - @overload - def evaluate( - self, - prompts: Prompt, - predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[True] = True, - return_seq: Literal[False] = False, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[True] = True, - ) -> Tuple[List[float], List[float], List[float]]: - ... + return Costs( + input_tokens=inputs_array, + output_tokens=outputs_array, + agg_input_tokens=agg_input_tokens, + agg_output_tokens=agg_output_tokens, + ) - @overload - def evaluate( - self, - prompts: Prompt, - predictor: "BasePredictor", - system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: Literal[False] = False, - return_seq: Literal[False] = False, - eval_strategy: Optional["EvalStrategy"] = None, - return_costs: Literal[True] = True, - ) -> Tuple[List[List[float]], List[List[float]], List[List[float]]]: - ... + @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 evaluate( self, prompts: Union[Prompt, List[Prompt]], predictor: "BasePredictor", system_prompts: Optional[Union[str, List[str]]] = None, - return_agg_scores: bool = True, - return_seq: bool = False, eval_strategy: Optional["EvalStrategy"] = None, - return_costs: bool = False, - ) -> Union[ - List[float], - List[List[float]], - Tuple[List[List[float]], List[List[str]]], - Tuple[List[float], List[float], List[float]], - Tuple[List[List[float]], List[List[float]], List[List[float]]], - ]: + ) -> EvalResult: """Evaluate a set of prompts using a given predictor. This method orchestrates subsampling, prediction, caching, and result collection. - - Note: Cannot return both aggregated scores and sequences (assertion will fail). + Sequences, token costs, raw scores, and aggregated scores are always returned. """ - assert not (return_agg_scores and return_seq), "Cannot return both aggregated scores and sequences" - assert not return_seq or not return_costs, "Token cost reporting is not supported together with sequences." - - prompts = [prompts] if isinstance(prompts, Prompt) else prompts + prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts) eval_strategy = eval_strategy or self.eval_strategy xs, ys = self.subsample(eval_strategy=eval_strategy) - batches = self._prepare_batch(prompts, xs, ys, eval_strategy=eval_strategy) - (prompts_to_evaluate, xs_to_evaluate, ys_to_evaluate) = ([], [], []) if not batches else zip(*batches) - - if prompts_to_evaluate: - preds_seqs = predictor.predict( - prompts=list(prompts_to_evaluate), - xs=list(xs_to_evaluate), - system_prompts=system_prompts, - return_seq=return_seq, - ) - else: - preds_seqs = ([], []) if return_seq else [] - - seqs: List[str] = [] - if return_seq: - preds, seqs = preds_seqs if isinstance(preds_seqs, tuple) else (preds_seqs, []) - else: - preds = cast(List[str], preds_seqs) + ( + 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: List[float] = self._evaluate(list(xs_to_evaluate), list(ys_to_evaluate), preds) - for i, cache_key in enumerate(batches): + 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] - if return_seq: - self.seq_cache[cache_key] = seqs[i] + self.seq_cache[cache_key] = str(pred_seqs[i]) - agg_scores = self._collect_results_from_cache( - prompts, + scores, agg_scores, seqs = self._collect_results_from_cache( + prompts_list, xs, ys, - return_agg_scores, - return_seq, ) - if not return_costs: - return agg_scores - - token_counter = get_token_counter(predictor.llm) - - per_prompt_inputs: List[List[float]] = [] - per_prompt_outputs: List[List[float]] = [] + costs = self._compute_costs(prompts_list, xs, ys, self.seq_cache, predictor) - input_token_counts = [float(token_counter(x)) for x in xs] - - for idx, prompt in enumerate(prompts): - prompt_tokens = float(token_counter(prompt.construct_prompt())) - start = idx * len(xs) - end = (idx + 1) * len(xs) - preds_for_prompt = preds[start:end] - output_token_counts = [float(token_counter(p)) for p in preds_for_prompt] - - prompt_input_tokens = [prompt_tokens + input_toks for input_toks in input_token_counts] - per_prompt_inputs.append(prompt_input_tokens) - per_prompt_outputs.append(output_token_counts) - - if return_agg_scores: - agg_scores_list = cast(List[float], agg_scores) - per_prompt_inputs_mean = [float(np.mean(tokens)) for tokens in per_prompt_inputs] - per_prompt_outputs_mean = [float(np.mean(tokens)) for tokens in per_prompt_outputs] - return agg_scores_list, per_prompt_inputs_mean, per_prompt_outputs_mean - - score_matrix = cast(List[List[float]], agg_scores) - return score_matrix, per_prompt_inputs, per_prompt_outputs + return EvalResult( + scores=scores, + agg_scores=agg_scores, + sequences=seqs, + costs=costs, + ) def pop_datapoints(self, n: Optional[int] = None, frac: Optional[float] = None) -> pd.DataFrame: """Pop a number of datapoints from the dataset. diff --git a/promptolution/tasks/classification_tasks.py b/promptolution/tasks/classification_tasks.py index e34c24f8..b5b5634e 100644 --- a/promptolution/tasks/classification_tasks.py +++ b/promptolution/tasks/classification_tasks.py @@ -67,9 +67,7 @@ def __init__( ) # 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]) -> List[float]: + def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarray: """Calculate the score for a single prediction.""" - scores = [] - for pred, y in zip(preds, ys): - scores.append(self.metric([y], [pred])) - return scores + scores = [self.metric([y], [pred]) for pred, y in zip(preds, ys)] + return np.asarray(scores, dtype=float) diff --git a/promptolution/tasks/judge_tasks.py b/promptolution/tasks/judge_tasks.py index 0f2fd4dc..19dca2f0 100644 --- a/promptolution/tasks/judge_tasks.py +++ b/promptolution/tasks/judge_tasks.py @@ -1,5 +1,6 @@ """Module for judge tasks.""" +import numpy as np import pandas as pd from typing import TYPE_CHECKING, List, Optional @@ -122,7 +123,7 @@ def _construct_judge_prompt(self, x: str, pred: str, y: Optional[str] = None) -> 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]) -> List[float]: + 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): @@ -145,4 +146,4 @@ def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> List[floa scores.append(score) - return scores + return np.asarray(scores, dtype=float) diff --git a/promptolution/tasks/reward_tasks.py b/promptolution/tasks/reward_tasks.py index cf92ed01..b09a06f3 100644 --- a/promptolution/tasks/reward_tasks.py +++ b/promptolution/tasks/reward_tasks.py @@ -1,6 +1,7 @@ """Module for Reward tasks.""" +import numpy as np import pandas as pd from typing import TYPE_CHECKING, Callable, List, Optional @@ -53,7 +54,7 @@ def __init__( config=config, ) - def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> List[float]: + def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarray: """Calculate the score for a single reward prediction using the reward function.""" rewards = [self.reward_function(pred) for pred in preds] - return rewards + return np.asarray(rewards, dtype=float) diff --git a/promptolution/utils/capo_utils.py b/promptolution/utils/capo_utils.py index 1e18c915..f9a8507d 100644 --- a/promptolution/utils/capo_utils.py +++ b/promptolution/utils/capo_utils.py @@ -43,7 +43,6 @@ def build_few_shot_examples( preds, seqs = predictor.predict( [instruction] * num_examples, list(sample_inputs), - return_seq=True, ) if isinstance(seqs, str): seqs = [seqs] diff --git a/promptolution/utils/prompt.py b/promptolution/utils/prompt.py index e9067981..935521f8 100644 --- a/promptolution/utils/prompt.py +++ b/promptolution/utils/prompt.py @@ -1,6 +1,8 @@ """Module defining the Prompt class and related utilities.""" -from typing import List, Optional, Tuple +import numpy as np + +from typing import List, Optional, Sequence, Tuple, Union from promptolution.utils.templates import DOWNSTREAM_TEMPLATE, DOWNSTREAM_TEMPLATE_W_FEWSHOTS @@ -48,25 +50,40 @@ def __str__(self) -> str: def sort_prompts_by_scores( - prompts: List[Prompt], scores: List[float], top_k: Optional[int] = None + prompts: List[Prompt], scores: Union[Sequence[float], np.ndarray], top_k: Optional[int] = None ) -> Tuple[List[Prompt], List[float]]: - """Sorts prompts based on their associated scores in descending order. + """Sort prompts by score, accepting scalar, 1D, or multi-dimensional scores. + + Scores can be provided as Python lists or NumPy arrays. If scores are multi- + dimensional (e.g., per-subsample results), they are aggregated with a + ``nanmean`` across all non-leading axes before sorting. Args: - prompts (List[Prompt]): List of Prompt objects. - scores (List[float]): Corresponding list of scores. - top_k (Optional[int]): If provided, limits the result to the top_k prompts. Defaults to None (returns all). + prompts (List[Prompt]): Prompt objects to sort. + scores (Sequence[float] | np.ndarray): Corresponding scores; can be nested lists or arrays. + top_k (Optional[int]): Limit the result to the top_k prompts. Returns: - Tuple[List[Prompt], List[float]]: A tuple containing prompts sorted by scores in descending order and their corresponding sorted scores. + Tuple[List[Prompt], List[float]]: Prompts and their aggregated scores, + sorted in descending order. """ - assert len(prompts) == len(scores), "Prompts and scores must have the same length." + scores_arr = np.asarray(scores, dtype=float) + if scores_arr.ndim == 0: + scores_arr = scores_arr.reshape(1) + + assert scores_arr.shape[0] == len(prompts), "Prompts and scores must have the same length." - sorted_prompts = [prompt for score, prompt in sorted(zip(scores, prompts), reverse=True, key=lambda x: x[0])] - sorted_scores = sorted(scores, reverse=True) + if scores_arr.ndim > 1: + axes_to_reduce = tuple(range(1, scores_arr.ndim)) + scores_arr = np.nanmean(scores_arr, axis=axes_to_reduce) + + prompt_score_pairs = list(zip(prompts, scores_arr.tolist())) + prompt_score_pairs.sort(key=lambda pair: pair[1], reverse=True) if top_k is not None: - sorted_prompts = sorted_prompts[:top_k] - sorted_scores = sorted_scores[:top_k] + prompt_score_pairs = prompt_score_pairs[:top_k] + + sorted_prompts = [p for p, _ in prompt_score_pairs] + sorted_scores = [s for _, s in prompt_score_pairs] return sorted_prompts, sorted_scores diff --git a/promptolution/utils/token_counter.py b/promptolution/utils/token_counter.py index 422e2771..75a6d408 100644 --- a/promptolution/utils/token_counter.py +++ b/promptolution/utils/token_counter.py @@ -13,7 +13,7 @@ logger = get_logger(__name__) -def get_token_counter(llm: "BaseLLM") -> Callable[[str], int]: +def get_token_counter(llm: "BaseLLM") -> Callable[[str], float]: """Get a token counter function for the given LLM. This function returns a callable that counts tokens based on the LLM's tokenizer @@ -28,7 +28,7 @@ def get_token_counter(llm: "BaseLLM") -> Callable[[str], int]: """ if llm.tokenizer is not None: tokenizer: "PreTrainedTokenizer" = llm.tokenizer - return lambda x: len(tokenizer.encode(x)) + return lambda x: float(len(tokenizer.encode(x))) else: logger.warning("⚠️ The LLM does not have a tokenizer. Using simple token count.") - return lambda x: len(x.split()) + return lambda x: float(len(x.split())) diff --git a/tests/helpers/test_helpers.py b/tests/helpers/test_helpers.py index d39ec385..76de258e 100644 --- a/tests/helpers/test_helpers.py +++ b/tests/helpers/test_helpers.py @@ -9,6 +9,7 @@ from tests.mocks.mock_task import MockTask from promptolution.helpers import run_evaluation, run_experiment, run_optimization +from promptolution.tasks.base_task import Costs, EvalResult from promptolution.utils import ExperimentConfig from promptolution.utils.prompt import Prompt @@ -200,7 +201,17 @@ def test_run_evaluation(mock_get_task, mock_get_predictor, mock_get_llm, sample_ prompts = [Prompt(p) for p in prompts] # Now this will work because mock_task is a MagicMock - mock_task.evaluate.return_value = np.array([0.8, 0.7, 0.9]) + mock_task.evaluate.return_value = EvalResult( + scores=np.array([[0.9], [0.8], [0.7]], dtype=float), + agg_scores=np.array([0.9, 0.8, 0.7], dtype=float), + sequences=np.array([["s1"], ["s2"], ["s3"]], dtype=object), + costs=Costs( + input_tokens=np.array([[10.0], [10.0], [10.0]], dtype=float), + output_tokens=np.array([[5.0], [5.0], [5.0]], dtype=float), + agg_input_tokens=np.array([10.0, 10.0, 10.0], dtype=float), + agg_output_tokens=np.array([5.0, 5.0, 5.0], dtype=float), + ), + ) # Run the function result = run_evaluation(sample_df, experiment_config, prompts) @@ -279,7 +290,19 @@ def test_helpers_integration(sample_df, experiment_config): # Use a MagicMock instead of MockTask mock_task = MagicMock() mock_task.classes = ["positive", "neutral", "negative"] - mock_task.evaluate = MagicMock(return_value=np.array([0.85, 0.75])) + mock_task.evaluate = MagicMock( + return_value=EvalResult( + scores=np.array([[0.9], [0.8]], dtype=float), + agg_scores=np.array([0.9, 0.8], dtype=float), + sequences=np.array([["s1"], ["s2"]], dtype=object), + costs=Costs( + input_tokens=np.array([[10.0], [10.0]], dtype=float), + output_tokens=np.array([[5.0], [5.0]], dtype=float), + agg_input_tokens=np.array([10.0, 10.0], dtype=float), + agg_output_tokens=np.array([5.0, 5.0], dtype=float), + ), + ) + ) mock_optimizer = MagicMock() diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py index 8650d2b8..4aab9f5f 100644 --- a/tests/optimizers/test_capoeira.py +++ b/tests/optimizers/test_capoeira.py @@ -1,8 +1,10 @@ from unittest.mock import MagicMock, patch +import numpy as np import pandas as pd from promptolution.optimizers.capoeira import Capoeira +from promptolution.tasks.base_task import Costs, EvalResult from promptolution.utils.capo_utils import perform_crossover, perform_mutation from promptolution.utils.prompt import Prompt from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE @@ -48,7 +50,19 @@ def test_capoeira_selection_prefers_better_score(mock_meta_llm, mock_predictor, df_few_shots=mock_df, ) candidates = [Prompt("short"), Prompt("longer prompt")] - optimizer.task.evaluate = MagicMock(return_value=([0.1, 0.9], [len("short"), len("longer prompt")], [5, 5])) # second candidate is better + optimizer.task.evaluate = MagicMock( + return_value=EvalResult( + scores=np.array([[0.4], [0.9]], dtype=float), + agg_scores=np.array([0.4, 0.9], dtype=float), + sequences=np.array([["s1"], ["s2"]], dtype=object), + costs=Costs( + input_tokens=np.array([[1.0], [1.0]], dtype=float), + output_tokens=np.array([[0.0], [0.0]], dtype=float), + agg_input_tokens=np.array([1.0, 1.0], dtype=float), + agg_output_tokens=np.array([0.0, 0.0], dtype=float), + ), + ) + ) objectives = optimizer._evaluate_candidates(candidates) selected, _ = optimizer._select_population(candidates, objectives) diff --git a/tests/predictors/test_base_predictor.py b/tests/predictors/test_base_predictor.py index 4bfeacdd..d20f51d4 100644 --- a/tests/predictors/test_base_predictor.py +++ b/tests/predictors/test_base_predictor.py @@ -8,7 +8,7 @@ def test_predictor_predict_flow(mock_predictor): prompts = ["Classify this text:"] # Call predict - predictions = mock_predictor.predict(prompts, xs) + predictions, _ = mock_predictor.predict(prompts, xs) # Verify shape and content of predictions assert predictions.shape == (1,) assert predictions[0] == "neutral" @@ -27,7 +27,7 @@ def test_predictor_with_return_seq(mock_predictor): xs = np.array(["This product is okay."]) # Call predict with return_seq=True - predictions, sequences = mock_predictor.predict(prompts, xs, return_seq=True) + predictions, sequences = mock_predictor.predict(prompts, xs) # Verify predictions assert predictions.shape == (1,) diff --git a/tests/predictors/test_predictors.py b/tests/predictors/test_predictors.py index 2f7e11fd..9fa5658f 100644 --- a/tests/predictors/test_predictors.py +++ b/tests/predictors/test_predictors.py @@ -14,7 +14,7 @@ def test_first_occurrence_classifier(mock_downstream_llm, mock_df): prompts = ["Classify:"] * len(xs) # Make predictions - predictions = classifier.predict(prompts, xs) + predictions, _ = classifier.predict(prompts, xs) # Verify shape and content assert len(predictions) == 4 @@ -39,7 +39,7 @@ def test_marker_based_classifier(mock_downstream_llm, mock_df): prompts = ["Classify:"] * len(xs) # Make predictions - predictions = classifier.predict(prompts, xs) + predictions, _ = classifier.predict(prompts, xs) # Verify shape and content assert len(predictions) == 3 @@ -49,7 +49,7 @@ def test_marker_based_classifier(mock_downstream_llm, mock_df): # Test with invalid class label invalid_input = np.array(["Broken item"] * len(prompts)) - invalid_predictions = classifier.predict(prompts, invalid_input) + invalid_predictions, _ = classifier.predict(prompts, invalid_input) # Should default to first class if invalid assert invalid_predictions[0] == "positive" @@ -70,7 +70,7 @@ def test_marker_based_without_classes(mock_downstream_llm): prompts = ["Classify:"] * len(xs) # Make predictions - predictions = predictor.predict(prompts, xs) + predictions, _ = predictor.predict(prompts, xs) # Verify shape and content - should accept any value between markers assert len(predictions) == 4 @@ -90,7 +90,7 @@ def test_multiple_prompts_with_classifiers(mock_downstream_llm, mock_df): xs = np.array(["I love this product!", "I hate this product!"] * 2) # Make predictions - predictions = classifier.predict(prompts, xs) + predictions, _ = classifier.predict(prompts, xs) # Verify shape and content assert len(predictions) == 4 @@ -110,7 +110,7 @@ def test_sequence_return_with_classifiers(mock_downstream_llm, mock_df): xs = np.array(["I love this product!"]) # Make predictions with sequences - predictions, sequences = classifier.predict(prompts, xs, return_seq=True) + predictions, sequences = classifier.predict(prompts, xs) # Verify predictions assert len(predictions) == 1 @@ -140,7 +140,7 @@ def test_marker_based_missing_markers(mock_downstream_llm): # When markers are missing, it should default to first class prompts = ["Classify:"] - xs = np.array(["Missing markers"]) - predictions = classifier.predict(prompts, xs) + xs = ["Missing markers"] + preds, seqs = classifier.predict(prompts, xs) - assert predictions[0] == "will" # Should default to first class + assert preds[0] == "will" # Should default to first class diff --git a/tests/tasks/test_classifications_tasks.py b/tests/tasks/test_classifications_tasks.py index 256a63db..9ab13bd9 100644 --- a/tests/tasks/test_classifications_tasks.py +++ b/tests/tasks/test_classifications_tasks.py @@ -21,44 +21,48 @@ def test_classification_task_initialization(mock_df): def test_task_evaluate(mock_classification_task_with_subsampling, mock_predictor): """Test the evaluate method of ClassificationTask.""" prompts = [Prompt("Classify sentiment:")] - scores = mock_classification_task_with_subsampling.evaluate(prompts, mock_predictor) + result = mock_classification_task_with_subsampling.evaluate(prompts, mock_predictor) + scores = result.agg_scores - assert isinstance(scores, list) - assert len(scores) == 1 + assert scores.shape == (1,) assert 0 <= scores[0] <= 1 prompts = ["Classify sentiment:", "Rate the text:"] prompts = [Prompt(p) for p in prompts] - scores = mock_classification_task_with_subsampling.evaluate(prompts, mock_predictor) + result = mock_classification_task_with_subsampling.evaluate(prompts, mock_predictor) + scores = result.agg_scores - assert len(scores) == 2 - assert all(0 <= score <= 1 for score in scores) + assert scores.shape == (2,) + assert np.all((scores >= 0) & (scores <= 1)) def test_task_evaluate_with_subsampling(mock_classification_task_with_subsampling, mock_predictor): """Test the evaluate method with subsampling.""" prompts = [Prompt("Classify sentiment:")] - scores = mock_classification_task_with_subsampling.evaluate( + scores_result = mock_classification_task_with_subsampling.evaluate( prompts, mock_predictor, ) + scores = scores_result.agg_scores - assert len(scores) == 1 + assert scores.shape == (1,) with pytest.raises(AssertionError, match=r".*Arrays are not equal.*"): np.random.seed(42) - scores1 = mock_classification_task_with_subsampling.evaluate( + scores1_result = mock_classification_task_with_subsampling.evaluate( prompts, mock_predictor, ) + scores1 = scores1_result.scores np.random.seed(43) - scores2 = mock_classification_task_with_subsampling.evaluate( + scores2_result = mock_classification_task_with_subsampling.evaluate( prompts, mock_predictor, ) + scores2 = scores2_result.scores np.testing.assert_array_equal(scores1, scores2) @@ -67,14 +71,13 @@ def test_task_evaluate_with_return_seq(mock_classification_task_with_subsampling """Test the evaluate method with return_seq=True.""" prompts = [Prompt("Classify sentiment:")] - scores, seqs = mock_classification_task_with_subsampling.evaluate( - prompts, mock_predictor, return_seq=True, return_agg_scores=False - ) + seq_result = mock_classification_task_with_subsampling.evaluate(prompts, mock_predictor) - assert len(scores) == 1 - assert len(scores[0]) == mock_classification_task_with_subsampling.n_subsamples - assert len(seqs) == 1 - assert len(seqs[0]) == mock_classification_task_with_subsampling.n_subsamples + assert seq_result.scores.shape == (1, mock_classification_task_with_subsampling.n_subsamples) + assert seq_result.sequences is not None + assert len(seq_result.sequences) == 1 + assert len(seq_result.sequences[0]) == mock_classification_task_with_subsampling.n_subsamples + assert seq_result.costs is not None def test_task_evaluate_with_system_prompts( @@ -85,11 +88,9 @@ def test_task_evaluate_with_system_prompts( prompts = [Prompt("Classify sentiment:")] system_prompts = ["Be concise"] - scores = mock_classification_task_with_subsampling.evaluate( - prompts, mock_predictor, system_prompts=system_prompts, return_agg_scores=True - ) + result = mock_classification_task_with_subsampling.evaluate(prompts, mock_predictor, system_prompts=system_prompts) - assert len(scores) == 1 + assert result.agg_scores.shape == (1,) assert any(call["system_prompts"] == system_prompts for call in mock_downstream_llm.call_history) @@ -97,7 +98,7 @@ def test_pop_datapoints(mock_df): task = ClassificationTask( df=mock_df, task_description="Sentiment classification task", - eval_strategy="sequential_blocks", + eval_strategy="sequential_block", ) df = task.pop_datapoints(n=1) @@ -108,7 +109,7 @@ def test_pop_datapoints(mock_df): def test_blocks(mock_df): task = ClassificationTask( - df=mock_df, task_description="Sentiment classification task", eval_strategy="sequential_blocks", n_subsamples=1 + df=mock_df, task_description="Sentiment classification task", eval_strategy="sequential_block", n_subsamples=1 ) task.increment_block_idx() diff --git a/tests/tasks/test_judge_task.py b/tests/tasks/test_judge_task.py index 3cf00664..ecb446b1 100644 --- a/tests/tasks/test_judge_task.py +++ b/tests/tasks/test_judge_task.py @@ -57,17 +57,18 @@ def test_judge_task_evaluate_with_ground_truth(mock_judge_task_with_y, mock_pred mock_predictor.call_history = [] mock_judge_llm.call_history = [] - scores_per_datapoint = mock_judge_task_with_y.evaluate(prompts, mock_predictor, return_agg_scores=False) + result = mock_judge_task_with_y.evaluate(prompts, mock_predictor) + scores_per_datapoint = result.scores - assert len(scores_per_datapoint) == len(prompts) + assert scores_per_datapoint.shape[0] == len(prompts) expected_scores = [1.0, 0, 0.5] np.testing.assert_allclose(scores_per_datapoint[0], expected_scores) mock_predictor.call_history = [] mock_judge_llm.call_history = [] - aggregated_scores = mock_judge_task_with_y.evaluate(prompts, mock_predictor, return_agg_scores=True) - assert len(aggregated_scores) == len(prompts) + aggregated_scores = result.agg_scores + assert aggregated_scores.shape[0] == len(prompts) expected_scores = [0.5, 0.4333333, 0.0] np.testing.assert_allclose(aggregated_scores, expected_scores) @@ -80,9 +81,10 @@ def test_judge_task_evaluate_no_ground_truth(mock_judge_task_no_y, mock_predicto mock_predictor.call_history = [] mock_judge_llm.call_history = [] - aggregated_scores = mock_judge_task_no_y.evaluate(prompts, mock_predictor, return_agg_scores=True) + aggregated_result = mock_judge_task_no_y.evaluate(prompts, mock_predictor) + aggregated_scores = aggregated_result.agg_scores - assert len(aggregated_scores) == len(prompts) + assert aggregated_scores.shape[0] == len(prompts) expected_scores = [0.5, 0.55, 0.35] np.testing.assert_allclose(aggregated_scores, expected_scores) @@ -92,9 +94,9 @@ def test_judge_task_evaluate_with_return_seq(mock_judge_task_with_y, mock_predic prompts = ["Evaluate this text:", "What is the sentiment?", "How would you classify this?"] prompts = [Prompt(p) for p in prompts] - scores, seqs = mock_judge_task_with_y.evaluate(prompts, mock_predictor, return_seq=True, return_agg_scores=False) + seq_result = mock_judge_task_with_y.evaluate(prompts, mock_predictor) - assert len(scores) == 3 - assert len(scores[0]) == len(mock_judge_task_with_y.xs) - assert len(seqs) == 3 - assert len(seqs[0]) == len(mock_judge_task_with_y.xs) + assert seq_result.scores.shape == (3, len(mock_judge_task_with_y.xs)) + assert seq_result.sequences is not None + assert seq_result.sequences.shape == (3, len(mock_judge_task_with_y.xs)) + assert seq_result.costs is not None diff --git a/tests/tasks/test_reward_tasks.py b/tests/tasks/test_reward_tasks.py index 76e35454..10299c53 100644 --- a/tests/tasks/test_reward_tasks.py +++ b/tests/tasks/test_reward_tasks.py @@ -24,7 +24,9 @@ def test_reward_task_evaluate_with_return_seq(mock_reward_task, mock_predictor): """Test the evaluate method with return_seq=True for RewardTask.""" prompts = [Prompt("Generate a short text:")] - scores, seqs = mock_reward_task.evaluate(prompts, mock_predictor, return_seq=True, return_agg_scores=False) + result = mock_reward_task.evaluate(prompts, mock_predictor) - assert len(scores) == 1 - assert len(seqs) == 1 + assert result.scores.shape[0] == 1 + assert result.sequences is not None + assert result.sequences.shape[0] == 1 + assert result.costs is not None diff --git a/tests/utils/test_prompt.py b/tests/utils/test_prompt.py index 3dc90bb1..9ee3c244 100644 --- a/tests/utils/test_prompt.py +++ b/tests/utils/test_prompt.py @@ -1,3 +1,5 @@ +import numpy as np + from promptolution.utils.prompt import Prompt, sort_prompts_by_scores @@ -39,3 +41,14 @@ def test_sort_prompts_by_scores(): # Verify sorting assert sorted_prompts == [prompt2, prompt1, prompt3] assert sorted_scores == [0.90, 0.75, 0.60] + + +def test_sort_prompts_by_scores_with_array(): + """Ensure sorting works when scores are numpy arrays (aggregated via mean).""" + prompts = [Prompt("p1"), Prompt("p2"), Prompt("p3")] + scores = np.array([[0.5, 0.7], [0.8, 0.9], [0.4, 0.6]]) + + sorted_prompts, sorted_scores = sort_prompts_by_scores(prompts, scores) + + assert sorted_prompts == [prompts[1], prompts[0], prompts[2]] + np.testing.assert_allclose(sorted_scores, [0.85, 0.6, 0.5]) From 6d461a975b1a428be8c76c6fd6f21bb33f729ed7 Mon Sep 17 00:00:00 2001 From: Tom Zehle Date: Mon, 29 Dec 2025 16:13:01 +0100 Subject: [PATCH 12/53] Delete .vscode/settings.json --- .vscode/settings.json | 7 ------- 1 file changed, 7 deletions(-) delete mode 100644 .vscode/settings.json diff --git a/.vscode/settings.json b/.vscode/settings.json deleted file mode 100644 index a3a18383..00000000 --- a/.vscode/settings.json +++ /dev/null @@ -1,7 +0,0 @@ -{ - "python.testing.pytestArgs": [ - "tests" - ], - "python.testing.unittestEnabled": false, - "python.testing.pytestEnabled": true -} From 4063f2478250123e4f3adb26fb7fbc45da6b9099 Mon Sep 17 00:00:00 2001 From: finitearth Date: Mon, 29 Dec 2025 15:16:32 +0000 Subject: [PATCH 13/53] Update coverage badge in README [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 184e0cd9..e62ac6a7 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@ -![Coverage](https://img.shields.io/badge/Coverage-91%25-brightgreen) +![Coverage](https://img.shields.io/badge/Coverage-90%25-brightgreen) [![CI](https://github.com/automl/promptolution/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/automl/promptolution/actions/workflows/ci.yml) [![Docs](https://github.com/automl/promptolution/actions/workflows/docs.yml/badge.svg?branch=main)](https://github.com/automl/promptolution/actions/workflows/docs.yml) ![Code Style](https://img.shields.io/badge/Code%20Style-black-black) From 84119336277ed9cf680ea52097997fa08d2c1634 Mon Sep 17 00:00:00 2001 From: finitearth Date: Sat, 3 Jan 2026 18:58:59 +0100 Subject: [PATCH 14/53] refining capoeira --- promptolution/optimizers/capo.py | 30 +--- promptolution/optimizers/capoeira.py | 253 ++++++++++++++++++++------- promptolution/utils/capo_utils.py | 58 +++--- 3 files changed, 217 insertions(+), 124 deletions(-) diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index 97f1f659..161552b8 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -128,15 +128,7 @@ def _initialize_population(self, initial_prompts: List[Prompt]) -> List[Prompt]: few_shots = build_few_shot_examples( instruction=prompt.instruction, num_examples=num_examples, - df_few_shots=self.df_few_shots, - x_column=self.task.x_column, - y_column=self.task.y_column, - predictor=self.predictor, - fewshot_template=CAPO_FEWSHOT_TEMPLATE, - target_begin_marker=self.target_begin_marker, - target_end_marker=self.target_end_marker, - check_fs_accuracy=self.check_fs_accuracy, - create_fs_reasoning=self.create_fs_reasoning, + optimizer=self, ) population.append(Prompt(prompt.instruction, few_shots)) @@ -197,24 +189,8 @@ def _pre_optimization_loop(self) -> None: def _step(self) -> List[Prompt]: """Perform a single optimization step.""" - offsprings = perform_crossover(self.prompts, self.crossovers_per_iter, self.crossover_template, self.meta_llm) - mutated = perform_mutation( - offsprings=offsprings, - mutation_template=self.mutation_template, - upper_shots=self.upper_shots, - meta_llm=self.meta_llm, - few_shot_kwargs=dict( - df_few_shots=self.df_few_shots, - x_column=self.task.x_column, - y_column=self.task.y_column, - predictor=self.predictor, - fewshot_template=CAPO_FEWSHOT_TEMPLATE, - target_begin_marker=self.target_begin_marker, - target_end_marker=self.target_end_marker, - check_fs_accuracy=self.check_fs_accuracy, - create_fs_reasoning=self.create_fs_reasoning, - ), - ) + offsprings = perform_crossover(self.prompts, optimizer=self) + mutated = perform_mutation(offsprings=offsprings, optimizer=self) combined = self.prompts + mutated self.prompts, self.scores = self._do_racing(combined, self.population_size) diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 9ea6b3e7..b1d58d44 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -7,7 +7,7 @@ import numpy as np import pandas as pd -from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, Dict, List, Optional, Tuple if TYPE_CHECKING: # pragma: no cover from promptolution.utils.callbacks import BaseCallback @@ -20,7 +20,7 @@ from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation from promptolution.utils.logging import get_logger from promptolution.utils.prompt import Prompt -from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE +from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_MUTATION_TEMPLATE from promptolution.utils.token_counter import get_token_counter logger = get_logger(__name__) @@ -78,12 +78,21 @@ def __init__( super().__init__(predictor, task, initial_prompts, callbacks, config) + self.incumbents: List[Prompt] = [] + self.challengers: List[Prompt] = [] + self.crossover_template = self._initialize_meta_template(crossover_template or CAPO_CROSSOVER_TEMPLATE) self.mutation_template = self._initialize_meta_template(mutation_template or CAPO_MUTATION_TEMPLATE) self.token_counter = get_token_counter(self.downstream_llm) self.df_few_shots = df_few_shots if df_few_shots is not None else task.pop_datapoints(frac=0.1) self.population_size = len(self.prompts) + if "block" not in self.task.eval_strategy: + logger.warning( + f"ℹ️ CAPO requires 'block' in the eval_strategy, but got {self.task.eval_strategy}. Setting eval_strategy to 'sequential_block'." + ) + self.task.eval_strategy = "sequential_block" + if hasattr(self.predictor, "begin_marker") and hasattr(self.predictor, "end_marker"): self.target_begin_marker = self.predictor.begin_marker # type: ignore self.target_end_marker = self.predictor.end_marker # type: ignore @@ -98,15 +107,7 @@ def _pre_optimization_loop(self) -> None: few_shots = build_few_shot_examples( instruction=prompt.instruction, num_examples=num_examples, - df_few_shots=self.df_few_shots, - x_column=self.task.x_column, - y_column=self.task.y_column, - predictor=self.predictor, - fewshot_template=CAPO_FEWSHOT_TEMPLATE, - target_begin_marker=self.target_begin_marker, - target_end_marker=self.target_end_marker, - check_fs_accuracy=self.check_fs_accuracy, - create_fs_reasoning=self.create_fs_reasoning, + optimizer=self, ) population.append(Prompt(prompt.instruction, few_shots)) @@ -115,28 +116,118 @@ def _pre_optimization_loop(self) -> None: self.max_prompt_length = ( max(self.token_counter(p.construct_prompt()) for p in self.prompts) if self.prompts else 1 ) - initial_vectors = self._evaluate_candidates(self.prompts) - self.prompts, selected_vectors = self._select_population(self.prompts, initial_vectors) - self.scores = (-selected_vectors[:, 0]).tolist() - - def _evaluate_candidates(self, candidates: List[Prompt]) -> np.ndarray: - result = self.task.evaluate( - candidates, - self.predictor, - eval_strategy=self.task.eval_strategy, - ) + initial_vectors = self._calculate_objective_vector(self.prompts) #TODO rename + fronts = self.fast_non_dominated_sort(initial_vectors) + self.incumbents = [self.prompts[i] for i in fronts[0]] + self.challengers = [self.prompts[i] for front in fronts[1:] for i in front] + + # keep self.prompts as a "view" if base class expects it + self.prompts = self.incumbents + self.challengers + self.scores = initial_vectors[:, 0].tolist() + + def _do_intensification(self, challenger: Prompt) -> None: + """ + Default MO-CAPO intensification (closest-incumbent comparison): + - evaluate challenger + incumbents on sequential blocks + - maintain running averages (challenger and incumbents) + - early reject if closest incumbent dominates challenger average + - if challenger survives all blocks: promote to incumbents and update front + """ + if not self.incumbents: + self.incumbents.append(challenger) + return + + # Start race from a consistent block index + self.task.reset_block_idx() # TODO this might need to change + + chal_hist: List[np.ndarray] = [] + inc_hist: Dict[int, List[np.ndarray]] = {i: [] for i in range(len(self.incumbents))} + + for _ in range(self.task.n_blocks): + joint_result = self.task.evaluate( + self.incumbents + [challenger], + self.predictor, + eval_strategy="sequential_block", + ) + joint_vecs = self._objective_vectors_from_result(joint_result) - scores = result.scores - input_tokens = result.costs.input_tokens - output_tokens = result.costs.output_tokens + inc_vecs = joint_vecs[:-1] + chal_vec = joint_vecs[-1] - score_vectors = np.column_stack( - [ - -scores, - self.cost_per_input_token * input_tokens + self.cost_per_output_token * output_tokens, - ] + chal_hist.append(chal_vec) + for i, v in enumerate(inc_vecs): + inc_hist[i].append(v) + + chal_avg = np.mean(chal_hist, axis=0) + + # Default: compare only against closest incumbent (in normalized objective space) + closest = self._get_closest_incumbent(chal_avg) + closest_idx = self.incumbents.index(closest) + closest_avg = np.mean(inc_hist[closest_idx], axis=0) + + if self._is_dominated(chal_avg, closest_avg): + # challenger loses -> goes to population + self.challengers.append(challenger) + self.task.reset_block_idx() + return + + self.task.increment_block_idx() + + # Survived full race -> promote and update incumbent front + self.incumbents.append(challenger) + self._update_incumbent_front() + self.task.reset_block_idx() + + + def _get_closest_incumbent(self, challenger_vec: np.ndarray): + """Finds the geometrically closest incumbent.""" + inc_vecs = self._calculate_objective_vector(self.incumbents, eval_strategy="sequential_block") + all_vecs = np.vstack([inc_vecs, challenger_vec[None, :]]) + min_b = np.min(all_vecs, axis=0) + max_b = np.max(all_vecs, axis=0) + rng = max_b - min_b + rng[rng == 0] = 1.0 # Avoid div/0 + + norm_chal = (challenger_vec - min_b) / rng + norm_incs = (inc_vecs - min_b) / rng + + dists = np.linalg.norm(norm_incs - norm_chal, axis=1) + return self.incumbents[np.argmin(dists)] + + + def _update_incumbent_front(self) -> None: + """ + After adding a challenger that survived a full race, recompute the incumbent Pareto front. + Default behavior: incumbents become front-0 (on current evaluation state), + all other incumbents are demoted to challengers. + """ + if not self.incumbents: + return + + vecs = self._calculate_objective_vector(self.incumbents, eval_strategy="sequential_block") + fronts = self.fast_non_dominated_sort(vecs) + + new_incumbents = [self.incumbents[i] for i in fronts[0]] + demoted = [self.incumbents[i] for front in fronts[1:] for i in front] + + self.incumbents = new_incumbents + self.challengers.extend(demoted) + + + def _calculate_objective_vector(self, prompts: List[Prompt], eval_strategy=None) -> np.ndarray: + eval_result = self.task.evaluate( + prompts=prompts, + predictor=self.predictor, + eval_strategy=eval_strategy, ) - return score_vectors + return self._objective_vectors_from_result(eval_result) + + def _objective_vectors_from_result(self, result) -> np.ndarray: + agg_scores = result.agg_scores + agg_input_tokens = result.costs.agg_input_tokens + agg_output_tokens = result.costs.agg_output_tokens + cost_scalar = self.cost_per_input_token * agg_input_tokens + self.cost_per_output_token * agg_output_tokens + return np.column_stack([agg_scores, -cost_scalar]) def _select_population( self, candidates: List[Prompt], score_vectors: np.ndarray @@ -158,44 +249,74 @@ def _select_population( selected_vectors = score_vectors[selected_indices] return selected_prompts, selected_vectors - def _step(self) -> List[Prompt]: - offsprings = perform_crossover(self.prompts, self.crossovers_per_iter, self.crossover_template, self.meta_llm) - mutated = perform_mutation( - offsprings=offsprings, - mutation_template=self.mutation_template, - upper_shots=self.upper_shots, - meta_llm=self.meta_llm, - few_shot_kwargs=dict( - df_few_shots=self.df_few_shots, - x_column=self.task.x_column, - y_column=self.task.y_column, - predictor=self.predictor, - fewshot_template=CAPO_FEWSHOT_TEMPLATE, - target_begin_marker=self.target_begin_marker, - target_end_marker=self.target_end_marker, - check_fs_accuracy=self.check_fs_accuracy, - create_fs_reasoning=self.create_fs_reasoning, - ), + + def _advance_one_incumbent(self) -> None: + """ + Default MO-CAPO step after processing a challenger: + evaluate one incumbent on one additional sequential block. + (With your current task API, this is the closest equivalent to the + "catch-up / new block" logic.) + """ + if not self.incumbents: + return + + chosen = random.choice(self.incumbents) + + _ = self.task.evaluate( # TODO might need to change + prompts=[chosen], + predictor=self.predictor, + eval_strategy="sequential_block", ) - combined = self.prompts + mutated + self.task.increment_block_idx() + + def _prune_population(self) -> None: + """ + Enforce |incumbents| + |challengers| <= population_size. + Default behavior: prune challengers first; if none, prune incumbents by crowding distance. + """ + while len(self.incumbents) + len(self.challengers) > self.population_size: + if self.challengers: + # simplest default: remove a random challenger + self.challengers.pop(random.randrange(len(self.challengers))) + else: + if len(self.incumbents) <= 1: + break + vecs = self._calculate_objective_vector(self.incumbents, eval_strategy="sequential_block") + dists = self.calculate_crowding_distance(vecs) + worst = int(np.argmin(dists)) + self.incumbents.pop(worst) + + + def _step(self) -> List[Prompt]: + # 1) generate challengers (random parent selection happens inside perform_crossover) + offsprings = perform_crossover(self.prompts, optimizer=self) + new_challengers = perform_mutation(offsprings=offsprings, optimizer=self) + + # 2) intensify each challenger; after each, advance incumbents + prune + for chal in new_challengers: + self._do_intensification(chal) + self._advance_one_incumbent() + self._prune_population() + + # 3) update "view" for base class / callbacks + self.prompts = self.incumbents + self.challengers + + # 4) logging scores: incumbents only (optional) + if self.incumbents: + vecs_inc = self._calculate_objective_vector(self.incumbents, eval_strategy="sequential_block") + self.scores = vecs_inc[:, 0].tolist() + else: + self.scores = [] - score_vectors = self._evaluate_candidates(combined) - self.prompts, selected_vectors = self._select_population(combined, score_vectors) - self.scores = (-selected_vectors[:, 0]).tolist() return self.prompts - def get_pareto_front(self) -> List[Dict[str, Any]]: - """Return the current Pareto front with objective values.""" - score_vectors = self._evaluate_candidates(self.prompts) - return [ - { - "prompt": prompt.construct_prompt(), - "score": float(score_vectors[i][0] * -1), - "cost": float(score_vectors[i][1]), - } - for i, prompt in enumerate(self.prompts) - ] + + @staticmethod + def _is_dominated(vec1, vec2): + """Returns True if vec2 dominates vec1 in a maximize-all setting.""" + return np.all(vec2 >= vec1) and np.any(vec2 > vec1) + @staticmethod def fast_non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: """Perform fast non-dominated sorting (NSGA-II) in a vectorized manner.""" @@ -203,9 +324,9 @@ def fast_non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: if num_solutions == 0: return [] - less = obj_vectors[:, None, :] < obj_vectors[None, :, :] - less_equal = obj_vectors[:, None, :] <= obj_vectors[None, :, :] - dominates = np.all(less_equal, axis=2) & np.any(less, axis=2) + greater = obj_vectors[:, None, :] > obj_vectors[None, :, :] + greater_equal = obj_vectors[:, None, :] >= obj_vectors[None, :, :] + dominates = np.all(greater_equal, axis=2) & np.any(greater, axis=2) domination_counts = dominates.sum(axis=0) dominated_solutions = [list(np.where(dominates[i])[0]) for i in range(num_solutions)] diff --git a/promptolution/utils/capo_utils.py b/promptolution/utils/capo_utils.py index f9a8507d..404a994d 100644 --- a/promptolution/utils/capo_utils.py +++ b/promptolution/utils/capo_utils.py @@ -4,10 +4,9 @@ import random -import pandas as pd - -from typing import List, Optional +from typing import List +from promptolution.utils.templates import CAPO_FEWSHOT_TEMPLATE from promptolution.utils.formatting import extract_from_tag from promptolution.utils.prompt import Prompt @@ -15,32 +14,26 @@ def build_few_shot_examples( instruction: str, num_examples: int, - df_few_shots: pd.DataFrame, - x_column: str, - y_column: Optional[str], - predictor, - fewshot_template: str, - target_begin_marker: str, - target_end_marker: str, - check_fs_accuracy: bool, - create_fs_reasoning: bool, + optimizer, ) -> List[str]: """Create few-shot examples with optional reasoning replacement.""" if num_examples == 0: return [] - few_shot_samples = df_few_shots.sample(num_examples, replace=False) - sample_inputs = few_shot_samples[x_column].values.astype(str) - sample_targets = few_shot_samples[y_column].values + few_shot_samples = optimizer.df_few_shots.sample(num_examples, replace=False) + sample_inputs = few_shot_samples[optimizer.task.x_column].values.astype(str) + sample_targets = few_shot_samples[optimizer.task.y_column].values few_shots = [ - fewshot_template.replace("", i).replace("", f"{target_begin_marker}{t}{target_end_marker}") + CAPO_FEWSHOT_TEMPLATE.replace("", i).replace( + "", f"{optimizer.target_begin_marker}{t}{optimizer.target_end_marker}" + ) for i, t in zip(sample_inputs, sample_targets) ] - if not create_fs_reasoning: + if not optimizer.create_fs_reasoning: return few_shots - preds, seqs = predictor.predict( + preds, seqs = optimizer.predictor.predict( [instruction] * num_examples, list(sample_inputs), ) @@ -51,25 +44,23 @@ def build_few_shot_examples( for j in range(num_examples): seqs[j] = seqs[j].replace(sample_inputs[j], "", 1).strip() - if preds[j] == sample_targets[j] or not check_fs_accuracy: - few_shots[j] = fewshot_template.replace("", sample_inputs[j]).replace("", seqs[j]) + if preds[j] == sample_targets[j] or not optimizer.check_fs_accuracy: + few_shots[j] = CAPO_FEWSHOT_TEMPLATE.replace("", sample_inputs[j]).replace("", seqs[j]) return few_shots def perform_crossover( parents: List[Prompt], - crossovers_per_iter: int, - template: str, - meta_llm, + optimizer, ) -> List[Prompt]: """Generate crossover offspring prompts.""" crossover_prompts: List[str] = [] offspring_few_shots: List[List[str]] = [] - for _ in range(crossovers_per_iter): + for _ in range(optimizer.crossovers_per_iter): mother, father = parents if len(parents) == 2 else random.sample(parents, 2) crossover_prompt = ( - template.replace("", mother.instruction).replace("", father.instruction).strip() + optimizer.crossover_template.replace("", mother.instruction).replace("", father.instruction).strip() ) crossover_prompts.append(crossover_prompt) combined_few_shots = mother.few_shots + father.few_shots @@ -77,7 +68,7 @@ def perform_crossover( offspring_few_shot = random.sample(combined_few_shots, num_few_shots) if combined_few_shots else [] offspring_few_shots.append(offspring_few_shot) - child_instructions = meta_llm.get_response(crossover_prompts) + child_instructions = optimizer.meta_llm.get_response(crossover_prompts) return [ Prompt(extract_from_tag(instr, "", ""), examples) for instr, examples in zip(child_instructions, offspring_few_shots) @@ -86,12 +77,13 @@ def perform_crossover( def perform_mutation( offsprings: List[Prompt], - mutation_template: str, - upper_shots: int, - meta_llm, - few_shot_kwargs: dict, + optimizer, ) -> List[Prompt]: """Mutate offspring prompts.""" + mutation_template = optimizer.mutation_template + meta_llm = optimizer.meta_llm + upper_shots = optimizer.upper_shots + mutation_prompts = [mutation_template.replace("", prompt.instruction) for prompt in offsprings] new_instructions = meta_llm.get_response(mutation_prompts) @@ -101,7 +93,11 @@ def perform_mutation( p = random.random() if p < 1 / 3 and len(prompt.few_shots) < upper_shots: - new_few_shot = build_few_shot_examples(new_instruction, 1, **few_shot_kwargs) + new_few_shot = build_few_shot_examples( + instruction=new_instruction, + num_examples=1, + optimizer=optimizer, + ) new_few_shots = prompt.few_shots + new_few_shot elif 1 / 3 <= p < 2 / 3 and len(prompt.few_shots) > 0: new_few_shots = random.sample(prompt.few_shots, len(prompt.few_shots) - 1) From 3bb14efac3947aa1a506226df3a2e0e912d77d93 Mon Sep 17 00:00:00 2001 From: finitearth Date: Mon, 5 Jan 2026 21:08:11 +0100 Subject: [PATCH 15/53] further refinements --- promptolution/optimizers/capoeira.py | 312 ++++++++++++-------- promptolution/tasks/__init__.py | 2 + promptolution/tasks/base_task.py | 53 ++-- promptolution/tasks/multi_objective_task.py | 161 ++++++++++ tests/helpers/test_helpers.py | 22 +- tests/optimizers/test_capoeira.py | 12 +- tests/tasks/test_classifications_tasks.py | 2 +- tests/tasks/test_judge_task.py | 2 +- tests/tasks/test_reward_tasks.py | 2 +- 9 files changed, 396 insertions(+), 172 deletions(-) create mode 100644 promptolution/tasks/multi_objective_task.py diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index b1d58d44..33fdfcc1 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -7,7 +7,7 @@ import numpy as np import pandas as pd -from typing import TYPE_CHECKING, Dict, List, Optional, Tuple +from typing import TYPE_CHECKING, List, Optional if TYPE_CHECKING: # pragma: no cover from promptolution.utils.callbacks import BaseCallback @@ -17,6 +17,8 @@ from promptolution.utils.config import ExperimentConfig from promptolution.optimizers.base_optimizer import BaseOptimizer +from promptolution.tasks.multi_objective_task import MultiObjectiveTask + from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation from promptolution.utils.logging import get_logger from promptolution.utils.prompt import Prompt @@ -116,8 +118,9 @@ def _pre_optimization_loop(self) -> None: self.max_prompt_length = ( max(self.token_counter(p.construct_prompt()) for p in self.prompts) if self.prompts else 1 ) - initial_vectors = self._calculate_objective_vector(self.prompts) #TODO rename - fronts = self.fast_non_dominated_sort(initial_vectors) + init_result = self.task.evaluate(prompts=self.prompts, predictor=self.predictor) + initial_vectors = self._get_objective_vectors(init_result) #TODO rename + fronts = self._non_dominated_sort(initial_vectors) self.incumbents = [self.prompts[i] for i in fronts[0]] self.challengers = [self.prompts[i] for front in fronts[1:] for i in front] @@ -125,6 +128,31 @@ def _pre_optimization_loop(self) -> None: self.prompts = self.incumbents + self.challengers self.scores = initial_vectors[:, 0].tolist() + + def _step(self) -> List[Prompt]: + # 1) generate challengers (random parent selection happens inside perform_crossover) + offsprings = perform_crossover(self.prompts, optimizer=self) + new_challengers = perform_mutation(offsprings=offsprings, optimizer=self) + + # 2) intensify each challenger; after each, advance incumbents + prune + for chal in new_challengers: + self._do_intensification(chal) + self._advance_one_incumbent() + self._prune_population() + + # 3) update "view" for base class / callbacks + self.prompts = self.incumbents + self.challengers + + # 4) logging scores: incumbents only (optional) + if self.incumbents: + inc_result = self.task.evaluate(prompts=self.incumbents, predictor=self.predictor, eval_strategy="evaluated") + vecs_inc = self._get_objective_vectors(inc_result) + self.scores = vecs_inc[:, 0].tolist() + else: + self.scores = [] + + return self.prompts + def _do_intensification(self, challenger: Prompt) -> None: """ Default MO-CAPO intensification (closest-incumbent comparison): @@ -137,62 +165,62 @@ def _do_intensification(self, challenger: Prompt) -> None: self.incumbents.append(challenger) return - # Start race from a consistent block index - self.task.reset_block_idx() # TODO this might need to change - chal_hist: List[np.ndarray] = [] - inc_hist: Dict[int, List[np.ndarray]] = {i: [] for i in range(len(self.incumbents))} + common_block_idx = 0 + while common_block_idx is not None: + common_block_idx = self._sample_common_block(self.incumbents) + self.task.set_block_idx(common_block_idx) # type: ignore - for _ in range(self.task.n_blocks): joint_result = self.task.evaluate( - self.incumbents + [challenger], - self.predictor, - eval_strategy="sequential_block", + prompts=self.incumbents + [challenger], + predictor=self.predictor ) - joint_vecs = self._objective_vectors_from_result(joint_result) - inc_vecs = joint_vecs[:-1] - chal_vec = joint_vecs[-1] + objective_vectors = self._get_objective_vectors(joint_result) + challenger_vec = objective_vectors[-1] + incumbent_vecs = objective_vectors[:-1] - chal_hist.append(chal_vec) - for i, v in enumerate(inc_vecs): - inc_hist[i].append(v) + closest_inc_vec = self._get_closest_incumbent(challenger_vec, incumbent_vecs) - chal_avg = np.mean(chal_hist, axis=0) - - # Default: compare only against closest incumbent (in normalized objective space) - closest = self._get_closest_incumbent(chal_avg) - closest_idx = self.incumbents.index(closest) - closest_avg = np.mean(inc_hist[closest_idx], axis=0) - - if self._is_dominated(chal_avg, closest_avg): + if self._is_dominated(challenger_vec, closest_inc_vec): # challenger loses -> goes to population self.challengers.append(challenger) - self.task.reset_block_idx() return - self.task.increment_block_idx() - - # Survived full race -> promote and update incumbent front self.incumbents.append(challenger) self._update_incumbent_front() - self.task.reset_block_idx() + def _sample_common_block(self, prompts: List[Prompt]) -> Optional[int]: + """Sample a block index that has been evaluated by all given prompts. + Returns None if no such block exists.""" + per_prompt = self.task.get_evaluated_blocks(prompts) # Dict[prompt -> Set[int]] + block_sets = list(per_prompt.values()) + + if not block_sets: + return random.randrange(self.task.n_blocks) + + common = set.intersection(*block_sets) + if not common: + return None - def _get_closest_incumbent(self, challenger_vec: np.ndarray): - """Finds the geometrically closest incumbent.""" - inc_vecs = self._calculate_objective_vector(self.incumbents, eval_strategy="sequential_block") - all_vecs = np.vstack([inc_vecs, challenger_vec[None, :]]) + return random.choice(tuple(common)) + + def _get_closest_incumbent( + self, challenger_vec: np.ndarray, incumbent_vecs: np.ndarray + ) -> np.ndarray: + """Return the vector of the geometrically closest incumbent.""" + all_vecs = np.vstack([incumbent_vecs, challenger_vec[None, :]]) min_b = np.min(all_vecs, axis=0) max_b = np.max(all_vecs, axis=0) rng = max_b - min_b rng[rng == 0] = 1.0 # Avoid div/0 norm_chal = (challenger_vec - min_b) / rng - norm_incs = (inc_vecs - min_b) / rng + norm_incs = (incumbent_vecs - min_b) / rng dists = np.linalg.norm(norm_incs - norm_chal, axis=1) - return self.incumbents[np.argmin(dists)] + idx = int(np.argmin(dists)) + return incumbent_vecs[idx] def _update_incumbent_front(self) -> None: @@ -204,8 +232,9 @@ def _update_incumbent_front(self) -> None: if not self.incumbents: return - vecs = self._calculate_objective_vector(self.incumbents, eval_strategy="sequential_block") - fronts = self.fast_non_dominated_sort(vecs) + vecs_result = self.task.evaluate(prompts=self.incumbents, predictor=self.predictor, eval_strategy="evaluated") + vecs = self._get_objective_vectors(vecs_result) + fronts = self._non_dominated_sort(vecs) new_incumbents = [self.incumbents[i] for i in fronts[0]] demoted = [self.incumbents[i] for front in fronts[1:] for i in front] @@ -214,141 +243,168 @@ def _update_incumbent_front(self) -> None: self.challengers.extend(demoted) - def _calculate_objective_vector(self, prompts: List[Prompt], eval_strategy=None) -> np.ndarray: - eval_result = self.task.evaluate( - prompts=prompts, - predictor=self.predictor, - eval_strategy=eval_strategy, - ) - return self._objective_vectors_from_result(eval_result) + def _get_objective_vectors(self, result) -> np.ndarray: - def _objective_vectors_from_result(self, result) -> np.ndarray: - agg_scores = result.agg_scores - agg_input_tokens = result.costs.agg_input_tokens - agg_output_tokens = result.costs.agg_output_tokens - cost_scalar = self.cost_per_input_token * agg_input_tokens + self.cost_per_output_token * agg_output_tokens - return np.column_stack([agg_scores, -cost_scalar]) - - def _select_population( - self, candidates: List[Prompt], score_vectors: np.ndarray - ) -> Tuple[List[Prompt], np.ndarray]: - selected_indices: List[int] = [] - fronts = self.fast_non_dominated_sort(score_vectors) - for front in fronts: - if len(selected_indices) + len(front) <= self.population_size: - selected_indices.extend(front) - else: - remaining = self.population_size - len(selected_indices) - front_vectors = score_vectors[front] - distances = self.calculate_crowding_distance(front_vectors) - sorted_front = [i for _, i in sorted(zip(distances, front), reverse=True)] - selected_indices.extend(sorted_front[:remaining]) - break + # If the task is multi-objective, include all objective dimensions, else single objective. + if isinstance(self.task, MultiObjectiveTask): + agg_scores = np.stack(result.agg_scores, axis=1) # shape: (n_prompts, n_objectives) + else: + agg_scores = np.atleast_2d(result.agg_scores).T # shape: (n_prompts, 1) - selected_prompts = [candidates[i] for i in selected_indices] - selected_vectors = score_vectors[selected_indices] - return selected_prompts, selected_vectors + agg_input_tokens = np.asarray(result.agg_input_tokens) + agg_output_tokens = np.asarray(result.agg_output_tokens) + cost_scalar = self.cost_per_input_token * agg_input_tokens + self.cost_per_output_token * agg_output_tokens + cost_scalar = cost_scalar.reshape(-1, 1) + return np.hstack([agg_scores, -cost_scalar]) def _advance_one_incumbent(self) -> None: """ Default MO-CAPO step after processing a challenger: evaluate one incumbent on one additional sequential block. - (With your current task API, this is the closest equivalent to the - "catch-up / new block" logic.) """ - if not self.incumbents: - return + # choose least evaluated incumbent + eval_counts = [ + len(self.task.get_evaluated_blocks([inc])) for inc in self.incumbents + ] + min_count = min(eval_counts) + candidates = [inc for inc, count in zip(self.incumbents, eval_counts) if count == min_count] + chosen = random.sample(candidates, k=1) + self.task.evaluate(prompts=chosen, predictor=self.predictor) - chosen = random.choice(self.incumbents) - - _ = self.task.evaluate( # TODO might need to change - prompts=[chosen], - predictor=self.predictor, - eval_strategy="sequential_block", - ) - self.task.increment_block_idx() def _prune_population(self) -> None: """ - Enforce |incumbents| + |challengers| <= population_size. - Default behavior: prune challengers first; if none, prune incumbents by crowding distance. + Enforce |incumbents| + |challengers| <= population_size using Pareto logic. + + Logic: + 1. Prune from Challengers first (they are less optimal than incumbents). + - If challengers have DIFFERENT evaluation blocks (Heterogeneous): + We cannot fairly compare their scores. Prune the one with the FEWEST evaluations + (least information/newest). + - If challengers have the SAME evaluation blocks (Homogeneous): + Perform Non-Dominated Sorting (NDS). Identify the worst front. + Use Crowding Distance to prune the most crowded (least unique) individual from that front. + + 2. If no Challengers, prune from Incumbents. + - Use Crowding Distance to remove the least unique incumbent. """ while len(self.incumbents) + len(self.challengers) > self.population_size: if self.challengers: - # simplest default: remove a random challenger - self.challengers.pop(random.randrange(len(self.challengers))) + # 1. Check Heterogeneity (Fairness Check) + chal_blocks_map = self.task.get_evaluated_blocks(self.challengers) + block_sets = list(chal_blocks_map.values()) + + # Ensure we have data to compare + if not block_sets: + self.challengers.pop(random.randrange(len(self.challengers))) + continue + + first_set = block_sets[0] + # Are all challengers evaluated on the exact same set of blocks? + is_homogeneous = all(s == first_set for s in block_sets) + + if not is_homogeneous: + # CASE A: Heterogeneous (Unfair comparison). + # Prune the prompt with the FEWEST evaluations (least reliable/least invested). + counts = [len(s) for s in block_sets] + min_count = min(counts) + + # Find all indices with the minimum count (handle ties randomly) + candidates = [i for i, c in enumerate(counts) if c == min_count] + victim_idx = random.choice(candidates) + + self.challengers.pop(victim_idx) + + else: + # CASE B: Homogeneous (Fair comparison). + # Use NDS + Crowding Distance. + + # Get objective vectors for all challengers (safe because blocks are identical) + res = self.task.evaluate( + self.challengers, + self.predictor, + eval_strategy="evaluated" + ) + vecs = self._get_objective_vectors(res) + + # Perform Non-Dominated Sort + fronts = self._non_dominated_sort(vecs) + + # Select the worst front (the last one) + worst_front_indices = fronts[-1] + + if len(worst_front_indices) == 1: + # Only one candidate in the worst front -> prune it + victim_idx = worst_front_indices[0] + else: + # Multiple candidates in worst front -> Prune by Crowding Distance + # We want to keep diversity (high CD), so we remove low CD. + worst_front_vecs = vecs[worst_front_indices] + dists = self._calculate_crowding_distance(worst_front_vecs) + + # Find index relative to the worst front list + local_worst_idx = int(np.argmin(dists)) + # Map back to the main challenger list index + victim_idx = worst_front_indices[local_worst_idx] + + self.challengers.pop(victim_idx) + else: + # --- PRUNE FROM INCUMBENTS --- + # Fallback: If we only have incumbents, remove the least unique one. if len(self.incumbents) <= 1: break - vecs = self._calculate_objective_vector(self.incumbents, eval_strategy="sequential_block") - dists = self.calculate_crowding_distance(vecs) - worst = int(np.argmin(dists)) - self.incumbents.pop(worst) - - - def _step(self) -> List[Prompt]: - # 1) generate challengers (random parent selection happens inside perform_crossover) - offsprings = perform_crossover(self.prompts, optimizer=self) - new_challengers = perform_mutation(offsprings=offsprings, optimizer=self) + + res = self.task.evaluate( + self.incumbents, + self.predictor, + eval_strategy="evaluated" + ) + vecs = self._get_objective_vectors(res) + dists = self._calculate_crowding_distance(vecs) + + # Remove the one with the smallest crowding distance + victim_idx = int(np.argmin(dists)) + self.incumbents.pop(victim_idx) - # 2) intensify each challenger; after each, advance incumbents + prune - for chal in new_challengers: - self._do_intensification(chal) - self._advance_one_incumbent() - self._prune_population() - - # 3) update "view" for base class / callbacks self.prompts = self.incumbents + self.challengers - # 4) logging scores: incumbents only (optional) - if self.incumbents: - vecs_inc = self._calculate_objective_vector(self.incumbents, eval_strategy="sequential_block") - self.scores = vecs_inc[:, 0].tolist() - else: - self.scores = [] - - return self.prompts - - - @staticmethod - def _is_dominated(vec1, vec2): - """Returns True if vec2 dominates vec1 in a maximize-all setting.""" - return np.all(vec2 >= vec1) and np.any(vec2 > vec1) - - @staticmethod - def fast_non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: + def _non_dominated_sort(self, obj_vectors: np.ndarray) -> List[List[int]]: """Perform fast non-dominated sorting (NSGA-II) in a vectorized manner.""" - num_solutions = obj_vectors.shape[0] - if num_solutions == 0: - return [] + n_solutions = obj_vectors.shape[0] greater = obj_vectors[:, None, :] > obj_vectors[None, :, :] greater_equal = obj_vectors[:, None, :] >= obj_vectors[None, :, :] dominates = np.all(greater_equal, axis=2) & np.any(greater, axis=2) domination_counts = dominates.sum(axis=0) - dominated_solutions = [list(np.where(dominates[i])[0]) for i in range(num_solutions)] + dominated_solutions = [list(np.where(dominates[i])[0]) for i in range(n_solutions)] fronts: List[List[int]] = [list(np.where(domination_counts == 0)[0])] - current_front = 0 - while current_front < len(fronts) and fronts[current_front]: + current_front = 0 + while current_front < len(fronts) and len(fronts[current_front]) > 0: next_front: List[int] = [] for i in fronts[current_front]: for dominated in dominated_solutions[i]: domination_counts[dominated] -= 1 if domination_counts[dominated] == 0: next_front.append(dominated) - if next_front: + if len(next_front) > 0: fronts.append(next_front) current_front += 1 return fronts @staticmethod - def calculate_crowding_distance(obj_vectors: np.ndarray) -> np.ndarray: + def _is_dominated(vec1, vec2): + """Returns True if vec2 dominates vec1 in a maximize-all setting.""" + return np.all(vec2 >= vec1) and np.any(vec2 > vec1) + + @staticmethod + def _calculate_crowding_distance(obj_vectors: np.ndarray) -> np.ndarray: """Calculate crowding distance for a set of solutions.""" num_solutions, num_obj = obj_vectors.shape if num_solutions <= 2: diff --git a/promptolution/tasks/__init__.py b/promptolution/tasks/__init__.py index 825dbadc..5f61ff1f 100644 --- a/promptolution/tasks/__init__.py +++ b/promptolution/tasks/__init__.py @@ -3,9 +3,11 @@ from promptolution.tasks.classification_tasks import ClassificationTask from promptolution.tasks.judge_tasks import JudgeTask from promptolution.tasks.reward_tasks import RewardTask +from promptolution.tasks.multi_objective_task import MultiObjectiveTask __all__ = [ "ClassificationTask", "JudgeTask", "RewardTask", + "MultiObjectiveTask", ] diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 11aa6364..d28ffb0f 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -18,22 +18,12 @@ from promptolution.utils.config import ExperimentConfig -TaskType = Literal["classification", "reward", "judge"] +TaskType = Literal["classification", "reward", "judge", "multi"] EvalStrategy = Literal["full", "subsample", "sequential_block", "random_block", "evaluated"] logger = get_logger(__name__) -@dataclass -class Costs: - """Token costs accounting for model inputs and outputs.""" - - 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 - - @dataclass class EvalResult: """Evaluation outputs including scores, sequences, and costs.""" @@ -41,7 +31,10 @@ class EvalResult: 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) - costs: Costs + 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 class BaseTask(ABC): @@ -97,6 +90,8 @@ def __init__( 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[str, set[int]] = {} # prompt_str: set of evaluated block indices + def subsample(self, eval_strategy: Optional["EvalStrategy"] = None) -> Tuple[List[str], List[str]]: """Subsample the dataset based on the specified parameters. @@ -190,7 +185,7 @@ def _compute_costs( ys: List[str], seq_cache: Dict[Tuple[str, str, str], str], predictor: "BasePredictor", - ) -> Costs: + ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: token_counter = get_token_counter(predictor.llm) per_prompt_inputs: List[np.ndarray] = [] @@ -218,12 +213,7 @@ def _compute_costs( agg_input_tokens = inputs_array.mean(axis=1) if inputs_array.size else np.array([]) agg_output_tokens = outputs_array.mean(axis=1) if outputs_array.size else np.array([]) - return Costs( - input_tokens=inputs_array, - output_tokens=outputs_array, - agg_input_tokens=agg_input_tokens, - agg_output_tokens=agg_output_tokens, - ) + 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: @@ -272,13 +262,22 @@ def evaluate( ys, ) - costs = self._compute_costs(prompts_list, xs, ys, self.seq_cache, predictor) + # Record evaluated block for block strategies + for prompt in prompts_list: + self.prompt_evaluated_blocks.setdefault(str(prompt), set()).add(self.block_idx) + + input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = self._compute_costs( + prompts_list, xs, ys, self.seq_cache, predictor + ) return EvalResult( scores=scores, agg_scores=agg_scores, sequences=seqs, - costs=costs, + 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: @@ -344,3 +343,15 @@ def reset_block_idx(self) -> None: 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 for subsampling (block strategies only).""" + if "block" not in self.eval_strategy: + raise ValueError("Block assignment is only valid for block subsampling.") + if self.n_blocks > 0: + self.block_idx = idx % self.n_blocks + else: + self.block_idx = 0 + + def get_evaluated_blocks(self, prompts: List[Prompt]) -> Dict[str, set[int]]: + return {str(p): set(self.prompt_evaluated_blocks.get(str(p), set())) for p in prompts} diff --git a/promptolution/tasks/multi_objective_task.py b/promptolution/tasks/multi_objective_task.py new file mode 100644 index 00000000..6d5e8f0e --- /dev/null +++ b/promptolution/tasks/multi_objective_task.py @@ -0,0 +1,161 @@ +"""Multi-objective task wrapper that evaluates prompts across multiple tasks.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import numpy as np + +from promptolution.tasks.base_task import BaseTask, EvalResult, EvalStrategy, TaskType +from promptolution.utils.prompt import Prompt + + +@dataclass +class MultiObjectiveEvalResult: + 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 + + +class MultiObjectiveTask(BaseTask): + """A task that aggregates evaluations across multiple underlying tasks.""" + + def __init__( + self, + tasks: List[BaseTask], + eval_strategy: Optional[EvalStrategy] = None, + ) -> None: + 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: TaskType = "multi" + self.tasks = tasks + + def evaluate( # type: ignore + self, + prompts: Prompt | List[Prompt], + predictor, + system_prompts: Optional[str | List[str]] = None, + eval_strategy: Optional[EvalStrategy] = None, + ) -> MultiObjectiveEvalResult: + """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, task.seq_cache, predictor + ) + + # Record evaluated block for block strategies + for prompt in prompts_list: + task.prompt_evaluated_blocks.setdefault(str(prompt), set()).add(task.block_idx) + + 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] + + # Mirror evaluated block bookkeeping using the first task for parity with BaseTask. + first_task = self.tasks[0] + self.prompt_evaluated_blocks = { + str(p): first_task.prompt_evaluated_blocks[str(p)] for p in prompts_list + } + + return MultiObjectiveEvalResult( + scores=stacked_scores, + agg_scores=stacked_agg_scores, + sequences=per_task_results[0].sequences, + input_tokens=per_task_results[0].input_tokens, + output_tokens=per_task_results[0].output_tokens, + agg_input_tokens=per_task_results[0].agg_input_tokens, + agg_output_tokens=per_task_results[0].agg_output_tokens, + ) + + def _evaluate(self, xs, ys, preds): # pragma: no cover + raise NotImplementedError("MultiObjectiveTask overrides evaluate directly") diff --git a/tests/helpers/test_helpers.py b/tests/helpers/test_helpers.py index 76de258e..b467cc2a 100644 --- a/tests/helpers/test_helpers.py +++ b/tests/helpers/test_helpers.py @@ -9,7 +9,7 @@ from tests.mocks.mock_task import MockTask from promptolution.helpers import run_evaluation, run_experiment, run_optimization -from promptolution.tasks.base_task import Costs, EvalResult +from promptolution.tasks.base_task import EvalResult from promptolution.utils import ExperimentConfig from promptolution.utils.prompt import Prompt @@ -205,12 +205,10 @@ def test_run_evaluation(mock_get_task, mock_get_predictor, mock_get_llm, sample_ scores=np.array([[0.9], [0.8], [0.7]], dtype=float), agg_scores=np.array([0.9, 0.8, 0.7], dtype=float), sequences=np.array([["s1"], ["s2"], ["s3"]], dtype=object), - costs=Costs( - input_tokens=np.array([[10.0], [10.0], [10.0]], dtype=float), - output_tokens=np.array([[5.0], [5.0], [5.0]], dtype=float), - agg_input_tokens=np.array([10.0, 10.0, 10.0], dtype=float), - agg_output_tokens=np.array([5.0, 5.0, 5.0], dtype=float), - ), + input_tokens=np.array([[10.0], [10.0], [10.0]], dtype=float), + output_tokens=np.array([[5.0], [5.0], [5.0]], dtype=float), + agg_input_tokens=np.array([10.0, 10.0, 10.0], dtype=float), + agg_output_tokens=np.array([5.0, 5.0, 5.0], dtype=float), ) # Run the function @@ -295,12 +293,10 @@ def test_helpers_integration(sample_df, experiment_config): scores=np.array([[0.9], [0.8]], dtype=float), agg_scores=np.array([0.9, 0.8], dtype=float), sequences=np.array([["s1"], ["s2"]], dtype=object), - costs=Costs( - input_tokens=np.array([[10.0], [10.0]], dtype=float), - output_tokens=np.array([[5.0], [5.0]], dtype=float), - agg_input_tokens=np.array([10.0, 10.0], dtype=float), - agg_output_tokens=np.array([5.0, 5.0], dtype=float), - ), + input_tokens=np.array([[10.0], [10.0]], dtype=float), + output_tokens=np.array([[5.0], [5.0]], dtype=float), + agg_input_tokens=np.array([10.0, 10.0], dtype=float), + agg_output_tokens=np.array([5.0, 5.0], dtype=float), ) ) diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py index 4aab9f5f..45602afe 100644 --- a/tests/optimizers/test_capoeira.py +++ b/tests/optimizers/test_capoeira.py @@ -4,7 +4,7 @@ import pandas as pd from promptolution.optimizers.capoeira import Capoeira -from promptolution.tasks.base_task import Costs, EvalResult +from promptolution.tasks.base_task import EvalResult from promptolution.utils.capo_utils import perform_crossover, perform_mutation from promptolution.utils.prompt import Prompt from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE @@ -55,12 +55,10 @@ def test_capoeira_selection_prefers_better_score(mock_meta_llm, mock_predictor, scores=np.array([[0.4], [0.9]], dtype=float), agg_scores=np.array([0.4, 0.9], dtype=float), sequences=np.array([["s1"], ["s2"]], dtype=object), - costs=Costs( - input_tokens=np.array([[1.0], [1.0]], dtype=float), - output_tokens=np.array([[0.0], [0.0]], dtype=float), - agg_input_tokens=np.array([1.0, 1.0], dtype=float), - agg_output_tokens=np.array([0.0, 0.0], dtype=float), - ), + input_tokens=np.array([[1.0], [1.0]], dtype=float), + output_tokens=np.array([[0.0], [0.0]], dtype=float), + agg_input_tokens=np.array([1.0, 1.0], dtype=float), + agg_output_tokens=np.array([0.0, 0.0], dtype=float), ) ) diff --git a/tests/tasks/test_classifications_tasks.py b/tests/tasks/test_classifications_tasks.py index 9ab13bd9..642d7c94 100644 --- a/tests/tasks/test_classifications_tasks.py +++ b/tests/tasks/test_classifications_tasks.py @@ -77,7 +77,7 @@ def test_task_evaluate_with_return_seq(mock_classification_task_with_subsampling assert seq_result.sequences is not None assert len(seq_result.sequences) == 1 assert len(seq_result.sequences[0]) == mock_classification_task_with_subsampling.n_subsamples - assert seq_result.costs is not None + assert seq_result.agg_input_tokens is not None def test_task_evaluate_with_system_prompts( diff --git a/tests/tasks/test_judge_task.py b/tests/tasks/test_judge_task.py index ecb446b1..15a6032c 100644 --- a/tests/tasks/test_judge_task.py +++ b/tests/tasks/test_judge_task.py @@ -99,4 +99,4 @@ def test_judge_task_evaluate_with_return_seq(mock_judge_task_with_y, mock_predic assert seq_result.scores.shape == (3, len(mock_judge_task_with_y.xs)) assert seq_result.sequences is not None assert seq_result.sequences.shape == (3, len(mock_judge_task_with_y.xs)) - assert seq_result.costs is not None + assert seq_result.agg_input_tokens is not None diff --git a/tests/tasks/test_reward_tasks.py b/tests/tasks/test_reward_tasks.py index 10299c53..e8dc48d4 100644 --- a/tests/tasks/test_reward_tasks.py +++ b/tests/tasks/test_reward_tasks.py @@ -29,4 +29,4 @@ def test_reward_task_evaluate_with_return_seq(mock_reward_task, mock_predictor): assert result.scores.shape[0] == 1 assert result.sequences is not None assert result.sequences.shape[0] == 1 - assert result.costs is not None + assert result.agg_input_tokens is not None From f0caca0573b5d852cc4238225bdc85572825c089 Mon Sep 17 00:00:00 2001 From: finitearth Date: Mon, 5 Jan 2026 21:20:40 +0100 Subject: [PATCH 16/53] minor clean up --- promptolution/optimizers/capoeira.py | 111 ++++++++++++--------------- 1 file changed, 47 insertions(+), 64 deletions(-) diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 33fdfcc1..88450484 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -102,6 +102,10 @@ def __init__( self.target_begin_marker = "" self.target_end_marker = "" + @property + def prompts(self) -> List[Prompt]: + return self.incumbents + self.challengers + def _pre_optimization_loop(self) -> None: population: List[Prompt] = [] for prompt in self.prompts: @@ -113,26 +117,23 @@ def _pre_optimization_loop(self) -> None: ) population.append(Prompt(prompt.instruction, few_shots)) - self.prompts = population - # TODO: align placement of the logic with capo self.max_prompt_length = ( - max(self.token_counter(p.construct_prompt()) for p in self.prompts) if self.prompts else 1 + max(self.token_counter(p.construct_prompt()) for p in population) if population else 1 ) - init_result = self.task.evaluate(prompts=self.prompts, predictor=self.predictor) + init_result = self.task.evaluate(population, self.predictor) initial_vectors = self._get_objective_vectors(init_result) #TODO rename fronts = self._non_dominated_sort(initial_vectors) - self.incumbents = [self.prompts[i] for i in fronts[0]] - self.challengers = [self.prompts[i] for front in fronts[1:] for i in front] + self.incumbents = [population[i] for i in fronts[0]] + self.challengers = [population[i] for front in fronts[1:] for i in front] # keep self.prompts as a "view" if base class expects it - self.prompts = self.incumbents + self.challengers self.scores = initial_vectors[:, 0].tolist() def _step(self) -> List[Prompt]: # 1) generate challengers (random parent selection happens inside perform_crossover) - offsprings = perform_crossover(self.prompts, optimizer=self) - new_challengers = perform_mutation(offsprings=offsprings, optimizer=self) + offsprings = perform_crossover(self.prompts, self) + new_challengers = perform_mutation(offsprings, self) # 2) intensify each challenger; after each, advance incumbents + prune for chal in new_challengers: @@ -140,9 +141,6 @@ def _step(self) -> List[Prompt]: self._advance_one_incumbent() self._prune_population() - # 3) update "view" for base class / callbacks - self.prompts = self.incumbents + self.challengers - # 4) logging scores: incumbents only (optional) if self.incumbents: inc_result = self.task.evaluate(prompts=self.incumbents, predictor=self.predictor, eval_strategy="evaluated") @@ -165,7 +163,6 @@ def _do_intensification(self, challenger: Prompt) -> None: self.incumbents.append(challenger) return - common_block_idx = 0 while common_block_idx is not None: common_block_idx = self._sample_common_block(self.incumbents) @@ -290,16 +287,11 @@ def _prune_population(self) -> None: - Use Crowding Distance to remove the least unique incumbent. """ while len(self.incumbents) + len(self.challengers) > self.population_size: - if self.challengers: + if len(self.challengers) > 0: # 1. Check Heterogeneity (Fairness Check) chal_blocks_map = self.task.get_evaluated_blocks(self.challengers) block_sets = list(chal_blocks_map.values()) - # Ensure we have data to compare - if not block_sets: - self.challengers.pop(random.randrange(len(self.challengers))) - continue - first_set = block_sets[0] # Are all challengers evaluated on the exact same set of blocks? is_homogeneous = all(s == first_set for s in block_sets) @@ -315,60 +307,51 @@ def _prune_population(self) -> None: victim_idx = random.choice(candidates) self.challengers.pop(victim_idx) + continue - else: - # CASE B: Homogeneous (Fair comparison). - # Use NDS + Crowding Distance. - - # Get objective vectors for all challengers (safe because blocks are identical) - res = self.task.evaluate( - self.challengers, - self.predictor, - eval_strategy="evaluated" - ) - vecs = self._get_objective_vectors(res) - - # Perform Non-Dominated Sort - fronts = self._non_dominated_sort(vecs) - - # Select the worst front (the last one) - worst_front_indices = fronts[-1] - - if len(worst_front_indices) == 1: - # Only one candidate in the worst front -> prune it - victim_idx = worst_front_indices[0] - else: - # Multiple candidates in worst front -> Prune by Crowding Distance - # We want to keep diversity (high CD), so we remove low CD. - worst_front_vecs = vecs[worst_front_indices] - dists = self._calculate_crowding_distance(worst_front_vecs) - - # Find index relative to the worst front list - local_worst_idx = int(np.argmin(dists)) - # Map back to the main challenger list index - victim_idx = worst_front_indices[local_worst_idx] - - self.challengers.pop(victim_idx) - - else: - # --- PRUNE FROM INCUMBENTS --- - # Fallback: If we only have incumbents, remove the least unique one. - if len(self.incumbents) <= 1: - break + # CASE B: Homogeneous (Fair comparison). + # Use NDS + Crowding Distance. + # Get objective vectors for all challengers (safe because blocks are identical) res = self.task.evaluate( - self.incumbents, + self.challengers, self.predictor, eval_strategy="evaluated" ) vecs = self._get_objective_vectors(res) - dists = self._calculate_crowding_distance(vecs) - # Remove the one with the smallest crowding distance - victim_idx = int(np.argmin(dists)) - self.incumbents.pop(victim_idx) + # Perform Non-Dominated Sort + fronts = self._non_dominated_sort(vecs) + + # Select the worst front (the last one) + worst_front_indices = fronts[-1] + + # Multiple candidates in worst front -> Prune by Crowding Distance + # We want to keep diversity (high CD), so we remove low CD. + worst_front_vecs = vecs[worst_front_indices] + dists = self._calculate_crowding_distance(worst_front_vecs) + + # Find index relative to the worst front list + local_worst_idx = int(np.argmin(dists)) + # Map back to the main challenger list index + victim_idx = worst_front_indices[local_worst_idx] + + self.challengers.pop(victim_idx) + continue - self.prompts = self.incumbents + self.challengers + # --- PRUNE FROM INCUMBENTS --- + # Fallback: If we only have incumbents, remove the least unique one. + res = self.task.evaluate( + self.incumbents, + self.predictor, + eval_strategy="evaluated" + ) + vecs = self._get_objective_vectors(res) + dists = self._calculate_crowding_distance(vecs) + + # Remove the one with the smallest crowding distance + victim_idx = int(np.argmin(dists)) + self.incumbents.pop(victim_idx) def _non_dominated_sort(self, obj_vectors: np.ndarray) -> List[List[int]]: From 09cd80505894ec0ab6d2fbcfaa2cc18e1de4ee41 Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 6 Jan 2026 17:25:34 +0100 Subject: [PATCH 17/53] refine testing --- .coverage | Bin 69632 -> 69632 bytes README.md | 2 +- promptolution/helpers.py | 11 +- promptolution/optimizers/capoeira.py | 183 +++++--- promptolution/tasks/base_task.py | 36 +- tests/helpers/test_helpers.py | 127 ++++- tests/llms/test_api_llm.py | 138 +++++- tests/llms/test_base_llm.py | 35 ++ tests/mocks/dummy_config.py | 15 + tests/mocks/mock_task.py | 76 ++- tests/optimizers/test_base_optimizer.py | 86 ++++ tests/optimizers/test_capo.py | 58 +-- tests/optimizers/test_capoeira.py | 574 +++++++++++++++++++++-- tests/predictors/test_base_predictor.py | 24 + tests/tasks/test_base_task.py | 146 ++++++ tests/tasks/test_multi_objective_task.py | 75 +++ tests/utils/test_prompt_creation.py | 19 + 17 files changed, 1409 insertions(+), 196 deletions(-) create mode 100644 tests/llms/test_base_llm.py create mode 100644 tests/mocks/dummy_config.py create mode 100644 tests/optimizers/test_base_optimizer.py create mode 100644 tests/tasks/test_base_task.py create mode 100644 tests/tasks/test_multi_objective_task.py diff --git a/.coverage b/.coverage index 35124674714a2ecbc725f89aa85e571a7946542f..e65bb0277db3ed62b8694bc1bfcdf4f8524131a0 100644 GIT binary patch delta 4157 zcmeHKTU1k58s0lOCzo?_&PhNLsekBpL#RsDNBVQL&c5va6kk zwkW&Hi0z0@yKEO;P-&}RRqV_ZZHqW{s;#t?PPvFE(4ta6$n2b=i`MaBW*+8c^0N1T zKiU8O_y7Inhluh+#K|CS&`il+B?S_p_>_2^7>KGw!DKVJopd7_h-?A~-w?V8P6;vv zBs>n6f}g>@y?Y}BdVi!&Scv2|Xq2Fnk$wcI2t~St#mK*(5oF9pCv}-{mZg<~Vdzlo zbZJ&*QDILq=ZRWJ=;h{>jF8CKRH6Us6rpHk8Ij2gMKgo(=$8^|3JV(;EaVl7b-i4G6wCUZnX zf7Eer4fI2GD>LD2q)197XZ=P8XN2@!LEzK1S@8);`T(M)RJAP|&=a|)WJ3?^BpmIz z<0irX$(`5U+>7r)ub6b(NZ>TsY>O7pgswQwqRBKUNSAy9r{PzbEe)=u^OM*(6|yO* zaKaH+!RctvkKy?0ut>i|1|5;%wr*-BJaKXq>5GsmUO_4i1kQrZ_qBT>bfM1#Zo}p zS+9_aw%k9^+t-?XB0hC#s;f+gWlp+;4N>uDmgP}eOewWr zuBf>KISt->Xl}H$4Q=b$G`x$mi@^+^m<03Aotp2if7EvUqw70e)tDEpF1YhvTUg11 z2MG_Pc`>cu+^=)*HPQvc&@mc|%^}hKEu{zZhaZ-{_5D!k!J2S$S((^z0TzWU*q8nD zmFCXs&a9u~eQL5p(5BBwCJIYjqN1w1FK1ufGU5=4$vKg^qm?VnZ@yi&Y`G@kh~}!) zE&{WCA`VP!9&LUz-nu`2wAo7@j(Ib}4@^XiU%ub`=DY4FC5(h-X@y_dtr|^FYfFu} zvAbq$bYBntf(z@%M!wP9d{k3YlQ%I^GxkSo&-h0})`su5x7U@Zd&YBHtPRGNv6NLY z13fjZ)^F>&#tTNw{oA|7zZinv^KqnQ^K0KTzg>2ZdTp|#(PAH00I%q#NB6|YESt+b91YzUUP zh4`K8Os=x-4~Z1d#k73x!*(hdZ&_cG5I(%G zWLuk36^Mnwfk{Cbuz@pyPbI;gtE_%$A{zoppV?+P@7QzV}7q zTQ~B$8@>&NV67=^NBpwa!}F?3`mFkM1H0_~u=y;%+N&QLi@QwY<_pb@=L_|!*_azV zyEelF?h6NxgQCI{BAK@C;TF$XSg>$bF8Ro$wc>E#ZriP4WnZ0ONB#Q0pGZyRn$pC+ zSVitTDwyMgH+r8t^dYNJFz5{)j^2v|zl?c<^06+i{Gnk2f;k#ZTQL~415GO08*4jm z@~$O!$218jFH9-DlE~fVCh<(0!T4(p4R)o7meE-Lt-`#cSlMb#G5+{DuB=Z{5_Vr=+yx*XR3qVEg4BX!Gkc z@N0@M2P%5q*^JVBKA$7Bcf*WXZn>sJE6kW^uTF(K`d4zUnCan~n+dLQ9ozhVp%@p; zaB)dm^h#47JlV!NW75sJHW+A|7GJt35IbR7?sN+>(`|awYo~7oAA=}AH%CmyIQku< zh2Viz48SEzul3#UVquff-fSd+fL37zN7ZON5OieNDKX)rY<%AlK>5TN3C6eE6cK+uS^nbyWiydz-5LX=?X`u*vM?XotxV zJH3hmfB@wWtN>iF<{GT`{3!*ND9wu4+0a8YPmYOP`RS6(22Gy^Jo@hXaWH1Rc=W<~ zYh#-=-P~8$=5%|v_2BT??ivtN&|TLNS7dVFIBey?6+7&n5Q8S`YN43L3^}_KQjECc` z$|kb$Y%IHgormAn_ZdGCb^&8!`9-JW7p<0GG#Y-1iQyNux<38jyK~o>fC^Y30fnx@ z%75Lp$&I803{CCb|7F)!Bqd47Vd==gRizu`m% zOrad8Y!ePL=Rfb`@~M4XLVy%x0|yzsvfuV|(oym$$U~X8)z9qd)FAxw>7MTLte%e5 z6<9qEg%0GRX#EEST??`>!&sl(c8U=B!50B{1W_H{{%6Gs(I;Y)G+*i{xh%>OvE+y3 z8e)t%Lc|I03N6AgL5HARFbm#*FV-h^ctSM)M?H+iXk;bSVbFw;!wL~{?{p%y{K?sl z*1eN{7TC}r*6`SbwiM8ifiY-LX9iUBxC&b`FJx|c0urM!Hqny{(uoQUpR~X(IjZhb zz$jiju}3bMya&Q4P?R+_h*sBpIhvvjbKLUlc2GBlZ}nEjZgBtw#j3SXQ^&9v;2{QOd5 zR$;;VoDJA?E-!d&D&t*m+LVta!T7w?yhBS22Pqiz6n-OBQ?Vm-n*@FNz3*V4skn3O zcR=$AShPXBSM>_CrOyBz>i^caL4Xl$=DMZm<^GE{)c4FAl15a_=Woax102g>#ODw8 z0qzX9k+Var*7N@}(Enlv8ks2yQAtFAJH%<(+w4B}IcAV~n|WT=FRPYmF;sqqo=bg6 zZKq~PFG=&Hw4`2=DG}k9`DT z-{-XvmU}X|4u;#7)wbFckzrUr)bf)IF0s|7bh0*+=TF+IM22}bkGzKJsq%b26bs#L z^X`-S6eWdV-t5ObrD)|`O#9ks=lsH~ye&b$tJ9UikLA1%#gsx#48kJs$F3+_wbC$< cpW7|v9P}guuw2IvytrFf4smcT$v26A0*I2iCIA2c delta 3576 zcmeHKdr(x@89(Rly>}n?waasr*CG#PQ639|yo5C&0wP3+ydp1!Aif}IQWL8SwXJ_R z3F>jON)$<=NoSkHG&CAW5YopWNsWn)ZG%-26@w&^5H(>Fx&7`2B2Jr5XF8qfAKjVV z`M%%z&Uenazx(a)JB@90V;lWOWLo4L&BvO@G#=`1^(yu@OIS10&onWRj`r`~oy3-f)2PWxmRXmOx6eszi5c7k#T8an*4I_nHdIE`KUEk; zOzClCbg36fv+{TjS!%7sv&m6wArm?kuPHCD-B7!!yuAJOMaM9k69@x=kSL8`UA?Zp zx?v;a$v2VioH){!%aZY7Z+2Ax#4I6SURgj^q#KEpuE+kcP!Z$jnhUr+Lc}a-Tg;H2 zbPXH?HRUTQiSSS`_JLhI2o~qPvT|c}d3{55Rqd)xbq%o2x|tAJIvL54D-ZcIGXPP} zrOf4osl6(DGEVNvKq`iL_F-=b3~>d1y;#6rgvm(L#BGjVySaMPlVoYeEbK|@GYS|F z2umT>ti_~z#K0E2gB$OfHSeZp`{m3Kw&*Z|k0c$Y{zz|@+{F{(YRwsqCsm7&QU>Hi zE$S+EFMEd#=kj&?bfelq?GDi(bl?J^igO4e-@&W+N^O`fnfV#xC%sLdlxo#I%y~`C z?%m0xaWLL)lJE?LGaOEEVY)p+!W=)1Gcy_ZyaoTX{u+3o{;tWDg%1K} zYQ}Fh4V*zLci7q(N$_^hXao$oP)9X~u1Hu5MO2bP$k+`t%@(ZB?R3G3Fa-x&EM z{6czfQu?WNIwOHVlA$b}{5HeDc$U(@ijw)G&slO}#EtPNp@Ene$``kvDeh&+8-pAT zw1~=cAQwv9->jfbwq)HWNV+VYmU^WlQnR#ODwZ;&6v-^{;!W|2cwT&8w2M2%$HkRm zzL+XnL^n|-IE5QRztAUi2z!JzLW*DzG`!4z#Sifp_+$J*eh0sj&*o!!A70{Q?i$y@ zy~OS2wsMbgC0rU8$O#yHC)9$};J#TI8?4p9qfM%E-p5Ij8l~rh<>%?E>?*6S`3}3oQA*5+MwO(GU z_6Sp^hB;a;)uncfjXAeAAMa{ryypNKHOEm=jc%z%UqkhEFSC5wOSZ9wkl8SjJ$nuF zU3HoaS)Jbt6U>HLj%=U$d(glO#6`aRx zM-}QIB_gQBs-C47v-+_)5R8Jr8%&#APK}13JGyOKL^L5P@Ah{Xp`O$fu;3#iM1`9m zxX9F*5jo|$PQ#i1X}Rf&Fp`1?65nzt`pq&`whzTlgD*Hog9q=6^R-{?di zCBP>y{g*w%c-Qk-oT>OT9Z@K)=b`tyDfJBS^)vdhUKaQi8kUy664ED&v7v!RQ9A#m&Jn@K{if(}vI`+gq>8cRY0l81XVJ3d1Hg zS`T)W{>>*U$8AM0uda0i9+R6-bjg;^cDZZf&zFvEWu1+3Yu78eNF5)uY}DW@NifGG zHF$2HP@^8XKw?DzM6r_$JR6`=qk2L$$I^|%#~6~>JR9rD%I3Ja9Y}*O;zk?8n`nq$ zM0Y+>EC-;0L#O(GsxYaa{K97Sq_;Uml`NJK^{%Jd4f=~hds*w@=oKfC7iz(@hbl_- z;*SXp^8qfnCpP~Dp_$YaH)`I0C^ynHxl#9El}lv)wdAMjM9KlFCn0epGEZ2W9@}YD3NxuK^0N_tbFvj*=kU?l@mUvmvVaNIFP# z10g}j)vLAZwMNZJwL9C%ma>TXHIq$$N$;W)RfDQss`;AI-|W6Zku8@?m}Hk6g6!H4 zX|M-Q0tt&(S}A*l4s)b_P{N61U~oB;P(e36pgg(n2ew@?G7DTHs6cuS-o1pY0ux_G zH>J9w0170+P>$WkVKp;<1>FRxLEPJh^d#e&#Khi*9Jwap7?<2DT^z*NpXD)!Eg;>G zFwr@5Q-Eu-JY>L8#Be>IiJaOtk%x5liz3$Hh<}i9#%p@hfc@@k1hA2;AASj&h~dVo z_gW2FX*HDo#U=GSN(^ME+LLUgBeTXZJ) zv9|V;` z6JdU28N2_$}@D?wPvgZmo2S?gl1)Rq|906%uN13V^4$p(* List[Promp ) config.prompts = [Prompt(p) for p in initial_prompts] - if config.optimizer in {"capo", "capoeira"} and ( - config.eval_strategy is None or "block" not in config.eval_strategy - ): - logger.warning("📌 CAPO-style optimizers require block evaluation strategy. Setting it to 'sequential_block'.") - config.eval_strategy = "sequential_block" - task = get_task(df, config, judge_llm=llm) optimizer = get_optimizer( predictor=predictor, @@ -287,7 +281,7 @@ def get_exemplar_selector( raise ValueError(f"Unknown exemplar selector: {name}") -def get_predictor(downstream_llm=None, type: "PredictorType" = "marker", *args, **kwargs) -> "BasePredictor": +def get_predictor(downstream_llm: Optional["BaseLLM"] = None, type: "PredictorType" = "marker", *args, **kwargs) -> "BasePredictor": """Create and return a predictor instance. This function supports three types of predictors: @@ -305,6 +299,7 @@ def get_predictor(downstream_llm=None, type: "PredictorType" = "marker", *args, Returns: An instance of FirstOccurrencePredictor or MarkerBasedPredictor. """ + assert downstream_llm is not None, "downstream_llm must be provided to create a predictor." if type == "first_occurrence": return FirstOccurrencePredictor(downstream_llm, *args, **kwargs) elif type == "marker": diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 88450484..7fd05144 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -1,7 +1,5 @@ """Implementation of the Capoeira (Multi-Objective CAPO) optimizer.""" -from __future__ import annotations - import random import numpy as np @@ -80,15 +78,15 @@ def __init__( super().__init__(predictor, task, initial_prompts, callbacks, config) - self.incumbents: List[Prompt] = [] - self.challengers: List[Prompt] = [] - self.crossover_template = self._initialize_meta_template(crossover_template or CAPO_CROSSOVER_TEMPLATE) self.mutation_template = self._initialize_meta_template(mutation_template or CAPO_MUTATION_TEMPLATE) self.token_counter = get_token_counter(self.downstream_llm) self.df_few_shots = df_few_shots if df_few_shots is not None else task.pop_datapoints(frac=0.1) - self.population_size = len(self.prompts) + self.incumbents: List[Prompt] = self.prompts + self.challengers: List[Prompt] = [] + self.population_size = len(self.prompts) + if "block" not in self.task.eval_strategy: logger.warning( f"ℹ️ CAPO requires 'block' in the eval_strategy, but got {self.task.eval_strategy}. Setting eval_strategy to 'sequential_block'." @@ -102,10 +100,6 @@ def __init__( self.target_begin_marker = "" self.target_end_marker = "" - @property - def prompts(self) -> List[Prompt]: - return self.incumbents + self.challengers - def _pre_optimization_loop(self) -> None: population: List[Prompt] = [] for prompt in self.prompts: @@ -121,11 +115,12 @@ def _pre_optimization_loop(self) -> None: max(self.token_counter(p.construct_prompt()) for p in population) if population else 1 ) init_result = self.task.evaluate(population, self.predictor) - initial_vectors = self._get_objective_vectors(init_result) #TODO rename + initial_vectors = self._get_objective_vectors(init_result) fronts = self._non_dominated_sort(initial_vectors) self.incumbents = [population[i] for i in fronts[0]] self.challengers = [population[i] for front in fronts[1:] for i in front] + # keep self.prompts as a "view" if base class expects it self.scores = initial_vectors[:, 0].tolist() @@ -139,7 +134,7 @@ def _step(self) -> List[Prompt]: for chal in new_challengers: self._do_intensification(chal) self._advance_one_incumbent() - self._prune_population() + self._select_survivors() # 4) logging scores: incumbents only (optional) if self.incumbents: @@ -152,55 +147,63 @@ def _step(self) -> List[Prompt]: return self.prompts def _do_intensification(self, challenger: Prompt) -> None: - """ - Default MO-CAPO intensification (closest-incumbent comparison): - - evaluate challenger + incumbents on sequential blocks - - maintain running averages (challenger and incumbents) - - early reject if closest incumbent dominates challenger average - - if challenger survives all blocks: promote to incumbents and update front - """ if not self.incumbents: self.incumbents.append(challenger) return - common_block_idx = 0 - while common_block_idx is not None: - common_block_idx = self._sample_common_block(self.incumbents) - self.task.set_block_idx(common_block_idx) # type: ignore + common_blocks = self._get_common_blocks(self.incumbents) - joint_result = self.task.evaluate( - prompts=self.incumbents + [challenger], - predictor=self.predictor - ) + # bootstrap if no common blocks yet + if not common_blocks: + b = random.randrange(self.task.n_blocks) + self.task.set_block_idx(b) + self.task.evaluate(self.incumbents + [challenger], self.predictor) + self.incumbents.append(challenger) + self._update_incumbent_front(blocks={b}) + return - objective_vectors = self._get_objective_vectors(joint_result) - challenger_vec = objective_vectors[-1] - incumbent_vecs = objective_vectors[:-1] + remaining_blocks = set(common_blocks) - closest_inc_vec = self._get_closest_incumbent(challenger_vec, incumbent_vecs) + challenger_mean: Optional[np.ndarray] = None + incumbents_mean: Optional[np.ndarray] = None + t = 0 - if self._is_dominated(challenger_vec, closest_inc_vec): - # challenger loses -> goes to population - self.challengers.append(challenger) - return + fold_vec: Optional[np.ndarray] = None - self.incumbents.append(challenger) - self._update_incumbent_front() + while remaining_blocks: + b = random.choice(tuple(remaining_blocks)) + remaining_blocks.remove(b) - def _sample_common_block(self, prompts: List[Prompt]) -> Optional[int]: - """Sample a block index that has been evaluated by all given prompts. - Returns None if no such block exists.""" - per_prompt = self.task.get_evaluated_blocks(prompts) # Dict[prompt -> Set[int]] - block_sets = list(per_prompt.values()) + # evaluate all incumbents + challenger on THIS block (cache will avoid recompute) + self.task.set_block_idx(b) + res = self.task.evaluate(self.incumbents + [challenger], self.predictor) + vecs = self._get_objective_vectors(res) # per-block vectors, shape (n_inc+1, n_obj) + incumbent_block = vecs[:-1] + challenger_block = vecs[-1] - if not block_sets: - return random.randrange(self.task.n_blocks) + # running means + t += 1 + if challenger_mean is None: + challenger_mean = challenger_block.copy() + incumbents_mean = incumbent_block.copy() + else: + challenger_mean += (challenger_block - challenger_mean) / t + incumbents_mean += (incumbent_block - incumbents_mean) / t # type: ignore - common = set.intersection(*block_sets) - if not common: - return None + # trigger comparisons CAPO/thesis-style + if fold_vec is not None and not self._is_dominated(fold_vec, challenger_mean): + continue + fold_vec = challenger_mean.copy() + + closest_inc = self._get_closest_incumbent(challenger_mean, incumbents_mean) # type: ignore + if self._is_dominated(challenger_mean, closest_inc): + self.challengers.append(challenger) + return + + # survived all common blocks -> admit and update front restricted to common_blocks + self.incumbents.append(challenger) + self._update_incumbent_front(blocks=common_blocks) - return random.choice(tuple(common)) def _get_closest_incumbent( self, challenger_vec: np.ndarray, incumbent_vecs: np.ndarray @@ -212,25 +215,26 @@ def _get_closest_incumbent( rng = max_b - min_b rng[rng == 0] = 1.0 # Avoid div/0 - norm_chal = (challenger_vec - min_b) / rng - norm_incs = (incumbent_vecs - min_b) / rng + challenger_norm = (challenger_vec - min_b) / rng + incumbents_norm = (incumbent_vecs - min_b) / rng - dists = np.linalg.norm(norm_incs - norm_chal, axis=1) + dists = np.linalg.norm(incumbents_norm - challenger_norm, axis=1) idx = int(np.argmin(dists)) return incumbent_vecs[idx] - def _update_incumbent_front(self) -> None: - """ - After adding a challenger that survived a full race, recompute the incumbent Pareto front. - Default behavior: incumbents become front-0 (on current evaluation state), - all other incumbents are demoted to challengers. - """ + def _update_incumbent_front(self, blocks: Optional[set[int]] = None) -> None: if not self.incumbents: return - vecs_result = self.task.evaluate(prompts=self.incumbents, predictor=self.predictor, eval_strategy="evaluated") - vecs = self._get_objective_vectors(vecs_result) + if blocks is None: + res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") + else: + self.task.set_block_idx(list(sorted(blocks))) # sorted for deterministic behaviour + res = self.task.evaluate(self.incumbents, self.predictor) + + vecs = self._get_objective_vectors(res) + fronts = self._non_dominated_sort(vecs) new_incumbents = [self.incumbents[i] for i in fronts[0]] @@ -241,7 +245,6 @@ def _update_incumbent_front(self) -> None: def _get_objective_vectors(self, result) -> np.ndarray: - # If the task is multi-objective, include all objective dimensions, else single objective. if isinstance(self.task, MultiObjectiveTask): agg_scores = np.stack(result.agg_scores, axis=1) # shape: (n_prompts, n_objectives) @@ -254,23 +257,43 @@ def _get_objective_vectors(self, result) -> np.ndarray: cost_scalar = cost_scalar.reshape(-1, 1) return np.hstack([agg_scores, -cost_scalar]) - + def _advance_one_incumbent(self) -> None: - """ - Default MO-CAPO step after processing a challenger: - evaluate one incumbent on one additional sequential block. - """ - # choose least evaluated incumbent - eval_counts = [ - len(self.task.get_evaluated_blocks([inc])) for inc in self.incumbents - ] + if not self.incumbents: + return + + blocks_map = self.task.get_evaluated_blocks(self.incumbents) # Dict[str -> Set[int]] + inc_keys = [str(inc) for inc in self.incumbents] + + # least evaluated incumbents + eval_counts = [len(blocks_map[k]) for k in inc_keys] min_count = min(eval_counts) - candidates = [inc for inc, count in zip(self.incumbents, eval_counts) if count == min_count] - chosen = random.sample(candidates, k=1) - self.task.evaluate(prompts=chosen, predictor=self.predictor) + least = [inc for inc, c in zip(self.incumbents, eval_counts) if c == min_count] + chosen_inc = random.choice(least) + + # union over incumbents + union_blocks: set[int] = set() + for inc in self.incumbents: + union_blocks |= set(blocks_map[str(inc)]) + chosen_blocks = set(blocks_map[str(chosen_inc)]) - def _prune_population(self) -> None: + # gap-first, else brand-new + gap_blocks = union_blocks - chosen_blocks + if gap_blocks: + b = random.choice(tuple(gap_blocks)) + else: + all_blocks = set(range(self.task.n_blocks)) + new_blocks = all_blocks - union_blocks + if not new_blocks: + return + b = random.choice(tuple(new_blocks)) + + self.task.set_block_idx(b) + self.task.evaluate(prompts=[chosen_inc], predictor=self.predictor) + + + def _select_survivors(self) -> None: """ Enforce |incumbents| + |challengers| <= population_size using Pareto logic. @@ -353,8 +376,20 @@ def _prune_population(self) -> None: victim_idx = int(np.argmin(dists)) self.incumbents.pop(victim_idx) + def _get_common_blocks(self, prompts: List[Prompt]) -> set: + """Get the set of block indices that have been evaluated by all given prompts.""" + per_prompt = self.task.get_evaluated_blocks(prompts) # Dict[prompt -> Set[int]] + block_sets = list(per_prompt.values()) - def _non_dominated_sort(self, obj_vectors: np.ndarray) -> List[List[int]]: + if not block_sets: + return set() + + common = set.intersection(*block_sets) + return common + + + @staticmethod + def _non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: """Perform fast non-dominated sorting (NSGA-II) in a vectorized manner.""" n_solutions = obj_vectors.shape[0] diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index d28ffb0f..19a3042f 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -83,7 +83,7 @@ def __init__( # If no y_column is provided, create a dummy y array self.ys = [""] * len(self.xs) - self.block_idx: int = 0 + self.block_idx: int | list[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) @@ -116,9 +116,18 @@ def subsample(self, eval_strategy: Optional["EvalStrategy"] = None) -> Tuple[Lis 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": - 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) + if isinstance(self.block_idx, list): + block_indices: List[int] = [] + for block_id in self.block_idx: + start_idx = block_id * self.n_subsamples + end_idx = min((block_id + 1) * self.n_subsamples, len(self.xs)) + block_indices.extend(range(start_idx, end_idx)) + indices = np.array(sorted(set(block_indices)), dtype=int) + 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}'") @@ -264,7 +273,10 @@ def evaluate( # Record evaluated block for block strategies for prompt in prompts_list: - self.prompt_evaluated_blocks.setdefault(str(prompt), set()).add(self.block_idx) + if isinstance(self.block_idx, list): + self.prompt_evaluated_blocks.setdefault(str(prompt), set()).update(self.block_idx) + else: + self.prompt_evaluated_blocks.setdefault(str(prompt), set()).add(self.block_idx) input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = self._compute_costs( prompts_list, xs, ys, self.seq_cache, predictor @@ -328,6 +340,7 @@ def increment_block_idx(self) -> None: """ 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 @@ -344,14 +357,17 @@ def reset_block_idx(self) -> None: 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 for subsampling (block strategies only).""" + def set_block_idx(self, idx: Union[int, List[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.") - if self.n_blocks > 0: - self.block_idx = idx % self.n_blocks + + if isinstance(idx, list): + assert all(0 <= i < self.n_blocks for i in idx), "Block indices must be integers within valid range" else: - self.block_idx = 0 + assert isinstance(idx, int), "Block index must be an integer" + + self.block_idx = idx def get_evaluated_blocks(self, prompts: List[Prompt]) -> Dict[str, set[int]]: return {str(p): set(self.prompt_evaluated_blocks.get(str(p), set())) for p in prompts} diff --git a/tests/helpers/test_helpers.py b/tests/helpers/test_helpers.py index b467cc2a..e6de09c7 100644 --- a/tests/helpers/test_helpers.py +++ b/tests/helpers/test_helpers.py @@ -8,8 +8,29 @@ from tests.mocks.mock_predictor import MockPredictor from tests.mocks.mock_task import MockTask -from promptolution.helpers import run_evaluation, run_experiment, run_optimization +from promptolution.exemplar_selectors.random_search_selector import RandomSearchSelector +from promptolution.exemplar_selectors.random_selector import RandomSelector +from promptolution.helpers import ( + get_exemplar_selector, + get_llm, + get_optimizer, + get_predictor, + get_task, + run_evaluation, + run_experiment, + run_optimization, +) +from promptolution.optimizers.capo import CAPO +from promptolution.optimizers.capoeira import Capoeira +from promptolution.optimizers.evoprompt_de import EvoPromptDE +from promptolution.optimizers.evoprompt_ga import EvoPromptGA +from promptolution.optimizers.opro import OPRO +from promptolution.predictors.first_occurrence_predictor import FirstOccurrencePredictor +from promptolution.predictors.maker_based_predictor import MarkerBasedPredictor from promptolution.tasks.base_task import EvalResult +from promptolution.tasks.classification_tasks import ClassificationTask +from promptolution.tasks.judge_tasks import JudgeTask +from promptolution.tasks.reward_tasks import RewardTask from promptolution.utils import ExperimentConfig from promptolution.utils.prompt import Prompt @@ -327,3 +348,107 @@ def test_helpers_integration(sample_df, experiment_config): # Verify evaluation was called mock_task.evaluate.assert_called() + + +def test_get_llm_variants(monkeypatch): + def factory(model_name=None, config=None, **kwargs): + created['name'] = model_name or kwargs.get("model_id") + created['config'] = config + return MockLLM() + + created = {} + + monkeypatch.setattr("promptolution.helpers.LocalLLM", factory) + monkeypatch.setattr("promptolution.helpers.VLLM", factory) + monkeypatch.setattr("promptolution.helpers.APILLM", factory) + + cfg = ExperimentConfig() + cfg.model_id = "local-foo" + res = get_llm(config=cfg) + assert isinstance(res, MockLLM) + assert created['name'] == "foo" + + cfg.model_id = "vllm-bar" + res = get_llm(config=cfg) + assert created['name'] == "bar" + + cfg.model_id = "api-model" + res = get_llm(config=cfg) + assert created['name'] == "api-model" + + with pytest.raises(ValueError): + get_llm() + + +def test_get_task_variants(sample_df): + cfg = ExperimentConfig() + cfg.task_type = "reward" + task = get_task(sample_df, cfg, reward_function=lambda _: 1.0) + + assert isinstance(task, RewardTask) + + cfg.task_type = "judge" + judge_task = get_task(sample_df, cfg, judge_llm=MockLLM()) + + assert isinstance(judge_task, JudgeTask) + + cfg.task_type = "classification" + cls_task = get_task(sample_df, cfg) + + assert isinstance(cls_task, ClassificationTask) + + +def test_get_optimizer_variants(): + pred = MockPredictor(llm=MockLLM()) + task = MockTask() + cfg = ExperimentConfig() + + opt = get_optimizer(pred, MockLLM(), task, optimizer="capo", config=cfg) + + assert isinstance(opt, CAPO) + + opt2 = get_optimizer(pred, MockLLM(), task, optimizer="capoeira", config=cfg) + + assert isinstance(opt2, Capoeira) + + opt3 = get_optimizer(pred, MockLLM(), task, optimizer="evopromptde", config=cfg) + + assert isinstance(opt3, EvoPromptDE) + + opt4 = get_optimizer(pred, MockLLM(), task, optimizer="evopromptga", config=cfg) + + assert isinstance(opt4, EvoPromptGA) + + opt5 = get_optimizer(pred, MockLLM(), task, optimizer="opro", config=cfg) + + assert isinstance(opt5, OPRO) + + with pytest.raises(ValueError): + get_optimizer(pred, MockLLM(), task, optimizer="unknown", config=cfg) + + +def test_get_exemplar_selector_variants(): + task = MockTask() + pred = MockPredictor() + + sel = get_exemplar_selector("random", task, pred) + assert isinstance(sel, RandomSelector) + + sel2 = get_exemplar_selector("random_search", task, pred) + assert isinstance(sel2, RandomSearchSelector) + + with pytest.raises(ValueError): + get_exemplar_selector("nope", task, pred) + + +def test_get_predictor_variants(): + llm = MockLLM() + + p1 = get_predictor(llm, type="first_occurrence", classes=["a", "b"]) + assert isinstance(p1, FirstOccurrencePredictor) + + p2 = get_predictor(llm, type="marker") + assert isinstance(p2, MarkerBasedPredictor) + + with pytest.raises(ValueError): + get_predictor(llm, type="bad") diff --git a/tests/llms/test_api_llm.py b/tests/llms/test_api_llm.py index 2d1cb5af..b8fc3fa4 100644 --- a/tests/llms/test_api_llm.py +++ b/tests/llms/test_api_llm.py @@ -1,8 +1,30 @@ -from unittest.mock import MagicMock, patch +import asyncio +from concurrent.futures import TimeoutError as FuturesTimeout +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest from promptolution.llms import APILLM +class _FakeSem: + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + + +def _make_api_stub(**attrs): + """Create an APILLM instance via __new__ with provided attributes.""" + api = APILLM.__new__(APILLM) + api._call_kwargs = {} + for key, value in attrs.items(): + setattr(api, key, value) + return api + + def test_api_llm_initialization(): """Test that APILLM initializes correctly.""" # Create patches for all dependencies @@ -34,3 +56,117 @@ def test_api_llm_initialization(): assert api_llm.api_url == "https://api.example.com" assert api_llm.model_id == "gpt-4" assert api_llm.max_concurrent_calls == 10 + + +def test_ainvoke_once_uses_client_and_timeout(monkeypatch): + response = SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))]) + create = AsyncMock(return_value=response) + client = SimpleNamespace(chat=SimpleNamespace(completions=SimpleNamespace(create=create))) + + api = _make_api_stub(model_id="m", max_tokens=11, call_timeout_s=0.5, _sem=_FakeSem(), client=client) + + out = asyncio.run(api._ainvoke_once("prompt", "system")) + + assert out is response + assert create.await_count == 1 + kwargs = create.await_args.kwargs + assert kwargs["model"] == "m" + assert kwargs["messages"][0]["role"] == "system" + assert kwargs["max_tokens"] == 11 + + +def test_ainvoke_with_retries_recovers(monkeypatch): + good = SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="done"))]) + api = _make_api_stub(max_retries=2, retry_base_delay_s=0) + api._ainvoke_once = AsyncMock(side_effect=[Exception("fail"), good]) + async def _sleep(_): + return None + + monkeypatch.setattr("promptolution.llms.api_llm.asyncio.sleep", _sleep) + + out = asyncio.run(api._ainvoke_with_retries("p", "s")) + + assert out == "done" + assert api._ainvoke_once.await_count == 2 + + +def test_ainvoke_with_retries_exhausts(monkeypatch): + api = _make_api_stub(max_retries=1, retry_base_delay_s=0) + api._ainvoke_once = AsyncMock(side_effect=[Exception("boom"), Exception("boom2")]) + async def _sleep(_): + return None + + monkeypatch.setattr("promptolution.llms.api_llm.asyncio.sleep", _sleep) + + with pytest.raises(Exception) as excinfo: + asyncio.run(api._ainvoke_with_retries("p", "s")) + + assert "boom2" in str(excinfo.value) + assert api._ainvoke_once.await_count == 2 + + +def test_aget_batch_success(monkeypatch): + api = _make_api_stub(gather_timeout_s=1) + api._ainvoke_with_retries = AsyncMock(side_effect=["a", "b"]) + monkeypatch.setattr("promptolution.llms.api_llm.asyncio.wait_for", asyncio.wait_for) + + outs = asyncio.run(api._aget_batch(["p1", "p2"], ["s1", "s2"])) + + assert outs == ["a", "b"] + assert api._ainvoke_with_retries.await_count == 2 + + +def test_aget_batch_raises_on_failure(monkeypatch): + api = _make_api_stub(gather_timeout_s=1) + api._ainvoke_with_retries = AsyncMock(side_effect=["ok", Exception("boom")]) + monkeypatch.setattr("promptolution.llms.api_llm.asyncio.wait_for", asyncio.wait_for) + + with pytest.raises(RuntimeError): + asyncio.run(api._aget_batch(["p1", "p2"], ["s1", "s2"])) + + +def test_get_response_success(monkeypatch): + api = _make_api_stub(gather_timeout_s=1) + api._aget_batch = AsyncMock() + + class _Future: + def __init__(self, value): + self.value = value + self.cancelled = False + + def result(self, timeout=None): + return self.value + + def cancel(self): + self.cancelled = True + + fut = _Future(["r1", "r2"]) + api._submit = MagicMock(return_value=fut) + + out = api._get_response(["p1", "p2"], ["s1", "s2"]) + + assert out == ["r1", "r2"] + api._submit.assert_called_once() + assert fut.cancelled is False + + +def test_get_response_times_out(): + api = _make_api_stub(gather_timeout_s=1) + + class _Future: + def __init__(self): + self.cancelled = False + + def result(self, timeout=None): + raise FuturesTimeout() + + def cancel(self): + self.cancelled = True + + fut = _Future() + api._submit = MagicMock(return_value=fut) + + with pytest.raises(TimeoutError): + api._get_response(["p"], ["s"]) + + assert fut.cancelled is True diff --git a/tests/llms/test_base_llm.py b/tests/llms/test_base_llm.py new file mode 100644 index 00000000..0a3d7745 --- /dev/null +++ b/tests/llms/test_base_llm.py @@ -0,0 +1,35 @@ +from tests.mocks.dummy_config import DummyConfig +from tests.mocks.mock_llm import MockLLM + + +def test_base_llm_token_count_and_reset(): + llm = MockLLM() + llm.update_token_count(["a b"], ["c d e"]) + counts = llm.get_token_count() + assert counts["input_tokens"] == 2 + assert counts["output_tokens"] == 3 + + llm.reset_token_count() + assert llm.get_token_count()["total_tokens"] == 0 + + +def test_base_llm_default_and_list_system_prompts(): + llm = MockLLM() + res_single = llm.get_response("hello") + assert res_single == ["Mock response for: hello"] + + res_multi = llm.get_response(["p1", "p2"], system_prompts=["s1", "s2"]) + assert res_multi == ["Mock response for: p1", "Mock response for: p2"] + + +def test_base_llm_config_applied(): + cfg = DummyConfig() + llm = MockLLM(predetermined_responses=["r1"], add_prompt_tags=False, config=cfg) + assert cfg.applied is True + assert getattr(llm, "applied") is True + + +def test_base_llm_set_generation_seed(): + llm = MockLLM() + llm.set_generation_seed(123) + assert llm._generation_seed == 123 diff --git a/tests/mocks/dummy_config.py b/tests/mocks/dummy_config.py new file mode 100644 index 00000000..80b6d81a --- /dev/null +++ b/tests/mocks/dummy_config.py @@ -0,0 +1,15 @@ +class DummyConfig: + """Lightweight config stub used across tests.""" + + def __init__(self, task_description=None): + self.applied = False + self.validated = False + self.task_description = task_description + + def apply_to(self, obj): + self.applied = True + obj.config_applied = True + obj.applied = True + + def validate(self): + self.validated = True diff --git a/tests/mocks/mock_task.py b/tests/mocks/mock_task.py index 9aeb46c7..9b70f30a 100644 --- a/tests/mocks/mock_task.py +++ b/tests/mocks/mock_task.py @@ -2,6 +2,7 @@ from unittest.mock import MagicMock +import math import pandas as pd from typing import List @@ -16,35 +17,65 @@ class MockTask(BaseTask): actual data or model inference. """ - def __init__(self, predetermined_scores=None): - """Initialize the MockTask with optional predetermined scores. + def __init__( + self, + predetermined_scores=None, + *, + df: pd.DataFrame | None = None, + n_subsamples: int = 1, + eval_strategy: str = "full", + n_blocks: int | None = None, + block_idx: int | list[int] = 0, + eval_blocks: dict[str, set[int]] | None = None, + task_description: str = "Mock classification task", + evaluate_fn=None, + config=None, + ): + """Initialize the MockTask with optional overrides for task settings. Args: - predetermined_scores: Dictionary mapping prompts to scores, - or a list of scores to return in sequence, or a function - that generates scores based on prompts. + predetermined_scores: Dict/list/callable for score generation used by _evaluate. + eval_strategy: Eval strategy to expose (defaults to "full"). + n_blocks: Number of blocks to report. + block_idx: Current block index (int or list). + eval_blocks: Mapping prompt->set of evaluated blocks for selection logic. + task_description: Description to attach to the task. + evaluate_fn: Optional callable to replace evaluate entirely for tests. """ + base_df = df if df is not None else pd.DataFrame( + {"x": ["Sample text 1", "Sample text 2", "Sample text 3"], "y": ["positive", "negative", "neutral"]} + ) + super().__init__( - df=pd.DataFrame( - {"x": ["Sample text 1", "Sample text 2", "Sample text 3"], "y": ["positive", "negative", "neutral"]} - ), + df=base_df, x_column="x", y_column="y", + eval_strategy=eval_strategy, + n_subsamples=n_subsamples, + config=config, ) self.predetermined_scores = predetermined_scores or {} self.call_history = [] self.score_index = 0 + self.eval_blocks: dict[str, set[int]] = eval_blocks or {} - self.x_column = "x" - self.y_column = "y" - # Default attributes similar to ClassificationTask - self.task_description = "Mock classification task" + self.task_description = task_description self.classes = ["positive", "neutral", "negative"] self.initial_prompts = ["Classify:", "Determine:"] - self.n_blocks = 10 - self.increment_block_idx = MagicMock() - self.reset_block_idx = MagicMock() + # Allow tests to control block metadata + self.n_blocks = n_blocks if n_blocks is not None else max(1, math.ceil(len(self.xs) / self.n_subsamples)) + self.block_idx = block_idx + + # Track block operations for assertions while keeping original behavior + self._reset_block_idx_impl = super().reset_block_idx + self.reset_block_idx = MagicMock(side_effect=self._reset_block_idx_impl) + self._increment_block_idx_impl = super().increment_block_idx + self.increment_block_idx = MagicMock(side_effect=self._increment_block_idx_impl) + + if evaluate_fn is not None: + # Replace evaluate for bespoke test logic + self.evaluate = evaluate_fn # type: ignore[assignment] def _evaluate(self, xs: List[str], ys: List[str], preds: List[str], **kwargs) -> List[float]: """Calculate the score for a single prediction. @@ -60,9 +91,20 @@ def _evaluate(self, xs: List[str], ys: List[str], preds: List[str], **kwargs) -> if isinstance(self.predetermined_scores, dict): return [self.predetermined_scores.get(pred, 0) for pred in preds] elif isinstance(self.predetermined_scores, list): - self.score_index += 1 - return self.predetermined_scores + if not self.predetermined_scores: + return [0 for _ in preds] + + scores = [ + self.predetermined_scores[(self.score_index + i) % len(self.predetermined_scores)] + for i in range(len(preds)) + ] + self.score_index += len(preds) + return scores elif callable(self.predetermined_scores): return self.predetermined_scores(xs) else: return [len(pred) for pred in preds] + + def get_evaluated_blocks(self, prompts): + """Return per-prompt evaluated block sets for testing selection logic.""" + return {str(p): set(self.prompt_evaluated_blocks.get(str(p), set())) for p in prompts} diff --git a/tests/optimizers/test_base_optimizer.py b/tests/optimizers/test_base_optimizer.py new file mode 100644 index 00000000..1b0c2f7f --- /dev/null +++ b/tests/optimizers/test_base_optimizer.py @@ -0,0 +1,86 @@ +import pytest + +from tests.mocks.dummy_config import DummyConfig +from tests.mocks.mock_predictor import MockPredictor +from tests.mocks.mock_task import MockTask + +from promptolution.optimizers.base_optimizer import BaseOptimizer +from promptolution.utils.callbacks import BaseCallback + + +class SimpleOptimizer(BaseOptimizer): + def __init__(self, predictor, task, **kwargs): + super().__init__(predictor=predictor, task=task, initial_prompts=["p1", "p2"], **kwargs) + self.prepared = False + self.steps = 0 + + def _pre_optimization_loop(self): + self.prepared = True + + def _step(self): + self.steps += 1 + return self.prompts + + +class FailingOptimizer(SimpleOptimizer): + def _step(self): + raise RuntimeError("boom") + + def _on_train_end(self): + self.cleaned = True + return None + + +class Stopper(BaseCallback): + def on_step_end(self, optimizer): + # stop after first step to exercise callback stop path + return False + + def on_train_end(self, optimizer): + optimizer.stopped = True + return True + + +@pytest.fixture +def predictor(): + return MockPredictor() + + +@pytest.fixture +def task(): + return MockTask() + + +def test_base_optimizer_runs_and_calls_callbacks(predictor: MockPredictor, task: MockTask): + opt = SimpleOptimizer(predictor=predictor, task=task) + opt.callbacks = [Stopper()] + opt.optimize(3) + + assert opt.prepared is True + assert opt.steps == 1 + assert getattr(opt, "stopped", False) is True + + +def test_base_optimizer_stops_on_exception(predictor: MockPredictor, task: MockTask): + opt = FailingOptimizer(predictor=predictor, task=task) + opt.optimize(2) + + assert opt.prepared is True + assert getattr(opt, "cleaned", False) is True + + +def test_base_optimizer_no_callbacks_continues(predictor: MockPredictor, task: MockTask): + opt = SimpleOptimizer(predictor=predictor, task=task) + opt.optimize(2) + assert opt.steps == 2 + + +def test_base_optimizer_config_validate_and_template(predictor: MockPredictor, task: MockTask): + cfg = DummyConfig(task_description="override desc") + opt = SimpleOptimizer(predictor=predictor, task=task, config=cfg) + opt.optimize(1) + assert cfg.validated is True + + templ = opt._initialize_meta_template("Hi ") + assert "override desc" in templ + assert getattr(opt, "config_applied", True) diff --git a/tests/optimizers/test_capo.py b/tests/optimizers/test_capo.py index 5fff4cb0..da227088 100644 --- a/tests/optimizers/test_capo.py +++ b/tests/optimizers/test_capo.py @@ -123,15 +123,7 @@ def test_create_few_shots(mock_meta_llm, mock_predictor, initial_prompts, mock_t few_shot_examples = build_few_shot_examples( instruction="Classify the sentiment of the text.", num_examples=2, - df_few_shots=mock_df, - x_column=mock_task.x_column, - y_column=mock_task.y_column, - predictor=mock_predictor, - fewshot_template=CAPO_FEWSHOT_TEMPLATE, - target_begin_marker=optimizer.target_begin_marker, - target_end_marker=optimizer.target_end_marker, - check_fs_accuracy=True, - create_fs_reasoning=True, + optimizer=optimizer, ) # Verify results @@ -141,15 +133,7 @@ def test_create_few_shots(mock_meta_llm, mock_predictor, initial_prompts, mock_t few_shot_examples = build_few_shot_examples( instruction="Classify the sentiment of the text.", num_examples=0, - df_few_shots=mock_df, - x_column=mock_task.x_column, - y_column=mock_task.y_column, - predictor=mock_predictor, - fewshot_template=CAPO_FEWSHOT_TEMPLATE, - target_begin_marker=optimizer.target_begin_marker, - target_end_marker=optimizer.target_end_marker, - check_fs_accuracy=True, - create_fs_reasoning=True, + optimizer=optimizer, ) assert len(few_shot_examples) == 0 @@ -167,9 +151,7 @@ def test_crossover(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mo offsprings = perform_crossover( [Prompt("Instruction 1", ["Example 1"]), Prompt("Instruction 2", ["Example 2"])], - optimizer.crossovers_per_iter, - optimizer.crossover_template, - optimizer.meta_llm, + optimizer=optimizer, ) assert len(offsprings) == 5 @@ -185,20 +167,7 @@ def test_mutate(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_ mutated = perform_mutation( offsprings=[Prompt("Instruction 1", ["Example 1"]), Prompt("Instruction 2", ["Example 2"])], - mutation_template=optimizer.mutation_template, - upper_shots=optimizer.upper_shots, - meta_llm=optimizer.meta_llm, - few_shot_kwargs=dict( - df_few_shots=mock_df, - x_column=mock_task.x_column, - y_column=mock_task.y_column, - predictor=mock_predictor, - fewshot_template=CAPO_FEWSHOT_TEMPLATE, - target_begin_marker=optimizer.target_begin_marker, - target_end_marker=optimizer.target_end_marker, - check_fs_accuracy=True, - create_fs_reasoning=True, - ), + optimizer=optimizer, ) assert len(mutated) == 2 @@ -222,7 +191,7 @@ def test_do_racing(mock_meta_llm, mock_predictor, initial_prompts, mock_df): assert "better instruction" in survivors[0].instruction assert mock_task.reset_block_idx.call_count == 2 - assert mock_task.increment_block_idx.call_count == 3 + assert mock_task.increment_block_idx.call_count == 2 def test_capo_crossover_prompt(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): @@ -237,7 +206,7 @@ def test_capo_crossover_prompt(mock_meta_llm, mock_predictor, initial_prompts, m mother = Prompt("Classify the sentiment of the text.", ["Input: I love this! Output: Positive"]) father = Prompt("Determine if the review is positive or negative.", ["Input: This is terrible. Output: Negative"]) - perform_crossover([mother, father], optimizer.crossovers_per_iter, optimizer.crossover_template, optimizer.meta_llm) + perform_crossover([mother, father], optimizer=optimizer) full_task_desc = mock_task.task_description + "\n" + optimizer.predictor.extraction_description @@ -269,20 +238,7 @@ def test_capo_mutate_prompt(mock_meta_llm, mock_predictor, initial_prompts, mock parent = Prompt("Classify the sentiment of the text.", ["Input: I love this! Output: Positive"]) perform_mutation( offsprings=[parent], - mutation_template=optimizer.mutation_template, - upper_shots=optimizer.upper_shots, - meta_llm=optimizer.meta_llm, - few_shot_kwargs=dict( - df_few_shots=mock_df, - x_column=mock_task.x_column, - y_column=mock_task.y_column, - predictor=mock_predictor, - fewshot_template=CAPO_FEWSHOT_TEMPLATE, - target_begin_marker=optimizer.target_begin_marker, - target_end_marker=optimizer.target_end_marker, - check_fs_accuracy=True, - create_fs_reasoning=True, - ), + optimizer=optimizer, ) expected_meta_prompt = CAPO_MUTATION_TEMPLATE.replace("", parent.instruction).replace( diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py index 45602afe..9ab5cf5c 100644 --- a/tests/optimizers/test_capoeira.py +++ b/tests/optimizers/test_capoeira.py @@ -1,13 +1,15 @@ -from unittest.mock import MagicMock, patch +from unittest.mock import patch import numpy as np import pandas as pd from promptolution.optimizers.capoeira import Capoeira from promptolution.tasks.base_task import EvalResult +from promptolution.tasks.multi_objective_task import MultiObjectiveEvalResult, MultiObjectiveTask from promptolution.utils.capo_utils import perform_crossover, perform_mutation from promptolution.utils.prompt import Prompt -from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE +from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_MUTATION_TEMPLATE +from tests.mocks.mock_task import MockTask def test_capoeira_initialization(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): @@ -41,7 +43,7 @@ def test_capoeira_initialize_population(mock_meta_llm, mock_predictor, initial_p assert all(isinstance(p, Prompt) for p in population) -def test_capoeira_selection_prefers_better_score(mock_meta_llm, mock_predictor, mock_task, mock_df): +def test_capoeira_objective_vectors_and_sort(mock_meta_llm, mock_predictor, mock_task, mock_df): optimizer = Capoeira( predictor=mock_predictor, task=mock_task, @@ -49,24 +51,27 @@ def test_capoeira_selection_prefers_better_score(mock_meta_llm, mock_predictor, initial_prompts=["short"], df_few_shots=mock_df, ) - candidates = [Prompt("short"), Prompt("longer prompt")] - optimizer.task.evaluate = MagicMock( - return_value=EvalResult( - scores=np.array([[0.4], [0.9]], dtype=float), - agg_scores=np.array([0.4, 0.9], dtype=float), - sequences=np.array([["s1"], ["s2"]], dtype=object), - input_tokens=np.array([[1.0], [1.0]], dtype=float), - output_tokens=np.array([[0.0], [0.0]], dtype=float), - agg_input_tokens=np.array([1.0, 1.0], dtype=float), - agg_output_tokens=np.array([0.0, 0.0], dtype=float), - ) + + result = EvalResult( + scores=np.array([[0.4], [0.9]], dtype=float), + agg_scores=np.array([0.4, 0.9], dtype=float), + sequences=np.array([["s1"], ["s2"]], dtype=object), + input_tokens=np.array([[1.0], [1.0]], dtype=float), + output_tokens=np.array([[0.0], [0.0]], dtype=float), + agg_input_tokens=np.array([10.0, 8.0], dtype=float), + agg_output_tokens=np.array([0.0, 0.0], dtype=float), ) - objectives = optimizer._evaluate_candidates(candidates) - selected, _ = optimizer._select_population(candidates, objectives) + vecs = optimizer._get_objective_vectors(result) - assert len(selected) == 1 - assert selected[0].instruction == "longer prompt" + assert vecs.shape == (2, 2) + assert np.allclose(vecs[:, 0], np.array([0.4, 0.9])) + assert np.allclose(vecs[:, 1], -np.array([10.0, 8.0])) + + fronts = optimizer._non_dominated_sort(vecs) + + assert fronts[0] == [1] + assert 0 in fronts[1] def test_capoeira_meta_prompts(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): @@ -81,7 +86,7 @@ def test_capoeira_meta_prompts(mock_meta_llm, mock_predictor, initial_prompts, m mother = Prompt("Instruction 1", ["Example 1"]) father = Prompt("Instruction 2", ["Example 2"]) - perform_crossover([mother, father], optimizer.crossovers_per_iter, optimizer.crossover_template, optimizer.meta_llm) + perform_crossover([mother, father], optimizer=optimizer) full_task_desc = mock_task.task_description + "\n" + optimizer.predictor.extraction_description expected_crossover = ( @@ -95,22 +100,525 @@ def test_capoeira_meta_prompts(mock_meta_llm, mock_predictor, initial_prompts, m parent = Prompt("Instruction 3", ["Example 3"]) perform_mutation( offsprings=[parent], - mutation_template=optimizer.mutation_template, - upper_shots=optimizer.upper_shots, - meta_llm=optimizer.meta_llm, - few_shot_kwargs=dict( - df_few_shots=mock_df, - x_column=mock_task.x_column, - y_column=mock_task.y_column, - predictor=mock_predictor, - fewshot_template=CAPO_FEWSHOT_TEMPLATE, - target_begin_marker=optimizer.target_begin_marker, - target_end_marker=optimizer.target_end_marker, - check_fs_accuracy=True, - create_fs_reasoning=True, - ), + optimizer=optimizer, ) expected_mutation = CAPO_MUTATION_TEMPLATE.replace("", parent.instruction).replace( "", full_task_desc ) assert expected_mutation in mock_meta_llm.call_history[0]["prompts"] + + +def test_capoeira_crowding_distance_edges(): + vecs = np.array([[1.0, 2.0], [3.0, 4.0]]) + dists = Capoeira._calculate_crowding_distance(vecs) + assert np.isinf(dists).all() + + +def test_capoeira_select_survivors_handles_heterogeneous_blocks(mock_meta_llm, mock_predictor): + def fake_evaluate(*_, **__): + return EvalResult( + scores=np.array([[0.5]], dtype=float), + agg_scores=np.array([0.5], dtype=float), + sequences=np.array([[""]], dtype=object), + input_tokens=np.array([[0.0]], dtype=float), + output_tokens=np.array([[0.0]], dtype=float), + agg_input_tokens=np.array([0.0], dtype=float), + agg_output_tokens=np.array([0.0], dtype=float), + ) + + task = MockTask( + eval_strategy="sequential_block", + n_blocks=2, + block_idx=0, + eval_blocks={}, + evaluate_fn=fake_evaluate, + ) + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["inc1", "inc2"], + df_few_shots=task.pop_datapoints(n=1), + ) + + c1, c2 = Prompt("c1"), Prompt("c2") + task.eval_blocks = {str(c1): {0}, str(c2): {0, 1}} + optimizer.incumbents = [Prompt("i1"), Prompt("i2")] + optimizer.challengers = [c1, c2] + optimizer.population_size = 3 + + optimizer._select_survivors() + + assert len(optimizer.challengers) == 1 + assert optimizer.challengers[0].instruction == "c2" + + +def test_capoeira_select_survivors_homogeneous_prunes_lowest(mock_meta_llm, mock_predictor): + next_result: dict[str, EvalResult | None] = {"value": None} + + def fake_evaluate(prompts, *_, **__): + return next_result["value"] # type: ignore[return-value] + + task = MockTask( + eval_strategy="sequential_block", + n_blocks=2, + block_idx=0, + eval_blocks={}, + evaluate_fn=fake_evaluate, + ) + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["inc"], + df_few_shots=task.pop_datapoints(n=1), + ) + + c1, c2 = Prompt("c1"), Prompt("c2") + task.eval_blocks = {str(c1): {0}, str(c2): {0}} + + next_result["value"] = EvalResult( + scores=np.array([[0.1], [0.2]], dtype=float), + agg_scores=np.array([0.1, 0.2], dtype=float), + sequences=np.array([["s1"], ["s2"]], dtype=object), + input_tokens=np.array([[0.0], [0.0]], dtype=float), + output_tokens=np.array([[0.0], [0.0]], dtype=float), + agg_input_tokens=np.array([0.0, 0.0], dtype=float), + agg_output_tokens=np.array([0.0, 0.0], dtype=float), + ) + + optimizer.incumbents = [Prompt("inc")] # keeps population pressure + optimizer.challengers = [c1, c2] + optimizer.population_size = 2 + + optimizer._select_survivors() + + assert len(optimizer.challengers) == 1 + assert optimizer.challengers[0].instruction == "c2" + + +def test_capoeira_select_survivors_prefers_lower_cost(mock_meta_llm, mock_predictor): + def fake_evaluate(prompts, *_, **__): + costs = np.array([1.0 if "cheap" in p.instruction else 5.0 for p in prompts], dtype=float) + return EvalResult( + scores=np.array([[0.4], [0.4]], dtype=float), + agg_scores=np.array([0.4, 0.4], dtype=float), + sequences=np.array([["s1"], ["s2"]], dtype=object), + input_tokens=costs.reshape(-1, 1), + output_tokens=np.zeros((len(prompts), 1)), + agg_input_tokens=costs, + agg_output_tokens=np.zeros(len(prompts)), + ) + + task = MockTask( + eval_strategy="sequential_block", + n_blocks=1, + block_idx=0, + eval_blocks={"cheap": {0}, "expensive": {0}}, + evaluate_fn=fake_evaluate, + ) + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["cheap", "expensive"], + df_few_shots=task.pop_datapoints(n=1), + ) + + optimizer.incumbents = [] + optimizer.challengers = [Prompt("cheap"), Prompt("expensive")] + optimizer.population_size = 1 + + optimizer._select_survivors() + + assert len(optimizer.challengers) == 1 + assert optimizer.challengers[0].instruction == "cheap" + + +def test_capoeira_step_invokes_hooks(mock_meta_llm, mock_predictor, mock_df): + task = MockTask() + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["p1", "p2"], + df_few_shots=mock_df, + ) + + def fake_eval(prompts, *_, **__): + n = len(prompts) + return EvalResult( + scores=np.zeros((n, 1), dtype=float), + agg_scores=np.arange(n, dtype=float), + sequences=np.array([[""] for _ in range(n)], dtype=object), + input_tokens=np.ones((n, 1)), + output_tokens=np.zeros((n, 1)), + agg_input_tokens=np.ones(n), + agg_output_tokens=np.zeros(n), + ) + + optimizer.task.evaluate = fake_eval # type: ignore[assignment] + optimizer.incumbents = [Prompt("inc")] + optimizer.prompts = [Prompt("p1"), Prompt("p2")] + + with patch("promptolution.optimizers.capoeira.perform_crossover", return_value=[Prompt("c1")]), patch( + "promptolution.optimizers.capoeira.perform_mutation", return_value=[Prompt("m1")] + ), patch.object(optimizer, "_do_intensification") as do_int, patch.object( + optimizer, "_advance_one_incumbent" + ) as adv_inc, patch.object(optimizer, "_select_survivors") as sel: + optimizer._step() + + assert do_int.call_count == 1 + assert adv_inc.call_count == 1 + assert sel.call_count == 1 + + +def test_capoeira_do_intensification_updates_incumbents(mock_meta_llm, mock_predictor): + def fake_eval(prompts, *_, **__): + n = len(prompts) + scores = np.arange(1, n + 1, dtype=float).reshape(n, 1) + return EvalResult( + scores=scores, + agg_scores=scores.flatten(), + sequences=np.array([[""] for _ in range(n)], dtype=object), + input_tokens=np.ones((n, 1)), + output_tokens=np.zeros((n, 1)), + agg_input_tokens=np.ones(n), + agg_output_tokens=np.zeros(n), + ) + + task = MockTask(eval_strategy="sequential_block", n_blocks=2, block_idx=0, evaluate_fn=fake_eval) + challenger = Prompt("chal") + inc1, inc2 = Prompt("i1"), Prompt("i2") + task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {0}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["p"], + df_few_shots=task.pop_datapoints(n=1), + ) + optimizer.incumbents = [inc1, inc2] + update_mock = patch.object(optimizer, "_update_incumbent_front", autospec=True).start() + + with patch("random.choice", side_effect=lambda seq: seq[0]): + optimizer._do_intensification(challenger) + + patch.stopall() + assert challenger in optimizer.incumbents + update_mock.assert_called_once() + + +def test_capoeira_do_intensification_bootstrap_no_common_blocks(mock_meta_llm, mock_predictor): + def fake_eval(prompts, *_, **__): + n = len(prompts) + return EvalResult( + scores=np.zeros((n, 1)), + agg_scores=np.zeros(n), + sequences=np.array([[ + "" + ] for _ in range(n)], dtype=object), + input_tokens=np.zeros((n, 1)), + output_tokens=np.zeros((n, 1)), + agg_input_tokens=np.zeros(n), + agg_output_tokens=np.zeros(n), + ) + + task = MockTask(eval_strategy="sequential_block", n_blocks=3, block_idx=0, evaluate_fn=fake_eval) + inc1, inc2, challenger = Prompt("i1"), Prompt("i2"), Prompt("chal") + task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {1}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["p"], + df_few_shots=task.pop_datapoints(n=1), + ) + optimizer.incumbents = [inc1, inc2] + with patch("random.randrange", return_value=2), patch.object( + optimizer, "_update_incumbent_front", autospec=True + ) as upd: + optimizer._do_intensification(challenger) + + assert task.block_idx == 2 + assert challenger in optimizer.incumbents + upd.assert_called_once_with(blocks={2}) + + +def test_capoeira_do_intensification_running_mean_path(monkeypatch, mock_meta_llm, mock_predictor): + task = MockTask(eval_strategy="sequential_block", n_blocks=2, block_idx=0) + inc1, inc2, challenger = Prompt("i1"), Prompt("i2"), Prompt("chal") + task.prompt_evaluated_blocks = {str(inc1): {0, 1}, str(inc2): {0, 1}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["p"], + df_few_shots=task.pop_datapoints(n=1), + ) + optimizer.incumbents = [inc1, inc2] + + vec1 = np.array([[0.1, -0.1], [0.2, -0.2], [0.15, -0.15]]) + vec2 = np.array([[0.2, -0.2], [0.3, -0.3], [0.25, -0.25]]) + + calls: list[tuple] = [] + + def fake_is_dom(_self, v1, v2): + calls.append((v1.copy(), v2.copy())) + return False + + monkeypatch.setattr(Capoeira, "_is_dominated", fake_is_dom) + + with patch.object(Capoeira, "_get_objective_vectors", side_effect=[vec1, vec2]), patch( + "random.choice", side_effect=lambda seq: list(seq)[0] + ), patch.object(optimizer, "_update_incumbent_front", autospec=True) as upd: + optimizer._do_intensification(challenger) + + # fold_vec path should call dominance check at least once + assert calls, "_is_dominated should be invoked when challenger_mean already set" + assert challenger in optimizer.incumbents + upd.assert_called_once() + + +def test_capoeira_do_intensification_dominated_challenger(monkeypatch, mock_meta_llm, mock_predictor): + task = MockTask(eval_strategy="sequential_block", n_blocks=1, block_idx=0) + inc1, inc2, challenger = Prompt("i1"), Prompt("i2"), Prompt("chal") + task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {0}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["p"], + df_few_shots=task.pop_datapoints(n=1), + ) + optimizer.incumbents = [inc1, inc2] + + dominated_vecs = np.array([[0.9, -0.1], [0.8, -0.1], [0.1, -0.1]]) + + with patch.object(Capoeira, "_get_objective_vectors", return_value=dominated_vecs), patch( + "random.choice", side_effect=lambda seq: list(seq)[0] + ): + optimizer._do_intensification(challenger) + + assert challenger in optimizer.challengers + assert challenger not in optimizer.incumbents + + +def test_capoeira_update_incumbent_front_demotes(mock_meta_llm, mock_predictor): + def fake_eval(prompts, *_, **__): + scores = np.array([0.3, 0.1], dtype=float) + return EvalResult( + scores=scores.reshape(-1, 1), + agg_scores=scores, + sequences=np.array([["s1"], ["s2"]], dtype=object), + input_tokens=np.zeros((2, 1)), + output_tokens=np.zeros((2, 1)), + agg_input_tokens=np.zeros(2), + agg_output_tokens=np.zeros(2), + ) + + task = MockTask(eval_strategy="sequential_block", n_blocks=1, evaluate_fn=fake_eval) + inc1, inc2 = Prompt("best"), Prompt("worst") + task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {0}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["p"], + df_few_shots=task.pop_datapoints(n=1), + ) + optimizer.incumbents = [inc1, inc2] + + optimizer._update_incumbent_front() + + assert optimizer.incumbents == [inc1] + assert inc2 in optimizer.challengers + + +def test_capoeira_advance_one_incumbent_no_gapblocks(mock_meta_llm, mock_predictor): + task = MockTask(eval_strategy="sequential_block", n_blocks=2, block_idx=0) + inc = Prompt("p1") + task.prompt_evaluated_blocks = {str(inc): {0, 1}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["p1"], + df_few_shots=task.pop_datapoints(n=1), + ) + optimizer.incumbents = [inc] + + called = {"evaluate": 0} + + def no_call(*args, **kwargs): + called["evaluate"] += 1 + raise AssertionError("evaluate should not be called when no new blocks") + + task.evaluate = no_call # type: ignore[assignment] + + optimizer._advance_one_incumbent() + + assert called["evaluate"] == 0 + + +def test_capoeira_get_closest_incumbent_returns_nearest(): + challenger = np.array([0.5, 0.5]) + incumbents = np.array([[0.0, 0.0], [0.6, 0.6]]) + res = Capoeira._get_closest_incumbent(None, challenger, incumbents) + assert np.allclose(res, incumbents[1]) + + +def test_capoeira_objective_vectors_multiobjective(mock_meta_llm, mock_predictor, mock_df): + t1 = MockTask(df=mock_df, n_subsamples=1, n_blocks=1) + t2 = MockTask(df=mock_df, n_subsamples=1, n_blocks=1) + multi_task = MultiObjectiveTask(tasks=[t1, t2]) + + optimizer = Capoeira( + predictor=mock_predictor, + task=multi_task, + meta_llm=mock_meta_llm, + initial_prompts=["p"], + df_few_shots=mock_df, + ) + + result = MultiObjectiveEvalResult( + scores=[np.array([[0.1], [0.2]]), np.array([[0.3], [0.4]])], + agg_scores=[np.array([0.1, 0.2]), np.array([0.3, 0.4])], + sequences=np.array([["s1"], ["s2"]], dtype=object), + input_tokens=np.array([[1.0], [2.0]]), + output_tokens=np.array([[0.0], [0.0]]), + agg_input_tokens=np.array([1.0, 2.0]), + agg_output_tokens=np.array([0.0, 0.0]), + ) + + vecs = optimizer._get_objective_vectors(result) + assert vecs.shape == (2, 3) + assert np.allclose(vecs[:, 0], [0.1, 0.2]) + assert np.allclose(vecs[:, 1], [0.3, 0.4]) + assert np.allclose(vecs[:, 2], -np.array([1.0, 2.0])) + + +def test_capoeira_advance_one_incumbent_chooses_gap(mock_meta_llm, mock_predictor): + def fake_eval(*_, **__): + return EvalResult( + scores=np.array([[0.0]]), + agg_scores=np.array([0.0]), + sequences=np.array([[""]], dtype=object), + input_tokens=np.array([[0.0]]), + output_tokens=np.array([[0.0]]), + agg_input_tokens=np.array([0.0]), + agg_output_tokens=np.array([0.0]), + ) + + task = MockTask(eval_strategy="sequential_block", n_blocks=3, block_idx=0, evaluate_fn=fake_eval) + p1, p2 = Prompt("p1"), Prompt("p2") + task.prompt_evaluated_blocks = {str(p1): {0}, str(p2): {0, 2}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["p1", "p2"], + df_few_shots=task.pop_datapoints(n=1), + ) + optimizer.incumbents = [p1, p2] + + with patch("random.choice", side_effect=lambda seq: list(seq)[0]): + optimizer._advance_one_incumbent() + + assert task.block_idx == 2 + + +def test_capoeira_select_survivors_heterogeneous_removes_lowest(mock_meta_llm, mock_predictor): + task = MockTask(eval_strategy="sequential_block", n_blocks=3) + c1, c2 = Prompt("c1"), Prompt("c2") + task.prompt_evaluated_blocks = {str(c1): {0}, str(c2): {1}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["inc"], + df_few_shots=task.pop_datapoints(n=1), + ) + optimizer.incumbents = [] + optimizer.challengers = [c1, c2] + optimizer.population_size = 1 + + with patch("random.choice", side_effect=lambda seq: list(seq)[0]): + optimizer._select_survivors() + + assert len(optimizer.challengers) == 1 + + +def test_capoeira_select_survivors_incumbent_only(mock_meta_llm, mock_predictor): + def fake_eval(prompts, *_, **__): + n = len(prompts) + vals = np.linspace(0.1, 0.2, n) + return EvalResult( + scores=np.tile(vals.reshape(n, 1), (1, 1)), + agg_scores=vals, + sequences=np.array([[""] for _ in range(n)], dtype=object), + input_tokens=np.ones((n, 1)), + output_tokens=np.zeros((n, 1)), + agg_input_tokens=np.ones(n), + agg_output_tokens=np.zeros(n), + ) + + task = MockTask(eval_strategy="sequential_block", n_blocks=2, evaluate_fn=fake_eval) + inc1, inc2 = Prompt("i1"), Prompt("i2") + task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {0}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["i1", "i2"], + df_few_shots=task.pop_datapoints(n=1), + ) + optimizer.incumbents = [inc1, inc2] + optimizer.challengers = [] + optimizer.population_size = 1 + + optimizer._select_survivors() + + assert len(optimizer.incumbents) == 1 + + +def test_capoeira_get_common_blocks(mock_meta_llm, mock_predictor): + task = MockTask(eval_strategy="sequential_block", n_blocks=2) + p1, p2 = Prompt("p1"), Prompt("p2") + task.prompt_evaluated_blocks = {str(p1): {0, 1}, str(p2): {1}} + + optimizer = Capoeira( + predictor=mock_predictor, + task=task, + meta_llm=mock_meta_llm, + initial_prompts=["p1", "p2"], + df_few_shots=task.pop_datapoints(n=1), + ) + + common = optimizer._get_common_blocks([p1, p2]) + assert common == {1} + + +def test_capoeira_is_dominated_logic(): + assert Capoeira._is_dominated(np.array([0.1, 0.1]), np.array([0.2, 0.2])) + assert not Capoeira._is_dominated(np.array([0.3, 0.2]), np.array([0.3, 0.2])) + assert not Capoeira._is_dominated(np.array([0.4, 0.5]), np.array([0.3, 0.6])) + + +def test_capoeira_calculate_crowding_distance_three_points(): + vecs = np.array([[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]]) + dists = Capoeira._calculate_crowding_distance(vecs) + assert np.isinf(dists[[0, -1]]).all() + assert dists[1] > 0 diff --git a/tests/predictors/test_base_predictor.py b/tests/predictors/test_base_predictor.py index d20f51d4..1d657184 100644 --- a/tests/predictors/test_base_predictor.py +++ b/tests/predictors/test_base_predictor.py @@ -1,5 +1,9 @@ import numpy as np +from tests.mocks.dummy_config import DummyConfig +from tests.mocks.mock_llm import MockLLM +from tests.mocks.mock_predictor import MockPredictor + def test_predictor_predict_flow(mock_predictor): """Test the basic prediction flow from prompt to final prediction.""" @@ -37,3 +41,23 @@ def test_predictor_with_return_seq(mock_predictor): assert len(sequences) == 1 assert isinstance(sequences, list) assert "This product is okay." in sequences[0] + + +def test_predictor_accepts_string_prompt(mock_predictor): + preds, seqs = mock_predictor.predict("solo", ["input"], system_prompts="sys") + assert preds.shape[0] == 1 + assert seqs[0].startswith("input\n") + + +def test_predictor_system_prompt_string_converted(mock_predictor): + preds, seqs = mock_predictor.predict(["p1", "p2"], ["x1", "x2"], system_prompts="sys") + assert len(preds) == 2 + # call_history should show system_prompts broadcasted + assert mock_predictor.llm.call_history[-1]["system_prompts"] == ["sys", "sys"] + + +def test_predictor_applies_config(): + cfg = DummyConfig() + predictor = MockPredictor(llm=MockLLM(), config=cfg) + assert cfg.applied is True + assert getattr(predictor, "config_applied") is True diff --git a/tests/tasks/test_base_task.py b/tests/tasks/test_base_task.py new file mode 100644 index 00000000..f3387e77 --- /dev/null +++ b/tests/tasks/test_base_task.py @@ -0,0 +1,146 @@ +import numpy as np +import pandas as pd +import pytest + +from tests.mocks.dummy_config import DummyConfig +from tests.mocks.mock_llm import MockLLM +from tests.mocks.mock_predictor import MockPredictor +from tests.mocks.mock_task import MockTask + +from promptolution.tasks.base_task import BaseTask +from promptolution.utils.prompt import Prompt + + +@pytest.fixture +def predictor(): + return MockPredictor(llm=MockLLM()) + + +@pytest.fixture +def small_task(): + df = pd.DataFrame({"x": ["a", "b", "c"], "y": ["1", "0", "1"]}) + return MockTask(df=df, eval_strategy="sequential_block", n_subsamples=1) + + +@pytest.fixture +def cost_task(): + df = pd.DataFrame({"x": ["m", "n", "o"], "y": ["1", "0", "1"]}) + return MockTask(df=df, eval_strategy="full", n_subsamples=3) + + +def test_subsample_and_block_controls(small_task): + task = small_task + + xs, ys = task.subsample() + assert len(xs) == 1 + + task.increment_block_idx() + assert task.block_idx == 1 % task.n_blocks if task.n_blocks else 0 + + task.set_block_idx([0, 1, 2]) + xs2, _ = task.subsample() + assert set(xs2) == set(task.xs) + + task.set_block_idx(0) + popped = task.pop_datapoints(n=1) + assert len(popped) == 1 + assert len(task.xs) == 2 + + task.reset_block_idx() + assert task.block_idx == 0 + + task.eval_strategy = "full" + with pytest.raises(ValueError): + task.increment_block_idx() + with pytest.raises(ValueError): + task.reset_block_idx() + + +def test_prepare_batch_and_evaluated_strategy(small_task): + task = small_task + prompts = [Prompt("p1"), Prompt("p2")] + xs, ys = task.subsample() + + to_eval = task._prepare_batch(prompts, xs, ys, eval_strategy="evaluated") + assert to_eval == ([], [], [], []) + + normal = task._prepare_batch(prompts, xs, ys) + assert len(normal[0]) == len(prompts) * len(xs) + + +def test_pop_datapoints_clears_cache_and_frac(small_task): + task = small_task + p = Prompt("p") + key = (str(p), task.xs[0], task.ys[0]) + task.eval_cache[key] = 0.5 + task.seq_cache[key] = "seq" + + popped = task.pop_datapoints(frac=0.5) + assert len(popped) > 0 + assert not task.eval_cache + assert not task.seq_cache + + +def test_unknown_strategy_raises(small_task): + task = small_task + task.eval_strategy = "unknown" + with pytest.raises(ValueError): + task.subsample() + + +def test_set_block_idx_validation(small_task): + task = small_task + with pytest.raises(AssertionError): + task.set_block_idx("bad") # type: ignore + + +def test_pop_datapoints_requires_arg(small_task): + task = small_task + with pytest.raises(AssertionError): + task.pop_datapoints(n=1, frac=0.1) + + +def test_get_evaluated_blocks_mapping(small_task): + task = small_task + prompt = Prompt("p") + task.prompt_evaluated_blocks[str(prompt)] = {0, 1} + mapping = task.get_evaluated_blocks([prompt]) + assert mapping[str(prompt)] == {0, 1} + + +def test_compute_costs_shapes(predictor, cost_task): + task = cost_task + prompts = [Prompt("inst"), Prompt("inst2")] + result = task.evaluate(prompts, predictor) + + assert result.input_tokens.shape[0] == len(prompts) + assert result.output_tokens.shape[0] == len(prompts) + + +def test_evaluate_with_block_list_updates_blocks(predictor, small_task): + task = small_task + task.block_idx = [0, 1] + prompts = [Prompt("p1"), Prompt("p2")] + task.evaluate(prompts, predictor) + for p in prompts: + assert task.prompt_evaluated_blocks[str(p)] == {0, 1} + + +def test_task_config_applied(): + cfg = DummyConfig() + df = pd.DataFrame({"x": ["a", "b", "c"], "y": ["1", "0", "1"]}) + task = MockTask(df=df, eval_strategy="sequential_block", n_subsamples=1, config=cfg) + assert cfg.applied is True + assert hasattr(task, "config_applied") + + +def test_block_wraparound_and_get_cache_keys(): + df = pd.DataFrame({"x": ["a", "b"], "y": ["1", "0"]}) + task = MockTask(df=df, eval_strategy="sequential_block", n_subsamples=1) + task.block_idx = task.n_blocks - 1 + task.increment_block_idx() + assert task.block_idx == 0 + + prompt = Prompt("hi") + key = task._cache_key(prompt, "x", "y") + assert key[0].startswith(prompt.instruction) \ No newline at end of file diff --git a/tests/tasks/test_multi_objective_task.py b/tests/tasks/test_multi_objective_task.py new file mode 100644 index 00000000..228b7c3d --- /dev/null +++ b/tests/tasks/test_multi_objective_task.py @@ -0,0 +1,75 @@ +import pandas as pd +import numpy as np + +import pytest + +from tests.mocks.mock_predictor import MockPredictor +from tests.mocks.mock_task import MockTask +from tests.mocks.mock_llm import MockLLM + +from promptolution.tasks.multi_objective_task import MultiObjectiveTask +from promptolution.utils.prompt import Prompt + + +def test_multi_objective_single_prediction_flow(): + task1 = MockTask() + task2 = MockTask() + predictor = MockPredictor(llm=MockLLM()) + + prompt = Prompt("classify") + result = MultiObjectiveTask([task1, task2]).evaluate([prompt], predictor=predictor) + + assert len(result.agg_scores) == 2 + assert result.agg_scores[0].shape == (1,) + assert result.sequences.shape[0] == 1 + assert MultiObjectiveTask([task1, task2]).tasks[0].n_subsamples == task1.n_subsamples + + +def test_multi_objective_shares_block_and_caches(): + df = pd.DataFrame({"x": ["u", "v"], "y": ["1", "0"]}) + t1 = MockTask(df=df, eval_strategy="sequential_block", n_subsamples=1, n_blocks=len(df), block_idx=0) + t2 = MockTask(df=df, eval_strategy="sequential_block", n_subsamples=1, n_blocks=len(df), block_idx=0) + + predictor = MockPredictor(llm=MockLLM()) + prompt = Prompt("judge") + + multi = MultiObjectiveTask([t1, t2]) + multi.block_idx = 1 + res = multi.evaluate(prompt, predictor=predictor) + + assert len(t1.eval_cache) == len(t2.eval_cache) + assert res.input_tokens.shape[0] == 1 + assert multi.prompt_evaluated_blocks[str(prompt)] == {1} + + +def test_multi_objective_requires_tasks(): + with pytest.raises(ValueError): + MultiObjectiveTask([]) + + +def test_multi_objective_matches_individual_results(): + df = pd.DataFrame({"x": ["u", "v"], "y": ["1", "0"]}) + + def make_task(): + return MockTask(df=df, eval_strategy="sequential_block", n_subsamples=1, n_blocks=len(df), block_idx=0) + + t1 = make_task() + t2 = make_task() + predictor = MockPredictor(llm=MockLLM()) + + prompt = Prompt("judge") + multi = MultiObjectiveTask([t1, t2]) + multi.block_idx = 1 + multi_res = multi.evaluate([prompt], predictor=predictor) + + # Fresh tasks/predictor to mirror a single-task call + s1 = make_task() + s2 = make_task() + single_pred = MockPredictor(llm=MockLLM()) + res1 = s1.evaluate([prompt], predictor=single_pred) + res2 = s2.evaluate([prompt], predictor=single_pred) + + assert np.allclose(multi_res.agg_scores[0], res1.agg_scores) + assert np.allclose(multi_res.agg_scores[1], res2.agg_scores) + assert multi_res.sequences.shape == res1.sequences.shape + assert multi.prompt_evaluated_blocks[str(prompt)] == {1} diff --git a/tests/utils/test_prompt_creation.py b/tests/utils/test_prompt_creation.py index 1c8c9506..faefd4ae 100644 --- a/tests/utils/test_prompt_creation.py +++ b/tests/utils/test_prompt_creation.py @@ -1,3 +1,5 @@ +import numpy as np + from promptolution.tasks.base_task import BaseTask from promptolution.tasks.classification_tasks import ClassificationTask from promptolution.utils.prompt_creation import create_prompt_variation, create_prompts_from_samples @@ -143,3 +145,20 @@ def test_create_prompts_from_samples_multiple_prompts(mock_df, mock_meta_llm): assert len(generated_prompts) == n_prompts assert len(mock_meta_llm.call_history) == 1 + + +def test_create_prompts_from_samples_uniform_labels(mock_df, mock_meta_llm): + """Ensure uniform-label sampling includes every class in the meta-prompt examples.""" + task = ClassificationTask(df=mock_df, x_column="x", y_column="y") + task.xs = np.asarray(task.xs) + task.ys = np.asarray(task.ys) + + mock_meta_llm.reset() + + prompts = create_prompts_from_samples(task, mock_meta_llm, n_samples=2, n_prompts=1, get_uniform_labels=True) + + assert len(prompts) == 1 + # The constructed meta-prompt should include at least one example per label + sent_prompt = mock_meta_llm.call_history[0]["prompts"][0] + for label in ["positive", "negative", "neutral"]: + assert f"Output: {label}" in sent_prompt From 505702da2225d747e16cc569bad97e7264a83257 Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 6 Jan 2026 18:19:58 +0100 Subject: [PATCH 18/53] parse kwargs to reward functions" --- .coverage | Bin 69632 -> 69632 bytes promptolution/helpers.py | 4 +- promptolution/optimizers/capo.py | 2 +- promptolution/optimizers/capoeira.py | 80 ++++++-------------- promptolution/tasks/__init__.py | 2 +- promptolution/tasks/base_task.py | 15 ++-- promptolution/tasks/multi_objective_task.py | 19 +++-- promptolution/tasks/reward_tasks.py | 15 +++- promptolution/utils/capo_utils.py | 6 +- tests/helpers/test_helpers.py | 12 +-- tests/llms/test_api_llm.py | 2 + tests/mocks/dummy_config.py | 8 +- tests/mocks/mock_task.py | 13 +++- tests/optimizers/test_capo.py | 2 +- tests/optimizers/test_capoeira.py | 11 +-- tests/tasks/test_base_task.py | 4 +- tests/tasks/test_multi_objective_task.py | 5 +- tests/tasks/test_reward_tasks.py | 22 ++++++ 18 files changed, 124 insertions(+), 98 deletions(-) diff --git a/.coverage b/.coverage index e65bb0277db3ed62b8694bc1bfcdf4f8524131a0..442105beb59832cb9b8953cee070fcd08af6c906 100644 GIT binary patch delta 791 zcmZozz|ydQWrMB*W8G#w2XPLL%1XoH!W6@T$-W-f*dj$(7#htd2Rb@T)_pI^lqfv; zqnFp@d{6Ps4?RQJnCm%JHd_SnB5R7c&Jx7})Uq(%U7#qnJh3Pxz9g|YyI8NFQe^Un zC~t`7NL4k)hm+YG6jJ;7n)x#MeEDqo^!TLsczNIQ-s8Q>dz5z@?^51=-V9y`UIw1; zJnwiO@m%Fu&NGjvnkR|Jhewe6H}_-iQ`~F07jjSMZs*SBj^;Muy32K%YdhBhu8Cab zT(MksTv}X`T-=<0IUjBo6bR*<{G?M$p2vZOkyFZHR^6`<|Gjzq_CTJ3{kvyxD|Z)f ztKP>xS-ZDu^xJ!D?>+loGw=EHv-9WQXPjKomFFL2Uw>}?Z2SKuXYId#)jj|J z&!3vJ-#%}yU-#4A-v0WZ&-UNn|FM6+eeVC_{qOh0zpgs>`|2zpd)D{WD{AhYx1axi zXZ-2kcR&5V^ZwNDna`Ogn|3F2aTo$!AY`~{o%ZB~-LgEM~s31pRQ)%K(``7jid+yub<=XK7|L#|}f3M&D_x;!H z|F*4dwEOYCvOoJD=Yz*-|G4T{!owIhAL}c0WCW(83n=O622ttg=3YLe1RJ5W`NWGV WHb!_toBS{`deR47Qj_*ic`*R7X!L^WU5?2LR9G5+pHkT9^ z9~T4Xlg)wx5uB4Bc52CTIj}HtiaA{JpZ|)H-yX`k(pp z=V$->Isf4XW5@;nSY;FT6SA(@~Lk35L<1abd7cv>$lXu z|Ct#c?7bJevHq>{+0PkI*1xW~cgOhq+cN>PuVp{mdB>QMAv)Ep!qcbMJbPy8>;F-^ z-`5Ch0kvyrtva_gzUXP{-+jk-?ahyxoYa%gBE)b&eDc+vU82HrASH5HsR92P6aL+s z8O~VoE^zYF-Wl3bB0w24k*u@z*S?>vUu}H)`+oap^~e7upR;G=6$Z-52yfzQ{>OZv zUTSh?pPHG808lJTU{h)0PW#vP41exkUd!`+r>h>-^hNFPi_)`@ZAZ zd`AKGs+0$Oue07f=$;4Wxv_^?^|kk#KJA "BasePredictor": +def get_predictor( + downstream_llm: Optional["BaseLLM"] = None, type: "PredictorType" = "marker", *args, **kwargs +) -> "BasePredictor": """Create and return a predictor instance. This function supports three types of predictors: diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index 161552b8..961cd9f6 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -19,7 +19,7 @@ from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation from promptolution.utils.logging import get_logger from promptolution.utils.prompt import Prompt, sort_prompts_by_scores -from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE +from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_MUTATION_TEMPLATE from promptolution.utils.test_statistics import get_test_statistic_func from promptolution.utils.token_counter import get_token_counter diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 7fd05144..b9ec0330 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -16,7 +16,6 @@ from promptolution.optimizers.base_optimizer import BaseOptimizer from promptolution.tasks.multi_objective_task import MultiObjectiveTask - from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation from promptolution.utils.logging import get_logger from promptolution.utils.prompt import Prompt @@ -86,7 +85,7 @@ def __init__( self.incumbents: List[Prompt] = self.prompts self.challengers: List[Prompt] = [] self.population_size = len(self.prompts) - + if "block" not in self.task.eval_strategy: logger.warning( f"ℹ️ CAPO requires 'block' in the eval_strategy, but got {self.task.eval_strategy}. Setting eval_strategy to 'sequential_block'." @@ -111,20 +110,16 @@ def _pre_optimization_loop(self) -> None: ) population.append(Prompt(prompt.instruction, few_shots)) - self.max_prompt_length = ( - max(self.token_counter(p.construct_prompt()) for p in population) if population else 1 - ) + self.max_prompt_length = max(self.token_counter(p.construct_prompt()) for p in population) if population else 1 init_result = self.task.evaluate(population, self.predictor) initial_vectors = self._get_objective_vectors(init_result) fronts = self._non_dominated_sort(initial_vectors) self.incumbents = [population[i] for i in fronts[0]] self.challengers = [population[i] for front in fronts[1:] for i in front] - # keep self.prompts as a "view" if base class expects it self.scores = initial_vectors[:, 0].tolist() - def _step(self) -> List[Prompt]: # 1) generate challengers (random parent selection happens inside perform_crossover) offsprings = perform_crossover(self.prompts, self) @@ -138,7 +133,9 @@ def _step(self) -> List[Prompt]: # 4) logging scores: incumbents only (optional) if self.incumbents: - inc_result = self.task.evaluate(prompts=self.incumbents, predictor=self.predictor, eval_strategy="evaluated") + inc_result = self.task.evaluate( + prompts=self.incumbents, predictor=self.predictor, eval_strategy="evaluated" + ) vecs_inc = self._get_objective_vectors(inc_result) self.scores = vecs_inc[:, 0].tolist() else: @@ -204,10 +201,7 @@ def _do_intensification(self, challenger: Prompt) -> None: self.incumbents.append(challenger) self._update_incumbent_front(blocks=common_blocks) - - def _get_closest_incumbent( - self, challenger_vec: np.ndarray, incumbent_vecs: np.ndarray - ) -> np.ndarray: + def _get_closest_incumbent(self, challenger_vec: np.ndarray, incumbent_vecs: np.ndarray) -> np.ndarray: """Return the vector of the geometrically closest incumbent.""" all_vecs = np.vstack([incumbent_vecs, challenger_vec[None, :]]) min_b = np.min(all_vecs, axis=0) @@ -222,7 +216,6 @@ def _get_closest_incumbent( idx = int(np.argmin(dists)) return incumbent_vecs[idx] - def _update_incumbent_front(self, blocks: Optional[set[int]] = None) -> None: if not self.incumbents: return @@ -230,9 +223,9 @@ def _update_incumbent_front(self, blocks: Optional[set[int]] = None) -> None: if blocks is None: res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") else: - self.task.set_block_idx(list(sorted(blocks))) # sorted for deterministic behaviour + self.task.set_block_idx(list(sorted(blocks))) # sorted for deterministic behaviour res = self.task.evaluate(self.incumbents, self.predictor) - + vecs = self._get_objective_vectors(res) fronts = self._non_dominated_sort(vecs) @@ -243,7 +236,6 @@ def _update_incumbent_front(self, blocks: Optional[set[int]] = None) -> None: self.incumbents = new_incumbents self.challengers.extend(demoted) - def _get_objective_vectors(self, result) -> np.ndarray: # If the task is multi-objective, include all objective dimensions, else single objective. if isinstance(self.task, MultiObjectiveTask): @@ -257,7 +249,7 @@ def _get_objective_vectors(self, result) -> np.ndarray: cost_scalar = cost_scalar.reshape(-1, 1) return np.hstack([agg_scores, -cost_scalar]) - + def _advance_one_incumbent(self) -> None: if not self.incumbents: return @@ -292,29 +284,14 @@ def _advance_one_incumbent(self) -> None: self.task.set_block_idx(b) self.task.evaluate(prompts=[chosen_inc], predictor=self.predictor) - def _select_survivors(self) -> None: - """ - Enforce |incumbents| + |challengers| <= population_size using Pareto logic. - - Logic: - 1. Prune from Challengers first (they are less optimal than incumbents). - - If challengers have DIFFERENT evaluation blocks (Heterogeneous): - We cannot fairly compare their scores. Prune the one with the FEWEST evaluations - (least information/newest). - - If challengers have the SAME evaluation blocks (Homogeneous): - Perform Non-Dominated Sorting (NDS). Identify the worst front. - Use Crowding Distance to prune the most crowded (least unique) individual from that front. - - 2. If no Challengers, prune from Incumbents. - - Use Crowding Distance to remove the least unique incumbent. - """ + """Prune population via Pareto logic to enforce size constraints.""" while len(self.incumbents) + len(self.challengers) > self.population_size: if len(self.challengers) > 0: # 1. Check Heterogeneity (Fairness Check) chal_blocks_map = self.task.get_evaluated_blocks(self.challengers) block_sets = list(chal_blocks_map.values()) - + first_set = block_sets[0] # Are all challengers evaluated on the exact same set of blocks? is_homogeneous = all(s == first_set for s in block_sets) @@ -324,54 +301,46 @@ def _select_survivors(self) -> None: # Prune the prompt with the FEWEST evaluations (least reliable/least invested). counts = [len(s) for s in block_sets] min_count = min(counts) - + # Find all indices with the minimum count (handle ties randomly) candidates = [i for i, c in enumerate(counts) if c == min_count] victim_idx = random.choice(candidates) - + self.challengers.pop(victim_idx) continue - + # CASE B: Homogeneous (Fair comparison). # Use NDS + Crowding Distance. - + # Get objective vectors for all challengers (safe because blocks are identical) - res = self.task.evaluate( - self.challengers, - self.predictor, - eval_strategy="evaluated" - ) + res = self.task.evaluate(self.challengers, self.predictor, eval_strategy="evaluated") vecs = self._get_objective_vectors(res) - + # Perform Non-Dominated Sort fronts = self._non_dominated_sort(vecs) - + # Select the worst front (the last one) worst_front_indices = fronts[-1] - + # Multiple candidates in worst front -> Prune by Crowding Distance # We want to keep diversity (high CD), so we remove low CD. worst_front_vecs = vecs[worst_front_indices] dists = self._calculate_crowding_distance(worst_front_vecs) - + # Find index relative to the worst front list local_worst_idx = int(np.argmin(dists)) # Map back to the main challenger list index victim_idx = worst_front_indices[local_worst_idx] - + self.challengers.pop(victim_idx) continue # --- PRUNE FROM INCUMBENTS --- # Fallback: If we only have incumbents, remove the least unique one. - res = self.task.evaluate( - self.incumbents, - self.predictor, - eval_strategy="evaluated" - ) + res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") vecs = self._get_objective_vectors(res) dists = self._calculate_crowding_distance(vecs) - + # Remove the one with the smallest crowding distance victim_idx = int(np.argmin(dists)) self.incumbents.pop(victim_idx) @@ -387,7 +356,6 @@ def _get_common_blocks(self, prompts: List[Prompt]) -> set: common = set.intersection(*block_sets) return common - @staticmethod def _non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: """Perform fast non-dominated sorting (NSGA-II) in a vectorized manner.""" @@ -420,7 +388,7 @@ def _non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: def _is_dominated(vec1, vec2): """Returns True if vec2 dominates vec1 in a maximize-all setting.""" return np.all(vec2 >= vec1) and np.any(vec2 > vec1) - + @staticmethod def _calculate_crowding_distance(obj_vectors: np.ndarray) -> np.ndarray: """Calculate crowding distance for a set of solutions.""" diff --git a/promptolution/tasks/__init__.py b/promptolution/tasks/__init__.py index 5f61ff1f..7dadf4f0 100644 --- a/promptolution/tasks/__init__.py +++ b/promptolution/tasks/__init__.py @@ -2,8 +2,8 @@ from promptolution.tasks.classification_tasks import ClassificationTask from promptolution.tasks.judge_tasks import JudgeTask -from promptolution.tasks.reward_tasks import RewardTask from promptolution.tasks.multi_objective_task import MultiObjectiveTask +from promptolution.tasks.reward_tasks import RewardTask __all__ = [ "ClassificationTask", diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 19a3042f..65c9b9ed 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -63,7 +63,11 @@ def __init__( seed (int): Random seed for reproducibility. config (ExperimentConfig, optional): Configuration for the task, overriding defaults. """ - self.df: pd.DataFrame = df + 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_description: Optional[str] = task_description @@ -75,10 +79,10 @@ def __init__( if config is not None: config.apply_to(self) - self.xs: List[str] = df[self.x_column].values.astype(str).tolist() + 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] = df[y_column].values.astype(str).tolist() + 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) @@ -356,7 +360,7 @@ def reset_block_idx(self) -> None: 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: Union[int, List[int]]) -> None: """Set the block index (or indices) for block subsampling strategies.""" if "block" not in self.eval_strategy: @@ -366,8 +370,9 @@ def set_block_idx(self, idx: Union[int, List[int]]) -> None: assert all(0 <= i < self.n_blocks for i in idx), "Block indices must be integers within valid range" else: assert isinstance(idx, int), "Block index must be an integer" - + self.block_idx = idx def get_evaluated_blocks(self, prompts: List[Prompt]) -> Dict[str, set[int]]: + """Return mapping of prompt string to evaluated block indices.""" return {str(p): set(self.prompt_evaluated_blocks.get(str(p), set())) for p in prompts} diff --git a/promptolution/tasks/multi_objective_task.py b/promptolution/tasks/multi_objective_task.py index 6d5e8f0e..32ffcb16 100644 --- a/promptolution/tasks/multi_objective_task.py +++ b/promptolution/tasks/multi_objective_task.py @@ -3,19 +3,22 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Dict, List, Optional, Tuple import numpy as np +from typing import Dict, List, Optional, Tuple + from promptolution.tasks.base_task import BaseTask, EvalResult, EvalStrategy, TaskType from promptolution.utils.prompt import Prompt @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 + sequences: np.ndarray input_tokens: np.ndarray output_tokens: np.ndarray agg_input_tokens: np.ndarray @@ -30,6 +33,7 @@ def __init__( 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") @@ -64,7 +68,6 @@ def evaluate( # type: ignore eval_strategy: Optional[EvalStrategy] = None, ) -> MultiObjectiveEvalResult: """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 @@ -124,7 +127,11 @@ def evaluate( # type: ignore # Record evaluated block for block strategies for prompt in prompts_list: - task.prompt_evaluated_blocks.setdefault(str(prompt), set()).add(task.block_idx) + block_set = task.prompt_evaluated_blocks.setdefault(str(prompt), set()) + if isinstance(task.block_idx, list): + block_set.update(task.block_idx) + else: + block_set.add(task.block_idx) per_task_results.append( EvalResult( @@ -143,9 +150,7 @@ def evaluate( # type: ignore # Mirror evaluated block bookkeeping using the first task for parity with BaseTask. first_task = self.tasks[0] - self.prompt_evaluated_blocks = { - str(p): first_task.prompt_evaluated_blocks[str(p)] for p in prompts_list - } + self.prompt_evaluated_blocks = {str(p): first_task.prompt_evaluated_blocks[str(p)] for p in prompts_list} return MultiObjectiveEvalResult( scores=stacked_scores, diff --git a/promptolution/tasks/reward_tasks.py b/promptolution/tasks/reward_tasks.py index b09a06f3..67887e79 100644 --- a/promptolution/tasks/reward_tasks.py +++ b/promptolution/tasks/reward_tasks.py @@ -3,6 +3,7 @@ import numpy as np import pandas as pd +from collections import defaultdict from typing import TYPE_CHECKING, Callable, List, Optional @@ -25,6 +26,7 @@ def __init__( df: pd.DataFrame, reward_function: Callable[[str], float], x_column: str = "x", + reward_columns: Optional[List[str]] = None, task_description: Optional[str] = None, n_subsamples: int = 30, eval_strategy: "EvalStrategy" = "full", @@ -35,8 +37,9 @@ def __init__( Args: df (pd.DataFrame): Input DataFrame containing the data. - reward_function (Callable): Function that takes a prediction and returns a reward score. Note: The optimizers aim to maximize. + 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". + 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". @@ -44,6 +47,7 @@ def __init__( 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, @@ -54,7 +58,12 @@ def __init__( config=config, ) + # 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 the score for a single reward prediction using the reward function.""" - rewards = [self.reward_function(pred) for pred in preds] + """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) diff --git a/promptolution/utils/capo_utils.py b/promptolution/utils/capo_utils.py index 404a994d..b4d8d9d2 100644 --- a/promptolution/utils/capo_utils.py +++ b/promptolution/utils/capo_utils.py @@ -6,9 +6,9 @@ from typing import List -from promptolution.utils.templates import CAPO_FEWSHOT_TEMPLATE from promptolution.utils.formatting import extract_from_tag from promptolution.utils.prompt import Prompt +from promptolution.utils.templates import CAPO_FEWSHOT_TEMPLATE def build_few_shot_examples( @@ -60,7 +60,9 @@ def perform_crossover( for _ in range(optimizer.crossovers_per_iter): mother, father = parents if len(parents) == 2 else random.sample(parents, 2) crossover_prompt = ( - optimizer.crossover_template.replace("", mother.instruction).replace("", father.instruction).strip() + optimizer.crossover_template.replace("", mother.instruction) + .replace("", father.instruction) + .strip() ) crossover_prompts.append(crossover_prompt) combined_few_shots = mother.few_shots + father.few_shots diff --git a/tests/helpers/test_helpers.py b/tests/helpers/test_helpers.py index e6de09c7..c77b2748 100644 --- a/tests/helpers/test_helpers.py +++ b/tests/helpers/test_helpers.py @@ -352,10 +352,10 @@ def test_helpers_integration(sample_df, experiment_config): def test_get_llm_variants(monkeypatch): def factory(model_name=None, config=None, **kwargs): - created['name'] = model_name or kwargs.get("model_id") - created['config'] = config + created["name"] = model_name or kwargs.get("model_id") + created["config"] = config return MockLLM() - + created = {} monkeypatch.setattr("promptolution.helpers.LocalLLM", factory) @@ -366,15 +366,15 @@ def factory(model_name=None, config=None, **kwargs): cfg.model_id = "local-foo" res = get_llm(config=cfg) assert isinstance(res, MockLLM) - assert created['name'] == "foo" + assert created["name"] == "foo" cfg.model_id = "vllm-bar" res = get_llm(config=cfg) - assert created['name'] == "bar" + assert created["name"] == "bar" cfg.model_id = "api-model" res = get_llm(config=cfg) - assert created['name'] == "api-model" + assert created["name"] == "api-model" with pytest.raises(ValueError): get_llm() diff --git a/tests/llms/test_api_llm.py b/tests/llms/test_api_llm.py index b8fc3fa4..cda8c748 100644 --- a/tests/llms/test_api_llm.py +++ b/tests/llms/test_api_llm.py @@ -79,6 +79,7 @@ def test_ainvoke_with_retries_recovers(monkeypatch): good = SimpleNamespace(choices=[SimpleNamespace(message=SimpleNamespace(content="done"))]) api = _make_api_stub(max_retries=2, retry_base_delay_s=0) api._ainvoke_once = AsyncMock(side_effect=[Exception("fail"), good]) + async def _sleep(_): return None @@ -93,6 +94,7 @@ async def _sleep(_): def test_ainvoke_with_retries_exhausts(monkeypatch): api = _make_api_stub(max_retries=1, retry_base_delay_s=0) api._ainvoke_once = AsyncMock(side_effect=[Exception("boom"), Exception("boom2")]) + async def _sleep(_): return None diff --git a/tests/mocks/dummy_config.py b/tests/mocks/dummy_config.py index 80b6d81a..cf0ac9e4 100644 --- a/tests/mocks/dummy_config.py +++ b/tests/mocks/dummy_config.py @@ -1,15 +1,21 @@ +"""Lightweight config stub used across tests.""" + + class DummyConfig: - """Lightweight config stub used across tests.""" + """Minimal config object that tracks apply/validate calls.""" def __init__(self, task_description=None): + """Initialize the dummy config with an optional task description.""" self.applied = False self.validated = False self.task_description = task_description def apply_to(self, obj): + """Mark the target object as having config applied.""" self.applied = True obj.config_applied = True obj.applied = True def validate(self): + """Record that validation was executed.""" self.validated = True diff --git a/tests/mocks/mock_task.py b/tests/mocks/mock_task.py index 9b70f30a..7566a764 100644 --- a/tests/mocks/mock_task.py +++ b/tests/mocks/mock_task.py @@ -1,8 +1,8 @@ """Mock task for testing purposes.""" +import math from unittest.mock import MagicMock -import math import pandas as pd from typing import List @@ -35,6 +35,9 @@ def __init__( Args: predetermined_scores: Dict/list/callable for score generation used by _evaluate. + config: Optional ExperimentConfig applied to the base class. + df: Optional dataframe override to seed the task. + n_subsamples: Number of subsamples to expose through BaseTask. eval_strategy: Eval strategy to expose (defaults to "full"). n_blocks: Number of blocks to report. block_idx: Current block index (int or list). @@ -42,8 +45,12 @@ def __init__( task_description: Description to attach to the task. evaluate_fn: Optional callable to replace evaluate entirely for tests. """ - base_df = df if df is not None else pd.DataFrame( - {"x": ["Sample text 1", "Sample text 2", "Sample text 3"], "y": ["positive", "negative", "neutral"]} + base_df = ( + df + if df is not None + else pd.DataFrame( + {"x": ["Sample text 1", "Sample text 2", "Sample text 3"], "y": ["positive", "negative", "neutral"]} + ) ) super().__init__( diff --git a/tests/optimizers/test_capo.py b/tests/optimizers/test_capo.py index da227088..b21b1c6a 100644 --- a/tests/optimizers/test_capo.py +++ b/tests/optimizers/test_capo.py @@ -7,7 +7,7 @@ from promptolution.optimizers.capo import CAPO from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation from promptolution.utils.prompt import Prompt -from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_FEWSHOT_TEMPLATE, CAPO_MUTATION_TEMPLATE +from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_MUTATION_TEMPLATE def test_capo_initialization(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py index 9ab5cf5c..919ee68b 100644 --- a/tests/optimizers/test_capoeira.py +++ b/tests/optimizers/test_capoeira.py @@ -3,13 +3,14 @@ import numpy as np import pandas as pd +from tests.mocks.mock_task import MockTask + from promptolution.optimizers.capoeira import Capoeira from promptolution.tasks.base_task import EvalResult from promptolution.tasks.multi_objective_task import MultiObjectiveEvalResult, MultiObjectiveTask from promptolution.utils.capo_utils import perform_crossover, perform_mutation from promptolution.utils.prompt import Prompt from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_MUTATION_TEMPLATE -from tests.mocks.mock_task import MockTask def test_capoeira_initialization(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): @@ -268,7 +269,9 @@ def fake_eval(prompts, *_, **__): "promptolution.optimizers.capoeira.perform_mutation", return_value=[Prompt("m1")] ), patch.object(optimizer, "_do_intensification") as do_int, patch.object( optimizer, "_advance_one_incumbent" - ) as adv_inc, patch.object(optimizer, "_select_survivors") as sel: + ) as adv_inc, patch.object( + optimizer, "_select_survivors" + ) as sel: optimizer._step() assert do_int.call_count == 1 @@ -319,9 +322,7 @@ def fake_eval(prompts, *_, **__): return EvalResult( scores=np.zeros((n, 1)), agg_scores=np.zeros(n), - sequences=np.array([[ - "" - ] for _ in range(n)], dtype=object), + sequences=np.array([[""] for _ in range(n)], dtype=object), input_tokens=np.zeros((n, 1)), output_tokens=np.zeros((n, 1)), agg_input_tokens=np.zeros(n), diff --git a/tests/tasks/test_base_task.py b/tests/tasks/test_base_task.py index f3387e77..16e9fa06 100644 --- a/tests/tasks/test_base_task.py +++ b/tests/tasks/test_base_task.py @@ -1,4 +1,3 @@ -import numpy as np import pandas as pd import pytest @@ -7,7 +6,6 @@ from tests.mocks.mock_predictor import MockPredictor from tests.mocks.mock_task import MockTask -from promptolution.tasks.base_task import BaseTask from promptolution.utils.prompt import Prompt @@ -143,4 +141,4 @@ def test_block_wraparound_and_get_cache_keys(): prompt = Prompt("hi") key = task._cache_key(prompt, "x", "y") - assert key[0].startswith(prompt.instruction) \ No newline at end of file + assert key[0].startswith(prompt.instruction) diff --git a/tests/tasks/test_multi_objective_task.py b/tests/tasks/test_multi_objective_task.py index 228b7c3d..76ebc811 100644 --- a/tests/tasks/test_multi_objective_task.py +++ b/tests/tasks/test_multi_objective_task.py @@ -1,11 +1,10 @@ -import pandas as pd import numpy as np - +import pandas as pd import pytest +from tests.mocks.mock_llm import MockLLM from tests.mocks.mock_predictor import MockPredictor from tests.mocks.mock_task import MockTask -from tests.mocks.mock_llm import MockLLM from promptolution.tasks.multi_objective_task import MultiObjectiveTask from promptolution.utils.prompt import Prompt diff --git a/tests/tasks/test_reward_tasks.py b/tests/tasks/test_reward_tasks.py index e8dc48d4..b6134573 100644 --- a/tests/tasks/test_reward_tasks.py +++ b/tests/tasks/test_reward_tasks.py @@ -1,3 +1,6 @@ +import pandas as pd + +from promptolution.tasks.reward_tasks import RewardTask from promptolution.utils.prompt import Prompt @@ -30,3 +33,22 @@ def test_reward_task_evaluate_with_return_seq(mock_reward_task, mock_predictor): assert result.sequences is not None assert result.sequences.shape[0] == 1 assert result.agg_input_tokens is not None + + +def test_reward_task_passes_reward_columns(): + """Ensure reward kwargs come from dataframe columns.""" + df = pd.DataFrame({"x": ["a", "b", "c"], "reward": [0.1, 0.2, 0.3]}) + + seen_rewards: list[float] = [] + def reward_fn(prediction: str, reward: float) -> float: + seen_rewards.append(reward) + return reward if prediction == "keep" else -1.0 + + task = RewardTask(df=df, reward_function=reward_fn, x_column="x", reward_columns=["reward"]) + + xs = ["a", "b", "c"] + preds = ["keep", "keep", "nope"] + scores = task._evaluate(xs, [""] * len(xs), preds) + + assert scores.tolist() == [0.1, 0.2, -1.0] + assert seen_rewards == [0.1, 0.2, 0.3] From d371a67d44df66140c94eb3ed7de8bfca0fe8ecf Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 6 Jan 2026 18:33:37 +0100 Subject: [PATCH 19/53] create scalarization fallback for multi objective --- .coverage | Bin 69632 -> 69632 bytes promptolution/optimizers/base_optimizer.py | 8 +++ promptolution/optimizers/capoeira.py | 2 + promptolution/tasks/base_task.py | 1 + promptolution/tasks/multi_objective_task.py | 29 ++++++++--- tests/tasks/test_multi_objective_task.py | 53 ++++++++++++++++++++ 6 files changed, 87 insertions(+), 6 deletions(-) diff --git a/.coverage b/.coverage index 442105beb59832cb9b8953cee070fcd08af6c906..298c04ecdc622f33dedb85f52421dc65773ac85d 100644 GIT binary patch delta 412 zcmZozz|ydQWrM2+M^UnYp{YrdymmbQc)s$y;d#JwnP(}_9G*&^cpgt4UhePQ_qmU8ujHP~J(;_i zJCi${+kopP*Kw}RTywa3xr(_WxvaRKZck|uPd4K0eKFWT*S*zTE`%f0@# zFPl+96DVq=xhf%*;R^qO*P9D^?lQ{B2m=MJgjdb_y*cmxpYl-i)8F^oKdV3f(*D`M z&;FA$`^?n$SML%{`2WBB>h|yP<-hO0&i|Vm&0PE8{%5)Ce}oSlKl?}Y568MSOp}lG i^)hbV9Mb<0IgoDz3cv$-vTuaa<`XY!k%RrmW^n-Gq_LF% delta 402 zcmZozz|ydQWrM2+M`fj9abb#K!De5NAN)wXh@I+;4=1xXsKxv8+4AY}N%8UWzU95g zdzJSn?>640y#2fxybin!Jl}cV@jT+W%Cnqj9#1t-5|0m$Aop+X$K0p5*K#l9p3dFQ zoy#4~ZNhbz>onJPt_55Zxyrd>x$L;KxFor_IsbA#+$<;%$~pOZr-dNDJqshJg8jQ^ zZ!32fZ>!$NzFDp7Ib*z~Hc+@kJB#(-rhoq#89v0lUw31FUDiD5-an;t&(E&Dwte=( zZP&7&?Yy&j+o>Fb&XtANHkH2qAGP~^jffUdtCrTPb6ew!o~Hi&E%*A{zHCNuO`xcO z=Bk8LhA+$?Hs|%+Wt5Q;28x*pubTCHbKdeuiyRm{nzdPwykZn`|-ZAKl>l&gU4zAxawHK!x$$Y>+5Blx;do(BXZE* X2o!(^?PT8wrOhW^)FKD=kImu$1g@?I diff --git a/promptolution/optimizers/base_optimizer.py b/promptolution/optimizers/base_optimizer.py index 68717b28..79f668c8 100644 --- a/promptolution/optimizers/base_optimizer.py +++ b/promptolution/optimizers/base_optimizer.py @@ -31,6 +31,8 @@ class BaseOptimizer(ABC): predictor: The predictor used for prompt evaluation (if applicable). """ + supports_multi_objective: bool = False + def __init__( self, predictor: "BasePredictor", @@ -50,6 +52,12 @@ def __init__( """ # Set up optimizer state self.prompts: List[Prompt] = [Prompt(p) for p in initial_prompts] if initial_prompts else [] + if task.task_type == "multi" and not self.supports_multi_objective: + logger.warning( + f"{self.__class__.__name__} does not support multi-objective tasks; objectives will be averaged equally.", + ) + task.activate_scalarized_objective() + self.task = task self.callbacks: List["BaseCallback"] = callbacks or [] self.predictor = predictor diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index b9ec0330..e62ef927 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -28,6 +28,8 @@ class Capoeira(BaseOptimizer): """Multi-objective variant of CAPO with Pareto-based selection.""" + supports_multi_objective = True + def __init__( self, predictor: "BasePredictor", diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 65c9b9ed..268efde7 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -70,6 +70,7 @@ def __init__( ) 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 diff --git a/promptolution/tasks/multi_objective_task.py b/promptolution/tasks/multi_objective_task.py index 32ffcb16..bc159e48 100644 --- a/promptolution/tasks/multi_objective_task.py +++ b/promptolution/tasks/multi_objective_task.py @@ -59,6 +59,11 @@ def __init__( ) self.task_type: TaskType = "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, @@ -66,7 +71,7 @@ def evaluate( # type: ignore predictor, system_prompts: Optional[str | List[str]] = None, eval_strategy: Optional[EvalStrategy] = None, - ) -> MultiObjectiveEvalResult: + ) -> 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 @@ -150,16 +155,28 @@ def evaluate( # type: ignore # Mirror evaluated block bookkeeping using the first task for parity with BaseTask. first_task = self.tasks[0] + first_result = per_task_results[0] self.prompt_evaluated_blocks = {str(p): first_task.prompt_evaluated_blocks[str(p)] for p in prompts_list} + 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=per_task_results[0].sequences, - input_tokens=per_task_results[0].input_tokens, - output_tokens=per_task_results[0].output_tokens, - agg_input_tokens=per_task_results[0].agg_input_tokens, - agg_output_tokens=per_task_results[0].agg_output_tokens, + 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 diff --git a/tests/tasks/test_multi_objective_task.py b/tests/tasks/test_multi_objective_task.py index 76ebc811..7efcae7d 100644 --- a/tests/tasks/test_multi_objective_task.py +++ b/tests/tasks/test_multi_objective_task.py @@ -6,6 +6,8 @@ from tests.mocks.mock_predictor import MockPredictor from tests.mocks.mock_task import MockTask +from promptolution.optimizers.base_optimizer import BaseOptimizer +from promptolution.tasks.base_task import BaseTask, EvalResult from promptolution.tasks.multi_objective_task import MultiObjectiveTask from promptolution.utils.prompt import Prompt @@ -72,3 +74,54 @@ def make_task(): assert np.allclose(multi_res.agg_scores[1], res2.agg_scores) assert multi_res.sequences.shape == res1.sequences.shape assert multi.prompt_evaluated_blocks[str(prompt)] == {1} + + +class ConstantTask(BaseTask): + """Simple task that returns a constant score for all predictions.""" + + def __init__(self, df: pd.DataFrame, value: float) -> None: + self._value = value + super().__init__( + df=df, + x_column="x", + y_column=None, + n_subsamples=len(df), + eval_strategy="full", + seed=0, + task_description="constant", + config=None, + ) + + def _evaluate(self, xs, ys, preds): + return np.full(len(preds), self._value, dtype=float) + + +class DummyOptimizer(BaseOptimizer): + """Non-multi-objective optimizer used to trigger fallback logic.""" + + def _pre_optimization_loop(self) -> None: + pass + + def _step(self): + return self.prompts + + +def test_multi_objective_fallback_warns_and_averages(caplog): + df = pd.DataFrame({"x": ["a", "b"]}) + t1 = ConstantTask(df.copy(), value=1.0) + t2 = ConstantTask(df.copy(), value=3.0) + mo_task = MultiObjectiveTask([t1, t2]) + + predictor = MockPredictor(llm=MockLLM(predetermined_responses=["p1", "p2"])) + + with caplog.at_level("WARNING"): + DummyOptimizer(predictor=predictor, task=mo_task) + + assert mo_task._scalarized_objective is True + assert any("averaged equally" in message for message in caplog.messages) + + result = mo_task.evaluate(prompts=[Prompt("p")], predictor=predictor) + + assert isinstance(result, EvalResult) + assert np.allclose(result.scores, np.array([[2.0, 2.0]])) + assert np.allclose(result.agg_scores, np.array([2.0])) From cb0b882c7346444b842ce1df6a3e88eaced9dffe Mon Sep 17 00:00:00 2001 From: finitearth Date: Wed, 7 Jan 2026 15:09:48 +0100 Subject: [PATCH 20/53] refine mocapo --- promptolution/optimizers/capoeira.py | 32 ++++++++++++++-------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index e62ef927..26671bdd 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -112,7 +112,6 @@ def _pre_optimization_loop(self) -> None: ) population.append(Prompt(prompt.instruction, few_shots)) - self.max_prompt_length = max(self.token_counter(p.construct_prompt()) for p in population) if population else 1 init_result = self.task.evaluate(population, self.predictor) initial_vectors = self._get_objective_vectors(init_result) fronts = self._non_dominated_sort(initial_vectors) @@ -128,20 +127,17 @@ def _step(self) -> List[Prompt]: new_challengers = perform_mutation(offsprings, self) # 2) intensify each challenger; after each, advance incumbents + prune - for chal in new_challengers: - self._do_intensification(chal) - self._advance_one_incumbent() + for challenger in new_challengers: + self._do_intensification(challenger) self._select_survivors() + self._advance_one_incumbent() - # 4) logging scores: incumbents only (optional) - if self.incumbents: - inc_result = self.task.evaluate( - prompts=self.incumbents, predictor=self.predictor, eval_strategy="evaluated" - ) - vecs_inc = self._get_objective_vectors(inc_result) - self.scores = vecs_inc[:, 0].tolist() - else: - self.scores = [] + inc_result = self.task.evaluate( + prompts=self.incumbents, predictor=self.predictor, eval_strategy="evaluated" + ) + vecs_inc = self._get_objective_vectors(inc_result) + self.scores = vecs_inc[:, 0].tolist() + self.prompts = self.incumbents return self.prompts @@ -189,10 +185,14 @@ def _do_intensification(self, challenger: Prompt) -> None: challenger_mean += (challenger_block - challenger_mean) / t incumbents_mean += (incumbent_block - incumbents_mean) / t # type: ignore - # trigger comparisons CAPO/thesis-style - if fold_vec is not None and not self._is_dominated(fold_vec, challenger_mean): + if fold_vec is None: + fold_vec = challenger_mean.copy() continue - fold_vec = challenger_mean.copy() + + if self._is_dominated(fold_vec, challenger_mean): + continue + + fold_vec = challenger_mean.copy() # TODO RENAME closest_inc = self._get_closest_incumbent(challenger_mean, incumbents_mean) # type: ignore if self._is_dominated(challenger_mean, closest_inc): From 2cd5ef249d2f8670602b656fd0ca812826b8607e Mon Sep 17 00:00:00 2001 From: finitearth Date: Thu, 8 Jan 2026 14:15:33 +0100 Subject: [PATCH 21/53] remove task description form get optimizer to fix circular imports --- promptolution/helpers.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/promptolution/helpers.py b/promptolution/helpers.py index 9f177233..4a249f7d 100644 --- a/promptolution/helpers.py +++ b/promptolution/helpers.py @@ -4,7 +4,6 @@ from promptolution.tasks.judge_tasks import JudgeTask from promptolution.tasks.reward_tasks import RewardTask -from promptolution.utils import ExperimentConfig from promptolution.utils.prompt import Prompt from promptolution.utils.prompt_creation import create_prompts_from_task_description @@ -17,6 +16,8 @@ from promptolution.tasks.base_task import TaskType from promptolution.optimizers.base_optimizer import OptimizerType from promptolution.predictors.base_predictor import PredictorType + from promptolution.utils import ExperimentConfig + import pandas as pd @@ -203,7 +204,6 @@ def get_optimizer( meta_llm: "BaseLLM", task: "BaseTask", optimizer: Optional["OptimizerType"] = None, - task_description: Optional[str] = None, config: Optional["ExperimentConfig"] = None, ) -> "BaseOptimizer": """Create and return an optimizer instance based on provided parameters. @@ -214,7 +214,6 @@ def get_optimizer( task: The task object used for evaluating prompts optimizer: String identifying which optimizer to use meta_prompt: Meta prompt text for the optimizer - task_description: Description of the task for the optimizer config: Configuration object with default parameters Returns: @@ -224,10 +223,6 @@ def get_optimizer( ValueError: If an unknown optimizer type is specified """ final_optimizer = optimizer or (config.optimizer if config else None) - if config is None: - config = ExperimentConfig() - if task_description is not None: - config.task_description = task_description if final_optimizer == "capo": return CAPO( From cfc0622a043aeaf97cd91c2f85d8a39119a7b7c5 Mon Sep 17 00:00:00 2001 From: finitearth Date: Thu, 8 Jan 2026 14:29:39 +0100 Subject: [PATCH 22/53] use task type to judge task type --- .coverage | Bin 69632 -> 69632 bytes promptolution/tasks/classification_tasks.py | 1 + promptolution/tasks/judge_tasks.py | 1 + promptolution/tasks/multi_objective_task.py | 2 +- promptolution/tasks/reward_tasks.py | 2 +- promptolution/utils/prompt_creation.py | 3 +-- 6 files changed, 5 insertions(+), 4 deletions(-) diff --git a/.coverage b/.coverage index 298c04ecdc622f33dedb85f52421dc65773ac85d..2c290b52e63162128cf4fe32f01b4e8e075e490f 100644 GIT binary patch delta 1930 zcmeH|e@xV69LJyMd*2`2_s8?S@BIdj1EREWkhH*{h#%ts?@%z~hbQN8D%BAvAtSOH zv`W^-FngLC4VcX4o}W z8oaO!ugc2A3F8LJ6NM-3cHA~@V#_R~lRVoQ8`o1Mk@$nG9DKp9gAv5D9){!E{2aW? zp|C5>#2JexFFnwnh#-qQ@)+D_3u8<~UOl^jlzS&M;uU1Rlq;Cld$v@%YwIdEdN zu1l`I%u2n;X6NFKyU}SLELTHRPPdlvBp(b|+hoh&>%=+4FGkm_n6GIL4NOCkpiV zdB5bN3XF6v|AzWJ#lix30JN)@gslzxg!`Fuu&=TsYzlLVX+(;26ii7Qg)Rw+M@5zB z5eoPTCYtV+kIoAN zN6>;jnY74L^;O%1Wky>uFfJHm(bK=*GSfUUa%vPyl~5)m@n`uj_{02ezKyTvo&1CRVxH6Rl=>I- z8K0slQ2V}rrGP@Gdt7l%uFF#m-m;6EKEz3_&n;pygT3^i16;X!ZVQ;0b^tJaA{zKV zo`K-T^&fVsn%h<28ih}_b{8Tm4b@g_oq_#L27&%GOBAWoqPpI^>DzR1=IP>PC#HKd zhMW3FGlm0Gg=P!s-e+-|cmQb7d^HHbv0#-e=wFf?NjlY$cxPKuQv?~tMAR8x3~-N%^ogzG;PXKblk@Q={jnT+nRpVa)``5^we@k4j5~Iw& nDxi5m{GrZQGs;l^6k^{uV^67WeKWl_5%>Q~a{DjIOhETLEfQ=9 delta 1954 zcmeH|ZA{fw9LLYO_y2J3|KXhb2ksLOSInRyC*VW{m8WremIqCc8+f?_Cd7kWRH9WE z=H}LhA)Xhs(v5O$iPg;8W-j14cbxdE7}#B&os3fHUA-B%meNS7sU>^lM`cDF2}-v zmgVH;PwUA@nm-#t)mD<9lt?<$MUtFuhQVZWQVLZeFo^V|*T6t>Gc$$k$uL4Q>CIV7 zI&&3LIxea5OypWdJ~WcZ%tVc|EGjx?&CXcTnE4oMpnFD=(F`l9t}Am?+pDW}^jK((6ClTl2vgUlIC5$ZJnj}+r4Uzl*LTDnZY)Ok*FQ@pG9ai!~rotLZn3hfR zEo5+#BRvy-tlq{@(r?+*bPrGtAjj-k6t0o=FJdTMT7s8?^`rNC7AnM1Y$f!W$Z)Cz zb+l_I#0v)=PE#Nvp2;i<61mmIvucP>mvW)5A%I8mS9k!wjazXuF2GhCho8V&bRBun z5b8&5=m4rjS!e@ViA+e8C*-U0CHaEJoFmU_sWGuKd69$@Re#XtWOMDQlOh_1Z0&#u|z2F>(i?Eoej5m?i*8*Kex-=Q@W&o9^ z)XzXgxS{)96M#D33=7V@&&&GxR{f`$-~tccQ=tGVEdfqF_Xh^Pfw*9Qs^k2--k5RM zj?6Z0TX%7~H(|^*FrF~xo60r$(d=qJn^6FO0nNAO0O*;kDw&&nT2N@lpd^b=!B8$n zM@6-cTsstyA& z)3fzwX5r str: """Construct the judge prompt based on whether ground truth is available.""" diff --git a/promptolution/tasks/multi_objective_task.py b/promptolution/tasks/multi_objective_task.py index bc159e48..1d557535 100644 --- a/promptolution/tasks/multi_objective_task.py +++ b/promptolution/tasks/multi_objective_task.py @@ -57,7 +57,7 @@ def __init__( seed=primary.seed, config=None, ) - self.task_type: TaskType = "multi" + self.task_type = "multi" self.tasks = tasks self._scalarized_objective: bool = False diff --git a/promptolution/tasks/reward_tasks.py b/promptolution/tasks/reward_tasks.py index 67887e79..2da558d4 100644 --- a/promptolution/tasks/reward_tasks.py +++ b/promptolution/tasks/reward_tasks.py @@ -57,7 +57,7 @@ def __init__( 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) diff --git a/promptolution/utils/prompt_creation.py b/promptolution/utils/prompt_creation.py index fd0087dc..0f0cff3d 100644 --- a/promptolution/utils/prompt_creation.py +++ b/promptolution/utils/prompt_creation.py @@ -13,7 +13,6 @@ from promptolution.llms.base_llm import BaseLLM from promptolution.tasks.base_task import BaseTask -from promptolution.tasks.classification_tasks import ClassificationTask from promptolution.utils.templates import ( PROMPT_CREATION_TEMPLATE, PROMPT_CREATION_TEMPLATE_FROM_TASK_DESCRIPTION, @@ -95,7 +94,7 @@ def create_prompts_from_samples( meta_prompts = [] for _ in range(n_prompts): - if isinstance(task, ClassificationTask) and get_uniform_labels: + if task.task_type == "classification" and get_uniform_labels: # if classification task sample such that all classes are represented unique_labels, counts = np.unique(task.ys, return_counts=True) proportions = counts / len(task.ys) From 79d654df1751a15736b101d06401e1e6413f2bd4 Mon Sep 17 00:00:00 2001 From: mo374z Date: Thu, 8 Jan 2026 15:43:20 +0100 Subject: [PATCH 23/53] change tokenizer handling to work with new hf interface --- promptolution/llms/vllm.py | 4 +--- promptolution/tasks/base_task.py | 7 ++++++- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/promptolution/llms/vllm.py b/promptolution/llms/vllm.py index f22ff528..6de4f937 100644 --- a/promptolution/llms/vllm.py +++ b/promptolution/llms/vllm.py @@ -13,7 +13,6 @@ logger = get_logger(__name__) try: - from transformers import AutoTokenizer # type: ignore from vllm import LLM, SamplingParams imports_successful = True @@ -113,8 +112,7 @@ def __init__( self.llm = LLM(**llm_params) - # Initialize tokenizer separately for potential pre-processing - self.tokenizer = AutoTokenizer.from_pretrained(model_id) + self.tokenizer = self.llm.get_tokenizer() if batch_size is None: cache_config = self.llm.llm_engine.model_executor.cache_config diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 268efde7..afc9199e 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -66,7 +66,7 @@ def __init__( 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." + 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 @@ -237,6 +237,11 @@ def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarra """ raise NotImplementedError + @abstractmethod + def activate_scalarized_objective(self) -> None: + """Activate scalarized objective for multi-objective tasks.""" + raise NotImplementedError + def evaluate( self, prompts: Union[Prompt, List[Prompt]], From 3719f8433a13c3885bf3e718d0d306c532b5c931 Mon Sep 17 00:00:00 2001 From: mo374z Date: Thu, 8 Jan 2026 16:29:31 +0100 Subject: [PATCH 24/53] change sampling params from vllm --- promptolution/llms/vllm.py | 3 ++- pyproject.toml | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/promptolution/llms/vllm.py b/promptolution/llms/vllm.py index 6de4f937..6844a604 100644 --- a/promptolution/llms/vllm.py +++ b/promptolution/llms/vllm.py @@ -13,7 +13,8 @@ logger = get_logger(__name__) try: - from vllm import LLM, SamplingParams + from vllm import LLM + from vllm.sampling_params import SamplingParams imports_successful = True except ImportError: diff --git a/pyproject.toml b/pyproject.toml index 487b398c..5b45f40b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -14,7 +14,7 @@ scikit-learn = ">=1.5.2" fastparquet = ">=2024.11.0" openai = {version = ">=1.0.0", optional = true} requests = {version = ">=2.31.0", optional = true} -vllm = {version = ">=0.10.1.1", optional = true} +vllm = {version = ">=0.13.0", optional = true} transformers = {version = ">=4.48.0", optional = true} scipy = ">=1.15" @@ -32,7 +32,7 @@ requests = ">=2.31.0" [tool.poetry.group.vllm] optional = true [tool.poetry.group.vllm.dependencies] -vllm = ">=0.10.1.1" +vllm = ">=0.13.0" [tool.poetry.group.transformers] optional = true @@ -52,7 +52,7 @@ pytest = ">=8.3.5" pytest-cov = ">=6.1.1" openai = ">=1.0.0" requests = ">=2.31.0" -vllm = "==0.10.1.1" +vllm = ">=0.13.0" transformers = ">=4.48.0" [tool.poetry.group.docs.dependencies] From 95332326ce5d54aa849240748c916a30696246d2 Mon Sep 17 00:00:00 2001 From: mo374z Date: Thu, 8 Jan 2026 17:07:36 +0100 Subject: [PATCH 25/53] change automatic batch size alignment of vllm --- promptolution/llms/vllm.py | 25 +++++++++++++------------ 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/promptolution/llms/vllm.py b/promptolution/llms/vllm.py index 6844a604..4fab6012 100644 --- a/promptolution/llms/vllm.py +++ b/promptolution/llms/vllm.py @@ -116,19 +116,20 @@ def __init__( self.tokenizer = self.llm.get_tokenizer() if batch_size is None: - cache_config = self.llm.llm_engine.model_executor.cache_config - if ( - cache_config.num_gpu_blocks is not None - and cache_config.block_size is not None - and self.max_model_len is not None - ): - self.batch_size = int( - (cache_config.num_gpu_blocks * cache_config.block_size / self.max_model_len) * 0.95 - ) - logger.info(f"🚀 Batch size set to {self.batch_size} based on GPU memory.") + max_num_seqs = int(llm_kwargs.get("max_num_seqs", 1)) + max_num_batched_tokens = llm_kwargs.get("max_num_batched_tokens", None) + + # Heuristic: if vLLM is capped by batched tokens, don't feed more seqs than fit. + if max_num_batched_tokens is not None and self.max_model_len is not None: + token_limited = max(1, int(max_num_batched_tokens) // int(self.max_model_len)) + self.batch_size = max(1, min(max_num_seqs, token_limited)) else: - self.batch_size = 1 - logger.warning("⚠️ Could not determine batch size from GPU memory. Using batch size of 1.") + self.batch_size = max(1, max_num_seqs) + + logger.info( + f"🚀 Batch size set to {self.batch_size} (max_num_seqs={max_num_seqs}, " + f"max_num_batched_tokens={max_num_batched_tokens}, max_model_len={self.max_model_len})." + ) else: self.batch_size = batch_size From 9e33e3a0d445e5e2d438cf9a7b12c7747dffc18e Mon Sep 17 00:00:00 2001 From: mo374z Date: Thu, 8 Jan 2026 19:20:10 +0100 Subject: [PATCH 26/53] remove abstractmethod --- promptolution/tasks/base_task.py | 1 - 1 file changed, 1 deletion(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index afc9199e..b2e378eb 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -237,7 +237,6 @@ def _evaluate(self, xs: List[str], ys: List[str], preds: List[str]) -> np.ndarra """ raise NotImplementedError - @abstractmethod def activate_scalarized_objective(self) -> None: """Activate scalarized objective for multi-objective tasks.""" raise NotImplementedError From 016b298f73ebcfc200ecbe5f6354ab557d296c16 Mon Sep 17 00:00:00 2001 From: finitearth Date: Fri, 9 Jan 2026 14:25:59 +0100 Subject: [PATCH 27/53] allow for tournament select --- .coverage | Bin 69632 -> 53248 bytes promptolution/optimizers/capoeira.py | 102 +++++++++++++++++++-------- promptolution/utils/capo_utils.py | 8 ++- tests/optimizers/test_capoeira.py | 28 ++++---- 4 files changed, 91 insertions(+), 47 deletions(-) diff --git a/.coverage b/.coverage index 2c290b52e63162128cf4fe32f01b4e8e075e490f..79b58f0563fb5933afab5f66c61780f23110febb 100644 GIT binary patch delta 1396 zcmY*Ydu&rx82`?F^|rnDoDSF{TWRmQN1-EQk^q5CdrO!EkRTlhbjGg1DI1Sr1qb1e zT^`Y32y7klGVVYGMpv^ZNdv__wE-iPlR?}H$oz~Bh z;$5Xj1i7+$^~RdDQ^N7u!5!iU{X)>ol>ixwPvDidM(Y)!ORx&7_{H2qZVy*tEs)FQ zZGg0sD(icePeiw^u(ja^qLEii zP4;7um-bu?qkLM_8%B9_AoK*bEC1n|OZ)FG6LTWMmQtEs7BJK7^rJuuWBs&>$=AZ@ zk7hsP%QmI4?CG?5p)|@{bcX~+`{VGOc-NBuQ4V=Tdi0t~`{vkbYu~~>NcdE!(1KoN zE0AlXhkQ#8k$SS7l#}^HCsW8cJd6kMFZdGv8n@$pcn3Cc6tBRG@e6n=&QpI=e^fiv z&()whSDnDtF|2edF=d~!HJ;*_q|}Y`3drj71}9e-BLqM;Kn7WG}C?V*w6rkXzz52lKf@{<`ji`~AuF z=lqXEpabuF>Jmd2o04yzY_WWX68&NQ26MpL(qn0*ZD z+U1<=OiWPI8SYH4jzA9^N>kHJGc8#9ZsJ!||C5r+%*@nS5$FJjuGuxQ&K!dYGG794 z_7oP8kuo?Jxc-V12QzXVorg1?_ytzbMi`VvW=%wzWW%P8PtSvd#ip9N`o>Q9YI2&w z^g<<;hGqct`_0Oxj5zHKpK13oowbv}!7i7M-)##jHs;8*#qJn2r&^*`E3?K~eW=2U zWv061n6g!8!vW{DF#Baov>-!v7$77dn^ofFf5+#ygIuI(1{~pKYV6La>i`*DMaE`P z(FB$^OX%FNJn1$^F{!fcKYdtXhs3i8%y&9Cc6vD9Hppv-o?;y#ClajIlEDL_Qq~BK z#_B(dKmtAikO2xv>^8xR?*kbk56N9}gY=P0|B0B+D0{DIYn!MPOG+xRXVkz;x=wnP%`&D5{Rgsb7uaSb7tnBcfar6 z-*@k~-1pskKkHVmbu0dc9CO*N=v9+Egb+pz!56tuLo}&L@Z5*%am3xo_2%s6=r}xV zCzge7p)Iww@tpin`tO;~(jPrbi3~=QF1&J8j*{LzJ(m}Pqna6XT6z|p_v{2Ln67<} zLAhci5GzI`Ohhc^qRvOAGOLZo#*-&aI zuBtGZBTTw-eMQ+Q1MrppfHyFB)f4cGb!LP1F+MZQ@3)4Hn$ekxUvjp#*cXOnFxC{0 zX*EYUt!>^E`tYn+dh2ZdWc%78sto2T?bt$#6`(2$RH+zLNU^T8bg{0u#5{rzPBRsCv#G&SJ`yJVd@k)Qiq?_*gd9g6e3+94}vVTdEf~OG(K-30I;M zXA?Jha=eZD3^!0l(FVbLM0jm&0(~=o2ICS!yGw-9AI;0_3d(3K8 zJZ%7Mhddt#XMib=IK%u$oFU(rCxr7S7vnWE9T++~$7@d!uCSzwen8YuR5#T@y+^%H zS*Zoo3@VWdqD161@+b0J@=J0L`4{r{WHnhv7Lv2b6f%hPA*Iq&(gV^)=|<^Nsaj%_ z+>>0B{3JOc*(P~cvRYCk$&`@dU&X!R4so-%QT(R(Rq--$t~iP~LmVYOCf+4B5>-S# zkwl1s2pPc<-4XrJ>LSY)Sth3Am?$=KEF>1i6)4Zb-0V7a>*dSUA8jeGm9b)CL`zwu#H-rq|4C!_*9Fg4S%dbRq; z?>_DM?$hs9Pfh_|a!UD?4|-yk-@Ti8S5%VPefrKW`DLxxQa*rrC4*d*e9zRlwCCU} z=PnMcZS3poY-&1>c_lHDq%Hbe$2u<5Unstn5!j?x(WPJUq=_K$PE4%7cvRnIzVDs@ z-LQmZgPUes-h5|W_AGVgHg%V1LOieo;~R!w9_)BC!}eCjV28gf4tTzC4a4ydkKXB+ z_~!d^dVVYq6V)hV^VOEFrnl9$s!sL5-Kq3TfA+O>4V`STnGZJ%Z2xfWxo&lnZJ+k| z(2uulXK#EpRNu!9mt5j%4IEW;^XZ3;M~|-aPdIkWKY_Mwvfal?6-azk%KaDSZMIq7 zQY8qZfH5IzPfI8f2`!Jvi*XgH)T9V#$44yMcbV>aJ*e+wcR}aQ&XRLy3r?LbEAR{l zVPyE+h$8DSg&q(s?Sy`C?bIzIkEtLCpE}nwmygikjgxD(BaL-cd!btx@ReaO0HI5V zzZbWE@=Du5PGYD9*wvxSc*A=JZ35K3=HULtw!sqz3jHjFeYdNPx#!D2J(ItyzR%Qs zM%jAAf7_RN$4~9Fm0UY?8*@`KDwWM$pJ+|L7#mV9vCz>8yBIgJ!b|i9{syGD_fVGXvxW;+VccKfRP0aa>4?k zkso-4r?DNzd$Z8hyFF%fb07<$f#Y&?ZT|xd4cMH1pP{wsQ!7 z=y>?g<*m0G1b&P@zlJn{FVpZngQZS^#=0qONmg4o|Epn3v8`i6x9jG$_px8^FAP6c z_e*_3+Hh@Oa^I58o&NW`Y_)wCG7eqe|MvI!i{{VoENd+bRDgn6icdLL`wSm^eJs1d ze9Js5nr_`-(b1h>$3$}jz)# z8TyHnn;V8&%g;YpF+cA3z+v@8OUGsP#o^nty}UuH^3L}p5Wwq=F3YoXI7 zfuwZOp5?lI>V7pkaCl9ddy&ny;^f`|?TVrE{S^<(3b?P_Jb&Bg`vgLfczev9@mhO7_E8V%4qYGj=18Pzf9}X(4sG2SVmVTG+4laA=~%ngRm(#h zu?t|lTm~}`r?+T=3_5|bwhz4!CL12&p}ji=1O+?~Dm?^K@TDZQy~)or(H8==gF{Fe zk#x!H=zw~n1lqBZW)Ca@UHV0hQ02}w!=DulqpKY)-eQm?ij{CrV-|#jNQM-Mqfsan znlzzM=t6+lpJ={ntvIsFgxafmI3f|S+(b&ONJI&tsSq|x+7_c_)UgDA9}e4l)m+(K?5|3KD|t4T9iN*c&TWDYrp z%plXrBr*nq$cIu3vnDJ(oh4~$EYWCKqE@pcHI*ePDJ)4&W=T>KOA-@Vl90fX_;{AY z#jzwdmL=1su_PvjCDG9=QK?uG6~$*_4i*{7b|WHK5+2TysZ&`J7RHj$P?ji_EC~r= zNpNuM-y5l@MFeg@o~RDPG45pUrY9$E{7Bw6=(Z?SSRgzoxFyKqzrkBcx=8o^i@DrL z&L);tRGKOvrwIR_$?N?7NnXdlgq1kbw*PQ$=VCIf(3$kwDxJB+9InM{(%6g+5^py{ zHFI!fl~vlMOu07IStuJ5>J1UG{?SRm)0ve?_xL+ zj^;Z!ba$SXN8Cgu7#*FNJX&?oEAm(2l##oVRdpka-PJncS<~I&Lhdmt*AwW2dk+!ouCPatya8#%#*%yS&O+ zYOlWRTUBf_IO?wS7x0<@`x^P~k3nJA(qPxvR9wz(|432x$bnR@(%841H7It%(LT?4 zVrrT!WHn57k}O-UUus}a6l|mx;Wg78v+#F96_@JFX5$iLF*{$g-2*CDrHob`PGGhI zV)~#l`H?NCG#T{9kxHAbR>%4POtdkKpAy}Q}JW0lY2RpHf$|_@-akaq&keFF`jodK@ zkJ8dI`=M7u6U5Q;@R|^Zc+!}-v;s;j?FdFS7q1DjViFL6nn_ zKXFU}JCTv`k&k%O7(`cTv|EUv!}HUe{~ll5UKcaI<07;lpO=08F7Vq)@v((z9fGGm zKTMg4*SLxt&kqzcd*;Yt{{L1T>it`FD3tqc@R#J`g?LR0Y+o7;7XdwMp)0In#z;Gx zPZNdtaK>5Wk_uCquBysd&TOi56tb(3U>+7QC1}c`@kM&aU(A-b^cV9x1yZ;~LeUml zpwayQ6$_rP#=lwocf4Z#k+-=BS9uEgh}ueNAU=KHxV<{80V=nMm!%21ocIE8~b5r7% None: initial_vectors = self._get_objective_vectors(init_result) fronts = self._non_dominated_sort(initial_vectors) self.incumbents = [population[i] for i in fronts[0]] - self.challengers = [population[i] for front in fronts[1:] for i in front] + self.non_incumbents = [population[i] for front in fronts[1:] for i in front] # keep self.prompts as a "view" if base class expects it self.scores = initial_vectors[:, 0].tolist() def _step(self) -> List[Prompt]: - # 1) generate challengers (random parent selection happens inside perform_crossover) - offsprings = perform_crossover(self.prompts, self) + # 1) generate challengers + offsprings = perform_crossover(self.prompts, self, self._tournament_selection) new_challengers = perform_mutation(offsprings, self) # 2) intensify each challenger; after each, advance incumbents + prune @@ -142,10 +147,6 @@ def _step(self) -> List[Prompt]: return self.prompts def _do_intensification(self, challenger: Prompt) -> None: - if not self.incumbents: - self.incumbents.append(challenger) - return - common_blocks = self._get_common_blocks(self.incumbents) # bootstrap if no common blocks yet @@ -163,7 +164,7 @@ def _do_intensification(self, challenger: Prompt) -> None: incumbents_mean: Optional[np.ndarray] = None t = 0 - fold_vec: Optional[np.ndarray] = None + fold_vec = np.full((self.n_objectives,), -np.inf) while remaining_blocks: b = random.choice(tuple(remaining_blocks)) @@ -185,27 +186,29 @@ def _do_intensification(self, challenger: Prompt) -> None: challenger_mean += (challenger_block - challenger_mean) / t incumbents_mean += (incumbent_block - incumbents_mean) / t # type: ignore - if fold_vec is None: - fold_vec = challenger_mean.copy() - continue - if self._is_dominated(fold_vec, challenger_mean): continue fold_vec = challenger_mean.copy() # TODO RENAME - closest_inc = self._get_closest_incumbent(challenger_mean, incumbents_mean) # type: ignore - if self._is_dominated(challenger_mean, closest_inc): - self.challengers.append(challenger) + closest_incumbent = self._get_closest_incumbent(challenger) # type: ignore + if self._is_dominated(challenger_mean, closest_incumbent): + self.non_incumbents.append(challenger) return # survived all common blocks -> admit and update front restricted to common_blocks self.incumbents.append(challenger) self._update_incumbent_front(blocks=common_blocks) - def _get_closest_incumbent(self, challenger_vec: np.ndarray, incumbent_vecs: np.ndarray) -> np.ndarray: + def _get_closest_incumbent(self, challenger) -> np.ndarray: """Return the vector of the geometrically closest incumbent.""" - all_vecs = np.vstack([incumbent_vecs, challenger_vec[None, :]]) + challenger_res = self.task.evaluate(challenger, self.predictor, eval_strategy="evaluated") + challenger_vec = self._get_objective_vectors(challenger_res) + + incumbent_res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") + incumbent_vecs = self._get_objective_vectors(incumbent_res) + + all_vecs = np.vstack([incumbent_vecs, challenger_vec]) min_b = np.min(all_vecs, axis=0) max_b = np.max(all_vecs, axis=0) rng = max_b - min_b @@ -219,9 +222,6 @@ def _get_closest_incumbent(self, challenger_vec: np.ndarray, incumbent_vecs: np. return incumbent_vecs[idx] def _update_incumbent_front(self, blocks: Optional[set[int]] = None) -> None: - if not self.incumbents: - return - if blocks is None: res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") else: @@ -236,7 +236,7 @@ def _update_incumbent_front(self, blocks: Optional[set[int]] = None) -> None: demoted = [self.incumbents[i] for front in fronts[1:] for i in front] self.incumbents = new_incumbents - self.challengers.extend(demoted) + self.non_incumbents.extend(demoted) def _get_objective_vectors(self, result) -> np.ndarray: # If the task is multi-objective, include all objective dimensions, else single objective. @@ -288,10 +288,10 @@ def _advance_one_incumbent(self) -> None: def _select_survivors(self) -> None: """Prune population via Pareto logic to enforce size constraints.""" - while len(self.incumbents) + len(self.challengers) > self.population_size: - if len(self.challengers) > 0: + while len(self.incumbents) + len(self.non_incumbents) > self.population_size: + if len(self.non_incumbents) > 0: # 1. Check Heterogeneity (Fairness Check) - chal_blocks_map = self.task.get_evaluated_blocks(self.challengers) + chal_blocks_map = self.task.get_evaluated_blocks(self.non_incumbents) block_sets = list(chal_blocks_map.values()) first_set = block_sets[0] @@ -308,14 +308,14 @@ def _select_survivors(self) -> None: candidates = [i for i, c in enumerate(counts) if c == min_count] victim_idx = random.choice(candidates) - self.challengers.pop(victim_idx) + self.non_incumbents.pop(victim_idx) continue # CASE B: Homogeneous (Fair comparison). # Use NDS + Crowding Distance. # Get objective vectors for all challengers (safe because blocks are identical) - res = self.task.evaluate(self.challengers, self.predictor, eval_strategy="evaluated") + res = self.task.evaluate(self.non_incumbents, self.predictor, eval_strategy="evaluated") vecs = self._get_objective_vectors(res) # Perform Non-Dominated Sort @@ -334,7 +334,7 @@ def _select_survivors(self) -> None: # Map back to the main challenger list index victim_idx = worst_front_indices[local_worst_idx] - self.challengers.pop(victim_idx) + self.non_incumbents.pop(victim_idx) continue # --- PRUNE FROM INCUMBENTS --- @@ -357,6 +357,46 @@ def _get_common_blocks(self, prompts: List[Prompt]) -> set: common = set.intersection(*block_sets) return common + + def _select_parent_from_pool(self, selection_pool: List[Prompt]) -> Prompt: + """Tournament-pick a parent, preferring incumbents and using crowding for ties.""" + p1, p2 = random.sample(selection_pool, 2) + + if p1 in self.incumbents and p2 in self.incumbents: + return self._pick_incumbent_by_crowding(p1, p2) + if p1 in self.incumbents: + return p1 + if p2 in self.incumbents: + return p2 + + return random.choice((p1, p2)) + + + def _pick_incumbent_by_crowding(self, p1: Prompt, p2: Prompt) -> Prompt: + """Break incumbent ties using crowding distance over common evaluated blocks.""" + res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") + inc_vectors = self._get_objective_vectors(res) + inc_distances = self._calculate_crowding_distance(inc_vectors) + + p1_idx = self.incumbents.index(p1) + p2_idx = self.incumbents.index(p2) + if inc_distances[p1_idx] > inc_distances[p2_idx]: + return p1 + if inc_distances[p2_idx] > inc_distances[p1_idx]: + return p2 + return random.choice((p1, p2)) + + + def _tournament_selection(self) -> Tuple[Prompt, Prompt]: + """Pick two distinct parents via tournament selection.""" + selection_pool = self.incumbents + self.non_incumbents + parent1 = self._select_parent_from_pool(selection_pool) + + parent2 = self._select_parent_from_pool(selection_pool) + while parent2 == parent1: + parent2 = self._select_parent_from_pool(selection_pool) + + return parent1, parent2 @staticmethod def _non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: diff --git a/promptolution/utils/capo_utils.py b/promptolution/utils/capo_utils.py index b4d8d9d2..9bda0c48 100644 --- a/promptolution/utils/capo_utils.py +++ b/promptolution/utils/capo_utils.py @@ -4,7 +4,7 @@ import random -from typing import List +from typing import List, Optional, Callable from promptolution.utils.formatting import extract_from_tag from promptolution.utils.prompt import Prompt @@ -53,12 +53,16 @@ def build_few_shot_examples( def perform_crossover( parents: List[Prompt], optimizer, + parent_select_func: Optional[Callable] = None, ) -> List[Prompt]: """Generate crossover offspring prompts.""" crossover_prompts: List[str] = [] offspring_few_shots: List[List[str]] = [] for _ in range(optimizer.crossovers_per_iter): - mother, father = parents if len(parents) == 2 else random.sample(parents, 2) + if parent_select_func: + mother, father = parent_select_func(parents) + else: + mother, father = random.sample(parents, 2) crossover_prompt = ( optimizer.crossover_template.replace("", mother.instruction) .replace("", father.instruction) diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py index 919ee68b..09955116 100644 --- a/tests/optimizers/test_capoeira.py +++ b/tests/optimizers/test_capoeira.py @@ -146,13 +146,13 @@ def fake_evaluate(*_, **__): c1, c2 = Prompt("c1"), Prompt("c2") task.eval_blocks = {str(c1): {0}, str(c2): {0, 1}} optimizer.incumbents = [Prompt("i1"), Prompt("i2")] - optimizer.challengers = [c1, c2] + optimizer.non_incumbents = [c1, c2] optimizer.population_size = 3 optimizer._select_survivors() - assert len(optimizer.challengers) == 1 - assert optimizer.challengers[0].instruction == "c2" + assert len(optimizer.non_incumbents) == 1 + assert optimizer.non_incumbents[0].instruction == "c2" def test_capoeira_select_survivors_homogeneous_prunes_lowest(mock_meta_llm, mock_predictor): @@ -191,13 +191,13 @@ def fake_evaluate(prompts, *_, **__): ) optimizer.incumbents = [Prompt("inc")] # keeps population pressure - optimizer.challengers = [c1, c2] + optimizer.non_incumbents = [c1, c2] optimizer.population_size = 2 optimizer._select_survivors() - assert len(optimizer.challengers) == 1 - assert optimizer.challengers[0].instruction == "c2" + assert len(optimizer.non_incumbents) == 1 + assert optimizer.non_incumbents[0].instruction == "c2" def test_capoeira_select_survivors_prefers_lower_cost(mock_meta_llm, mock_predictor): @@ -230,13 +230,13 @@ def fake_evaluate(prompts, *_, **__): ) optimizer.incumbents = [] - optimizer.challengers = [Prompt("cheap"), Prompt("expensive")] + optimizer.non_incumbents = [Prompt("cheap"), Prompt("expensive")] optimizer.population_size = 1 optimizer._select_survivors() - assert len(optimizer.challengers) == 1 - assert optimizer.challengers[0].instruction == "cheap" + assert len(optimizer.non_incumbents) == 1 + assert optimizer.non_incumbents[0].instruction == "cheap" def test_capoeira_step_invokes_hooks(mock_meta_llm, mock_predictor, mock_df): @@ -408,7 +408,7 @@ def test_capoeira_do_intensification_dominated_challenger(monkeypatch, mock_meta ): optimizer._do_intensification(challenger) - assert challenger in optimizer.challengers + assert challenger in optimizer.non_incumbents assert challenger not in optimizer.incumbents @@ -441,7 +441,7 @@ def fake_eval(prompts, *_, **__): optimizer._update_incumbent_front() assert optimizer.incumbents == [inc1] - assert inc2 in optimizer.challengers + assert inc2 in optimizer.non_incumbents def test_capoeira_advance_one_incumbent_no_gapblocks(mock_meta_llm, mock_predictor): @@ -552,13 +552,13 @@ def test_capoeira_select_survivors_heterogeneous_removes_lowest(mock_meta_llm, m df_few_shots=task.pop_datapoints(n=1), ) optimizer.incumbents = [] - optimizer.challengers = [c1, c2] + optimizer.non_incumbents = [c1, c2] optimizer.population_size = 1 with patch("random.choice", side_effect=lambda seq: list(seq)[0]): optimizer._select_survivors() - assert len(optimizer.challengers) == 1 + assert len(optimizer.non_incumbents) == 1 def test_capoeira_select_survivors_incumbent_only(mock_meta_llm, mock_predictor): @@ -587,7 +587,7 @@ def fake_eval(prompts, *_, **__): df_few_shots=task.pop_datapoints(n=1), ) optimizer.incumbents = [inc1, inc2] - optimizer.challengers = [] + optimizer.non_incumbents = [] optimizer.population_size = 1 optimizer._select_survivors() From e2cd177e3f356891ff06824a6be237fb404c295f Mon Sep 17 00:00:00 2001 From: finitearth Date: Fri, 9 Jan 2026 14:30:16 +0100 Subject: [PATCH 28/53] change init of f_old to inf --- promptolution/optimizers/capoeira.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 1c9ef137..a66a05bb 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -164,7 +164,7 @@ def _do_intensification(self, challenger: Prompt) -> None: incumbents_mean: Optional[np.ndarray] = None t = 0 - fold_vec = np.full((self.n_objectives,), -np.inf) + fold_vec = np.full((self.n_objectives,), np.inf) while remaining_blocks: b = random.choice(tuple(remaining_blocks)) From 8beca49a04bbe2b36c886657b04bcdda3b6dfdb0 Mon Sep 17 00:00:00 2001 From: finitearth Date: Mon, 12 Jan 2026 15:51:28 +0100 Subject: [PATCH 29/53] impelment comments --- promptolution/optimizers/capoeira.py | 58 ++++++++++++++++++++++++++-- 1 file changed, 55 insertions(+), 3 deletions(-) diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index a66a05bb..db22ead9 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -147,6 +147,12 @@ def _step(self) -> List[Prompt]: return self.prompts def _do_intensification(self, challenger: Prompt) -> None: + if challenger in self.incumbents: + return + if challenger in self.non_incumbents: + # remove from non-incumbents to re-evaluate + self.non_incumbents.remove(challenger) + common_blocks = self._get_common_blocks(self.incumbents) # bootstrap if no common blocks yet @@ -330,8 +336,9 @@ def _select_survivors(self) -> None: dists = self._calculate_crowding_distance(worst_front_vecs) # Find index relative to the worst front list - local_worst_idx = int(np.argmin(dists)) - # Map back to the main challenger list index + min_dist = np.min(dists) + tied_indices = np.where(dists == min_dist)[0] + local_worst_idx = np.random.choice(tied_indices) victim_idx = worst_front_indices[local_worst_idx] self.non_incumbents.pop(victim_idx) @@ -369,12 +376,52 @@ def _select_parent_from_pool(self, selection_pool: List[Prompt]) -> Prompt: if p2 in self.incumbents: return p2 + # both are non-incumbents + blocks_map = self.task.get_evaluated_blocks([p1, p2]) + blocks1 = blocks_map.get(str(p1), set()) + blocks2 = blocks_map.get(str(p2), set()) + + if blocks1 == blocks2: # both evaluated on same blocks + # use NDS + Crowding Distance + self.task.set_block_idx(list(sorted(blocks1))) + res = self.task.evaluate([p1, p2], self.predictor) + # check if dominated + vecs = self._get_objective_vectors(res) + if self._is_dominated(vecs[0], vecs[1]): + return p2 + if self._is_dominated(vecs[1], vecs[0]): + return p1 + # tie-breaker: crowding distance + distances = self._calculate_crowding_distance(vecs) + if distances[0] > distances[1]: + return p1 + if distances[1] > distances[0]: + return p2 + + # same crowding distance: random + + # use weaker dominance definition + # eval on common blocks only + common_blocks = blocks1 & blocks2 + if common_blocks: + self.task.set_block_idx(list(sorted(common_blocks))) + res = self.task.evaluate([p1, p2], self.predictor) + vecs = self._get_objective_vectors(res) + + if self._is_weakly_dominated(vecs[0], vecs[1]): + return p2 + if self._is_weakly_dominated(vecs[1], vecs[0]): + return p1 + return random.choice((p1, p2)) def _pick_incumbent_by_crowding(self, p1: Prompt, p2: Prompt) -> Prompt: """Break incumbent ties using crowding distance over common evaluated blocks.""" - res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") + common_blocks = self._get_common_blocks([p1, p2]) + if common_blocks: + self.task.set_block_idx(common_blocks) + res = self.task.evaluate(self.incumbents, self.predictor) inc_vectors = self._get_objective_vectors(res) inc_distances = self._calculate_crowding_distance(inc_vectors) @@ -430,6 +477,11 @@ def _non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: def _is_dominated(vec1, vec2): """Returns True if vec2 dominates vec1 in a maximize-all setting.""" return np.all(vec2 >= vec1) and np.any(vec2 > vec1) + + @staticmethod + def _is_weakly_dominated(vec1, vec2): + """Returns True if vec2 weakly dominates vec1 in a maximize-all setting.""" + return np.all(vec2 >= vec1) @staticmethod def _calculate_crowding_distance(obj_vectors: np.ndarray) -> np.ndarray: From f85062b4619dbb8a644a929dc9f8286a5413f47d Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 13 Jan 2026 13:00:45 +0100 Subject: [PATCH 30/53] incoeprated comments --- promptolution/optimizers/capoeira.py | 24 ++++++++++-------------- promptolution/utils/prompt.py | 16 ++++++++++++++++ 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index db22ead9..05343667 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -400,17 +400,18 @@ def _select_parent_from_pool(self, selection_pool: List[Prompt]) -> Prompt: # same crowding distance: random - # use weaker dominance definition - # eval on common blocks only - common_blocks = blocks1 & blocks2 - if common_blocks: - self.task.set_block_idx(list(sorted(common_blocks))) + # weaker dominance: larger eval set may dominate smaller on the smaller's blocks + elif blocks1.issubset(blocks2) and blocks1: + self.task.set_block_idx(list(sorted(blocks1))) res = self.task.evaluate([p1, p2], self.predictor) vecs = self._get_objective_vectors(res) - - if self._is_weakly_dominated(vecs[0], vecs[1]): + if self._is_dominated(vecs[0], vecs[1]): return p2 - if self._is_weakly_dominated(vecs[1], vecs[0]): + elif blocks2.issubset(blocks1) and blocks2: + self.task.set_block_idx(list(sorted(blocks2))) + res = self.task.evaluate([p1, p2], self.predictor) + vecs = self._get_objective_vectors(res) + if self._is_dominated(vecs[1], vecs[0]): return p1 return random.choice((p1, p2)) @@ -418,7 +419,7 @@ def _select_parent_from_pool(self, selection_pool: List[Prompt]) -> Prompt: def _pick_incumbent_by_crowding(self, p1: Prompt, p2: Prompt) -> Prompt: """Break incumbent ties using crowding distance over common evaluated blocks.""" - common_blocks = self._get_common_blocks([p1, p2]) + common_blocks = self._get_common_blocks(self.incumbents) if common_blocks: self.task.set_block_idx(common_blocks) res = self.task.evaluate(self.incumbents, self.predictor) @@ -477,11 +478,6 @@ def _non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: def _is_dominated(vec1, vec2): """Returns True if vec2 dominates vec1 in a maximize-all setting.""" return np.all(vec2 >= vec1) and np.any(vec2 > vec1) - - @staticmethod - def _is_weakly_dominated(vec1, vec2): - """Returns True if vec2 weakly dominates vec1 in a maximize-all setting.""" - return np.all(vec2 >= vec1) @staticmethod def _calculate_crowding_distance(obj_vectors: np.ndarray) -> np.ndarray: diff --git a/promptolution/utils/prompt.py b/promptolution/utils/prompt.py index 935521f8..bc76de57 100644 --- a/promptolution/utils/prompt.py +++ b/promptolution/utils/prompt.py @@ -48,6 +48,22 @@ def __str__(self) -> str: """Return the string representation of the prompt.""" return self.construct_prompt() + def __eq__(self, other: object) -> bool: + """Structural equality for use in lists, sets, and dict keys.""" + if not isinstance(other, Prompt): + return False + return ( + self.instruction == other.instruction + and self.few_shots == other.few_shots + and self.downstream_template == other.downstream_template + ) + + def __hash__(self) -> int: + return hash((self.instruction, tuple(self.few_shots), self.downstream_template)) + + + + def sort_prompts_by_scores( prompts: List[Prompt], scores: Union[Sequence[float], np.ndarray], top_k: Optional[int] = None From 52d29f636f8b194120f434e9d7923cad2b514d92 Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 13 Jan 2026 16:54:38 +0100 Subject: [PATCH 31/53] fix parent selection --- promptolution/optimizers/capoeira.py | 80 +++++++++++++++++----------- promptolution/tasks/base_task.py | 10 ++-- 2 files changed, 53 insertions(+), 37 deletions(-) diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py index 05343667..5b913f16 100644 --- a/promptolution/optimizers/capoeira.py +++ b/promptolution/optimizers/capoeira.py @@ -89,7 +89,7 @@ def __init__( self.population_size = len(self.prompts) if self.task.task_type == "multi": - self.n_objectives = len(self.task.tasks) + 1 # +1 for cost objective + self.n_objectives = len(self.task.tasks) + 1 # +1 for cost objective # type: ignore else: self.n_objectives = 2 # single objective + cost objective @@ -170,7 +170,7 @@ def _do_intensification(self, challenger: Prompt) -> None: incumbents_mean: Optional[np.ndarray] = None t = 0 - fold_vec = np.full((self.n_objectives,), np.inf) + old_scores = np.full((self.n_objectives,), np.inf) while remaining_blocks: b = random.choice(tuple(remaining_blocks)) @@ -192,10 +192,10 @@ def _do_intensification(self, challenger: Prompt) -> None: challenger_mean += (challenger_block - challenger_mean) / t incumbents_mean += (incumbent_block - incumbents_mean) / t # type: ignore - if self._is_dominated(fold_vec, challenger_mean): + if self._is_dominated(old_scores, challenger_mean): continue - fold_vec = challenger_mean.copy() # TODO RENAME + old_scores = challenger_mean.copy() # type: ignore closest_incumbent = self._get_closest_incumbent(challenger) # type: ignore if self._is_dominated(challenger_mean, closest_incumbent): @@ -262,7 +262,7 @@ def _advance_one_incumbent(self) -> None: if not self.incumbents: return - blocks_map = self.task.get_evaluated_blocks(self.incumbents) # Dict[str -> Set[int]] + blocks_map = self.task.get_evaluated_blocks(self.incumbents) inc_keys = [str(inc) for inc in self.incumbents] # least evaluated incumbents @@ -274,9 +274,9 @@ def _advance_one_incumbent(self) -> None: # union over incumbents union_blocks: set[int] = set() for inc in self.incumbents: - union_blocks |= set(blocks_map[str(inc)]) + union_blocks |= set(blocks_map[inc]) - chosen_blocks = set(blocks_map[str(chosen_inc)]) + chosen_blocks = set(blocks_map[chosen_inc]) # gap-first, else brand-new gap_blocks = union_blocks - chosen_blocks @@ -297,8 +297,8 @@ def _select_survivors(self) -> None: while len(self.incumbents) + len(self.non_incumbents) > self.population_size: if len(self.non_incumbents) > 0: # 1. Check Heterogeneity (Fairness Check) - chal_blocks_map = self.task.get_evaluated_blocks(self.non_incumbents) - block_sets = list(chal_blocks_map.values()) + blocks_map = self.task.get_evaluated_blocks(self.non_incumbents) + block_sets = list(blocks_map.values()) first_set = block_sets[0] # Are all challengers evaluated on the exact same set of blocks? @@ -354,15 +354,16 @@ def _select_survivors(self) -> None: victim_idx = int(np.argmin(dists)) self.incumbents.pop(victim_idx) - def _get_common_blocks(self, prompts: List[Prompt]) -> set: + def _get_common_blocks(self, prompts: List[Prompt]) -> set[int]: """Get the set of block indices that have been evaluated by all given prompts.""" - per_prompt = self.task.get_evaluated_blocks(prompts) # Dict[prompt -> Set[int]] + per_prompt = self.task.get_evaluated_blocks(prompts) # Dict[prompt -> Collection[int]] block_sets = list(per_prompt.values()) if not block_sets: return set() - common = set.intersection(*block_sets) + # Some task implementations may return lists instead of sets; normalize for typing and correctness. + common = set.intersection(*(set(s) for s in block_sets)) return common def _select_parent_from_pool(self, selection_pool: List[Prompt]) -> Prompt: @@ -378,27 +379,42 @@ def _select_parent_from_pool(self, selection_pool: List[Prompt]) -> Prompt: # both are non-incumbents blocks_map = self.task.get_evaluated_blocks([p1, p2]) - blocks1 = blocks_map.get(str(p1), set()) - blocks2 = blocks_map.get(str(p2), set()) + blocks1 = blocks_map[p1] + blocks2 = blocks_map[p2] - if blocks1 == blocks2: # both evaluated on same blocks - # use NDS + Crowding Distance - self.task.set_block_idx(list(sorted(blocks1))) - res = self.task.evaluate([p1, p2], self.predictor) - # check if dominated - vecs = self._get_objective_vectors(res) - if self._is_dominated(vecs[0], vecs[1]): - return p2 - if self._is_dominated(vecs[1], vecs[0]): - return p1 - # tie-breaker: crowding distance - distances = self._calculate_crowding_distance(vecs) - if distances[0] > distances[1]: - return p1 - if distances[1] > distances[0]: - return p2 - - # same crowding distance: random + if blocks1 == blocks2: # both evaluated on same blocks + # Use full NDS + crowding on all non-incumbents that share this block set + blocks_map = self.task.get_evaluated_blocks(self.non_incumbents) + same_block = [p for p in self.non_incumbents if blocks_map[p] == blocks1] + + if len(same_block) >= 2: + self.task.set_block_idx(list(sorted(blocks1))) + res = self.task.evaluate(same_block, self.predictor) + vecs = self._get_objective_vectors(res) + + fronts = self._non_dominated_sort(vecs) + idx1 = same_block.index(p1) + idx2 = same_block.index(p2) + + ranks = {idx: rank for rank, front in enumerate(fronts) for idx in front} + r1 = ranks[idx1] + r2 = ranks[idx2] + + if r1 < r2: + return p1 + if r2 < r1: + return p2 + + front_indices = fronts[r1] + front_vecs = vecs[front_indices] + dists = self._calculate_crowding_distance(front_vecs) + + pos1 = front_indices.index(idx1) + pos2 = front_indices.index(idx2) + if dists[pos1] > dists[pos2]: + return p1 + if dists[pos2] > dists[pos1]: + return p2 # weaker dominance: larger eval set may dominate smaller on the smaller's blocks elif blocks1.issubset(blocks2) and blocks1: diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index b2e378eb..d2257485 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -95,7 +95,7 @@ def __init__( 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[str, set[int]] = {} # prompt_str: set of evaluated block indices + self.prompt_evaluated_blocks: Dict[Prompt, List[int]] = {} # prompt_str: set of evaluated block indices def subsample(self, eval_strategy: Optional["EvalStrategy"] = None) -> Tuple[List[str], List[str]]: """Subsample the dataset based on the specified parameters. @@ -283,9 +283,9 @@ def evaluate( # Record evaluated block for block strategies for prompt in prompts_list: if isinstance(self.block_idx, list): - self.prompt_evaluated_blocks.setdefault(str(prompt), set()).update(self.block_idx) + self.prompt_evaluated_blocks.setdefault(prompt, []).extend(self.block_idx) else: - self.prompt_evaluated_blocks.setdefault(str(prompt), set()).add(self.block_idx) + self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx) input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = self._compute_costs( prompts_list, xs, ys, self.seq_cache, predictor @@ -378,6 +378,6 @@ def set_block_idx(self, idx: Union[int, List[int]]) -> None: self.block_idx = idx - def get_evaluated_blocks(self, prompts: List[Prompt]) -> Dict[str, set[int]]: + def get_evaluated_blocks(self, prompts: List[Prompt]) -> Dict[Prompt, List[int]]: """Return mapping of prompt string to evaluated block indices.""" - return {str(p): set(self.prompt_evaluated_blocks.get(str(p), set())) for p in prompts} + return {p: list(self.prompt_evaluated_blocks.get(p, [])) for p in prompts} \ No newline at end of file From 1a25c51ee6966fe33d22fd9dd271ae03b3ac43ef Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 13 Jan 2026 17:27:10 +0100 Subject: [PATCH 32/53] fix token counting --- promptolution/tasks/base_task.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index d2257485..d8dacf22 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -212,11 +212,11 @@ def _compute_costs( seq_token_counts: List[float] = [] for x, y in zip(xs, ys): cache_key = self._cache_key(prompt, x, str(y)) - seq_text = seq_cache.get(cache_key, "") + seq_text = seq_cache[cache_key] seq_token_counts.append(token_counter(seq_text)) prompt_input_tokens = prompt_tokens + input_token_counts - output_token_counts = np.maximum(np.array(seq_token_counts, dtype=float) - input_token_counts, 0.0) + output_token_counts = np.array(seq_token_counts, dtype=float) - input_token_counts per_prompt_inputs.append(np.asarray(prompt_input_tokens, dtype=float)) per_prompt_outputs.append(output_token_counts) @@ -224,8 +224,8 @@ def _compute_costs( inputs_array = np.vstack(per_prompt_inputs) outputs_array = np.vstack(per_prompt_outputs) - agg_input_tokens = inputs_array.mean(axis=1) if inputs_array.size else np.array([]) - agg_output_tokens = outputs_array.mean(axis=1) if outputs_array.size else np.array([]) + agg_input_tokens = inputs_array.mean(axis=1) + agg_output_tokens = outputs_array.mean(axis=1) return inputs_array, outputs_array, agg_input_tokens, agg_output_tokens From 98214cd9a0911688a09c7432a4b003e617bc0406 Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 13 Jan 2026 17:37:23 +0100 Subject: [PATCH 33/53] change tokenizer --- promptolution/utils/token_counter.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/promptolution/utils/token_counter.py b/promptolution/utils/token_counter.py index 75a6d408..8dfd1e00 100644 --- a/promptolution/utils/token_counter.py +++ b/promptolution/utils/token_counter.py @@ -26,9 +26,5 @@ def get_token_counter(llm: "BaseLLM") -> Callable[[str], float]: A callable that takes a text input and returns the token count. """ - if llm.tokenizer is not None: - tokenizer: "PreTrainedTokenizer" = llm.tokenizer - return lambda x: float(len(tokenizer.encode(x))) - else: - logger.warning("⚠️ The LLM does not have a tokenizer. Using simple token count.") - return lambda x: float(len(x.split())) + + return lambda x: float(len(x.split())) From 21a612d42b82b43cd720ffbf091fd48030e717df Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 13 Jan 2026 17:42:10 +0100 Subject: [PATCH 34/53] revert --- promptolution/utils/token_counter.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/promptolution/utils/token_counter.py b/promptolution/utils/token_counter.py index 8dfd1e00..75a6d408 100644 --- a/promptolution/utils/token_counter.py +++ b/promptolution/utils/token_counter.py @@ -26,5 +26,9 @@ def get_token_counter(llm: "BaseLLM") -> Callable[[str], float]: A callable that takes a text input and returns the token count. """ - - return lambda x: float(len(x.split())) + if llm.tokenizer is not None: + tokenizer: "PreTrainedTokenizer" = llm.tokenizer + return lambda x: float(len(tokenizer.encode(x))) + else: + logger.warning("⚠️ The LLM does not have a tokenizer. Using simple token count.") + return lambda x: float(len(x.split())) From c0d13cc0bacdc4e9816297723903a0224637d3f0 Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 13 Jan 2026 17:53:00 +0100 Subject: [PATCH 35/53] fix token counting --- promptolution/tasks/base_task.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index d8dacf22..8a8d9af1 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -206,14 +206,15 @@ def _compute_costs( per_prompt_outputs: List[np.ndarray] = [] input_token_counts = np.array([token_counter(x) for x in xs], dtype=float) - for prompt in prompts: prompt_tokens = token_counter(prompt.construct_prompt()) seq_token_counts: List[float] = [] for x, y in zip(xs, ys): cache_key = self._cache_key(prompt, x, str(y)) seq_text = seq_cache[cache_key] - seq_token_counts.append(token_counter(seq_text)) + prefix = f"{x}\n" + gen_text = seq_text[len(prefix):] if seq_text.startswith(prefix) else seq_text + seq_token_counts.append(token_counter(gen_text)) prompt_input_tokens = prompt_tokens + input_token_counts output_token_counts = np.array(seq_token_counts, dtype=float) - input_token_counts From a3affe3b96aa5b63be3d5bd5d8e4f14b32929e06 Mon Sep 17 00:00:00 2001 From: finitearth Date: Tue, 13 Jan 2026 18:18:36 +0100 Subject: [PATCH 36/53] revert --- promptolution/tasks/base_task.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 8a8d9af1..0e3f02ac 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -197,7 +197,6 @@ def _compute_costs( prompts: List[Prompt], xs: List[str], ys: List[str], - seq_cache: Dict[Tuple[str, str, str], str], predictor: "BasePredictor", ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: token_counter = get_token_counter(predictor.llm) @@ -206,15 +205,14 @@ def _compute_costs( per_prompt_outputs: List[np.ndarray] = [] input_token_counts = np.array([token_counter(x) for x in xs], dtype=float) + for prompt in prompts: prompt_tokens = token_counter(prompt.construct_prompt()) seq_token_counts: List[float] = [] for x, y in zip(xs, ys): cache_key = self._cache_key(prompt, x, str(y)) - seq_text = seq_cache[cache_key] - prefix = f"{x}\n" - gen_text = seq_text[len(prefix):] if seq_text.startswith(prefix) else seq_text - seq_token_counts.append(token_counter(gen_text)) + seq_text = self.seq_cache[cache_key] + seq_token_counts.append(token_counter(seq_text)) prompt_input_tokens = prompt_tokens + input_token_counts output_token_counts = np.array(seq_token_counts, dtype=float) - input_token_counts @@ -289,7 +287,7 @@ def evaluate( self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx) input_tokens, output_tokens, agg_input_tokens, agg_output_tokens = self._compute_costs( - prompts_list, xs, ys, self.seq_cache, predictor + prompts_list, xs, ys, predictor ) return EvalResult( From a4e2476db4f7265e86599667f7f23d3556c533bd Mon Sep 17 00:00:00 2001 From: finitearth Date: Wed, 14 Jan 2026 16:33:58 +0100 Subject: [PATCH 37/53] bye capoeira --- docs/examples/getting_started.md | 2 +- docs/examples/reward_task_tutorial.md | 2 +- promptolution/helpers.py | 9 - promptolution/optimizers/__init__.py | 2 - promptolution/optimizers/base_optimizer.py | 6 +- promptolution/optimizers/capoeira.py | 520 ---------------- promptolution/tasks/multi_objective_task.py | 13 +- promptolution/tasks/reward_tasks.py | 3 +- promptolution/utils/capo_utils.py | 2 +- promptolution/utils/prompt.py | 4 +- tests/helpers/test_helpers.py | 5 - tests/optimizers/test_capoeira.py | 625 -------------------- tests/tasks/test_base_task.py | 6 +- tests/tasks/test_multi_objective_task.py | 4 +- tests/tasks/test_reward_tasks.py | 1 + tutorials/getting_started.ipynb | 2 +- tutorials/reward_task_tutorial.ipynb | 2 +- 17 files changed, 22 insertions(+), 1186 deletions(-) delete mode 100644 promptolution/optimizers/capoeira.py delete mode 100644 tests/optimizers/test_capoeira.py diff --git a/docs/examples/getting_started.md b/docs/examples/getting_started.md index 2dfc1f14..47f359d5 100644 --- a/docs/examples/getting_started.md +++ b/docs/examples/getting_started.md @@ -83,7 +83,7 @@ api_key = "YOUR_API_KEY" # Replace with your Promptolution API key ``` Here's an explanation of each configuration parameter in the ExperimentConfig: -- `optimizer`: The algorithm used for prompt optimization. Currently we support "capo", "capoeira", "evopromptga", "evopromptde", and "opro". For this example, we use "capo" as it is capable of leveraging few-shot examples. +- `optimizer`: The algorithm used for prompt optimization. Currently we support "capo", "evopromptga", "evopromptde", and "opro". For this example, we use "capo" as it is capable of leveraging few-shot examples. - `task_description`: A string describing the task you're optimizing prompts for. This is used to provide the meta-llm with context about your task. - `prompts`: A list of initial prompt strings that will be used as the starting point for optimization. - `n_steps`: The number of optimization steps to run. Higher values allow more exploration and refinement but require more API calls and computational resources. diff --git a/docs/examples/reward_task_tutorial.md b/docs/examples/reward_task_tutorial.md index da51cfdd..82d0e973 100644 --- a/docs/examples/reward_task_tutorial.md +++ b/docs/examples/reward_task_tutorial.md @@ -102,7 +102,7 @@ api_key = "YOUR_API_KEY" # Replace with your Promptolution API key ``` Here's an explanation of each configuration parameter in the ExperimentConfig: -- `optimizer`: The algorithm used for prompt optimization. Currently we support "capo", "capoeira", "evopromptga", "evopromptde", and "opro". For this example, we use "capo" as it is capable of leveraging few-shot examples. +- `optimizer`: The algorithm used for prompt optimization. Currently we support "capo", "evopromptga", "evopromptde", and "opro". For this example, we use "capo" as it is capable of leveraging few-shot examples. - `task_description`: A string describing the task you're optimizing prompts for. This is used to provide the meta-llm with context about your task. - `prompts`: A list of initial prompt strings that will be used as the starting point for optimization. - `n_steps`: The number of optimization steps to run. Higher values allow more exploration and refinement but require more API calls and computational resources. diff --git a/promptolution/helpers.py b/promptolution/helpers.py index 4a249f7d..9cf13d6e 100644 --- a/promptolution/helpers.py +++ b/promptolution/helpers.py @@ -27,7 +27,6 @@ from promptolution.llms.local_llm import LocalLLM from promptolution.llms.vllm import VLLM from promptolution.optimizers.capo import CAPO -from promptolution.optimizers.capoeira import Capoeira from promptolution.optimizers.evoprompt_de import EvoPromptDE from promptolution.optimizers.evoprompt_ga import EvoPromptGA from promptolution.optimizers.opro import OPRO @@ -232,14 +231,6 @@ def get_optimizer( config=config, ) - if final_optimizer == "capoeira": - return Capoeira( - predictor=predictor, - meta_llm=meta_llm, - task=task, - config=config, - ) - if final_optimizer == "evopromptde": return EvoPromptDE(predictor=predictor, meta_llm=meta_llm, task=task, config=config) diff --git a/promptolution/optimizers/__init__.py b/promptolution/optimizers/__init__.py index 96e9a484..4b7a7dbb 100644 --- a/promptolution/optimizers/__init__.py +++ b/promptolution/optimizers/__init__.py @@ -1,14 +1,12 @@ """Module for prompt optimizers.""" from promptolution.optimizers.capo import CAPO -from promptolution.optimizers.capoeira import Capoeira from promptolution.optimizers.evoprompt_de import EvoPromptDE from promptolution.optimizers.evoprompt_ga import EvoPromptGA from promptolution.optimizers.opro import OPRO __all__ = [ "CAPO", - "Capoeira", "EvoPromptDE", "EvoPromptGA", "OPRO", diff --git a/promptolution/optimizers/base_optimizer.py b/promptolution/optimizers/base_optimizer.py index 79f668c8..69d163f2 100644 --- a/promptolution/optimizers/base_optimizer.py +++ b/promptolution/optimizers/base_optimizer.py @@ -15,7 +15,7 @@ logger = get_logger(__name__) -OptimizerType = Literal["evopromptde", "evopromptga", "opro", "capo", "capoeira"] +OptimizerType = Literal["evopromptde", "evopromptga", "opro", "capo"] class BaseOptimizer(ABC): @@ -54,10 +54,10 @@ def __init__( self.prompts: List[Prompt] = [Prompt(p) for p in initial_prompts] if initial_prompts else [] if task.task_type == "multi" and not self.supports_multi_objective: logger.warning( - f"{self.__class__.__name__} does not support multi-objective tasks; objectives will be averaged equally.", + f"{self.__class__.__name__} does not support multi-objective tasks, objectives will be averaged equally.", ) task.activate_scalarized_objective() - + self.task = task self.callbacks: List["BaseCallback"] = callbacks or [] self.predictor = predictor diff --git a/promptolution/optimizers/capoeira.py b/promptolution/optimizers/capoeira.py deleted file mode 100644 index 5b913f16..00000000 --- a/promptolution/optimizers/capoeira.py +++ /dev/null @@ -1,520 +0,0 @@ -"""Implementation of the Capoeira (Multi-Objective CAPO) optimizer.""" - -import random - -import numpy as np -import pandas as pd - -from typing import TYPE_CHECKING, List, Optional, Tuple - -if TYPE_CHECKING: # pragma: no cover - from promptolution.utils.callbacks import BaseCallback - from promptolution.llms.base_llm import BaseLLM - from promptolution.predictors.base_predictor import BasePredictor - from promptolution.tasks.base_task import BaseTask - from promptolution.utils.config import ExperimentConfig - -from promptolution.optimizers.base_optimizer import BaseOptimizer -from promptolution.tasks.multi_objective_task import MultiObjectiveTask -from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation -from promptolution.utils.logging import get_logger -from promptolution.utils.prompt import Prompt -from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_MUTATION_TEMPLATE -from promptolution.utils.token_counter import get_token_counter - -logger = get_logger(__name__) - - -class Capoeira(BaseOptimizer): - """Multi-objective variant of CAPO with Pareto-based selection.""" - - supports_multi_objective = True - - def __init__( - self, - predictor: "BasePredictor", - task: "BaseTask", - meta_llm: "BaseLLM", - initial_prompts: Optional[List[str]] = None, - crossover_template: Optional[str] = None, - mutation_template: Optional[str] = None, - crossovers_per_iter: int = 4, - upper_shots: int = 5, - cost_per_input_token: float = 1.0, - cost_per_output_token: float = 1.0, - check_fs_accuracy: bool = True, - create_fs_reasoning: bool = True, - df_few_shots: Optional[pd.DataFrame] = None, - callbacks: Optional[List["BaseCallback"]] = None, - config: Optional["ExperimentConfig"] = None, - ) -> None: - """Initialize the Capoeira optimizer. - - Args: - predictor: The predictor used to evaluate prompt performance. - task: The task instance containing data and evaluation settings. - meta_llm: Meta language model for crossover and mutation generation. - initial_prompts: Starting prompt strings to seed the population. - crossover_template: Optional meta-prompt template for crossover. - mutation_template: Optional meta-prompt template for mutation. - crossovers_per_iter: Number of crossover operations per iteration. - upper_shots: Maximum number of few-shot examples to attach. - cost_per_input_token: Weight applied to input token cost for the cost objective. - cost_per_output_token: Weight applied to output token cost for the cost objective. - check_fs_accuracy: Whether to verify few-shot correctness before use. - create_fs_reasoning: Whether to replace few-shots with model reasoning. - df_few_shots: Optional dataframe providing few-shot examples. If None, will pop 10% of datapoints from task. - callbacks: Optional list of optimization callbacks. - config: Optional experiment configuration object. - """ - self.meta_llm = meta_llm - self.downstream_llm = predictor.llm - self.crossovers_per_iter = crossovers_per_iter - self.upper_shots = upper_shots - - self.cost_per_input_token = cost_per_input_token - self.cost_per_output_token = cost_per_output_token - self.check_fs_accuracy = check_fs_accuracy - self.create_fs_reasoning = create_fs_reasoning - - super().__init__(predictor, task, initial_prompts, callbacks, config) - - self.crossover_template = self._initialize_meta_template(crossover_template or CAPO_CROSSOVER_TEMPLATE) - self.mutation_template = self._initialize_meta_template(mutation_template or CAPO_MUTATION_TEMPLATE) - self.token_counter = get_token_counter(self.downstream_llm) - self.df_few_shots = df_few_shots if df_few_shots is not None else task.pop_datapoints(frac=0.1) - - self.incumbents: List[Prompt] = self.prompts - self.non_incumbents: List[Prompt] = [] - self.population_size = len(self.prompts) - - if self.task.task_type == "multi": - self.n_objectives = len(self.task.tasks) + 1 # +1 for cost objective # type: ignore - else: - self.n_objectives = 2 # single objective + cost objective - - if "block" not in self.task.eval_strategy: - logger.warning( - f"ℹ️ CAPO requires 'block' in the eval_strategy, but got {self.task.eval_strategy}. Setting eval_strategy to 'sequential_block'." - ) - self.task.eval_strategy = "sequential_block" - - if hasattr(self.predictor, "begin_marker") and hasattr(self.predictor, "end_marker"): - self.target_begin_marker = self.predictor.begin_marker # type: ignore - self.target_end_marker = self.predictor.end_marker # type: ignore - else: - self.target_begin_marker = "" - self.target_end_marker = "" - - def _pre_optimization_loop(self) -> None: - population: List[Prompt] = [] - for prompt in self.prompts: - num_examples = random.randint(0, self.upper_shots) - few_shots = build_few_shot_examples( - instruction=prompt.instruction, - num_examples=num_examples, - optimizer=self, - ) - population.append(Prompt(prompt.instruction, few_shots)) - - init_result = self.task.evaluate(population, self.predictor) - initial_vectors = self._get_objective_vectors(init_result) - fronts = self._non_dominated_sort(initial_vectors) - self.incumbents = [population[i] for i in fronts[0]] - self.non_incumbents = [population[i] for front in fronts[1:] for i in front] - - # keep self.prompts as a "view" if base class expects it - self.scores = initial_vectors[:, 0].tolist() - - def _step(self) -> List[Prompt]: - # 1) generate challengers - offsprings = perform_crossover(self.prompts, self, self._tournament_selection) - new_challengers = perform_mutation(offsprings, self) - - # 2) intensify each challenger; after each, advance incumbents + prune - for challenger in new_challengers: - self._do_intensification(challenger) - self._select_survivors() - self._advance_one_incumbent() - - inc_result = self.task.evaluate( - prompts=self.incumbents, predictor=self.predictor, eval_strategy="evaluated" - ) - vecs_inc = self._get_objective_vectors(inc_result) - self.scores = vecs_inc[:, 0].tolist() - self.prompts = self.incumbents - - return self.prompts - - def _do_intensification(self, challenger: Prompt) -> None: - if challenger in self.incumbents: - return - if challenger in self.non_incumbents: - # remove from non-incumbents to re-evaluate - self.non_incumbents.remove(challenger) - - common_blocks = self._get_common_blocks(self.incumbents) - - # bootstrap if no common blocks yet - if not common_blocks: - b = random.randrange(self.task.n_blocks) - self.task.set_block_idx(b) - self.task.evaluate(self.incumbents + [challenger], self.predictor) - self.incumbents.append(challenger) - self._update_incumbent_front(blocks={b}) - return - - remaining_blocks = set(common_blocks) - - challenger_mean: Optional[np.ndarray] = None - incumbents_mean: Optional[np.ndarray] = None - t = 0 - - old_scores = np.full((self.n_objectives,), np.inf) - - while remaining_blocks: - b = random.choice(tuple(remaining_blocks)) - remaining_blocks.remove(b) - - # evaluate all incumbents + challenger on THIS block (cache will avoid recompute) - self.task.set_block_idx(b) - res = self.task.evaluate(self.incumbents + [challenger], self.predictor) - vecs = self._get_objective_vectors(res) # per-block vectors, shape (n_inc+1, n_obj) - incumbent_block = vecs[:-1] - challenger_block = vecs[-1] - - # running means - t += 1 - if challenger_mean is None: - challenger_mean = challenger_block.copy() - incumbents_mean = incumbent_block.copy() - else: - challenger_mean += (challenger_block - challenger_mean) / t - incumbents_mean += (incumbent_block - incumbents_mean) / t # type: ignore - - if self._is_dominated(old_scores, challenger_mean): - continue - - old_scores = challenger_mean.copy() # type: ignore - - closest_incumbent = self._get_closest_incumbent(challenger) # type: ignore - if self._is_dominated(challenger_mean, closest_incumbent): - self.non_incumbents.append(challenger) - return - - # survived all common blocks -> admit and update front restricted to common_blocks - self.incumbents.append(challenger) - self._update_incumbent_front(blocks=common_blocks) - - def _get_closest_incumbent(self, challenger) -> np.ndarray: - """Return the vector of the geometrically closest incumbent.""" - challenger_res = self.task.evaluate(challenger, self.predictor, eval_strategy="evaluated") - challenger_vec = self._get_objective_vectors(challenger_res) - - incumbent_res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") - incumbent_vecs = self._get_objective_vectors(incumbent_res) - - all_vecs = np.vstack([incumbent_vecs, challenger_vec]) - min_b = np.min(all_vecs, axis=0) - max_b = np.max(all_vecs, axis=0) - rng = max_b - min_b - rng[rng == 0] = 1.0 # Avoid div/0 - - challenger_norm = (challenger_vec - min_b) / rng - incumbents_norm = (incumbent_vecs - min_b) / rng - - dists = np.linalg.norm(incumbents_norm - challenger_norm, axis=1) - idx = int(np.argmin(dists)) - return incumbent_vecs[idx] - - def _update_incumbent_front(self, blocks: Optional[set[int]] = None) -> None: - if blocks is None: - res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") - else: - self.task.set_block_idx(list(sorted(blocks))) # sorted for deterministic behaviour - res = self.task.evaluate(self.incumbents, self.predictor) - - vecs = self._get_objective_vectors(res) - - fronts = self._non_dominated_sort(vecs) - - new_incumbents = [self.incumbents[i] for i in fronts[0]] - demoted = [self.incumbents[i] for front in fronts[1:] for i in front] - - self.incumbents = new_incumbents - self.non_incumbents.extend(demoted) - - def _get_objective_vectors(self, result) -> np.ndarray: - # If the task is multi-objective, include all objective dimensions, else single objective. - if isinstance(self.task, MultiObjectiveTask): - agg_scores = np.stack(result.agg_scores, axis=1) # shape: (n_prompts, n_objectives) - else: - agg_scores = np.atleast_2d(result.agg_scores).T # shape: (n_prompts, 1) - - agg_input_tokens = np.asarray(result.agg_input_tokens) - agg_output_tokens = np.asarray(result.agg_output_tokens) - cost_scalar = self.cost_per_input_token * agg_input_tokens + self.cost_per_output_token * agg_output_tokens - cost_scalar = cost_scalar.reshape(-1, 1) - - return np.hstack([agg_scores, -cost_scalar]) - - def _advance_one_incumbent(self) -> None: - if not self.incumbents: - return - - blocks_map = self.task.get_evaluated_blocks(self.incumbents) - inc_keys = [str(inc) for inc in self.incumbents] - - # least evaluated incumbents - eval_counts = [len(blocks_map[k]) for k in inc_keys] - min_count = min(eval_counts) - least = [inc for inc, c in zip(self.incumbents, eval_counts) if c == min_count] - chosen_inc = random.choice(least) - - # union over incumbents - union_blocks: set[int] = set() - for inc in self.incumbents: - union_blocks |= set(blocks_map[inc]) - - chosen_blocks = set(blocks_map[chosen_inc]) - - # gap-first, else brand-new - gap_blocks = union_blocks - chosen_blocks - if gap_blocks: - b = random.choice(tuple(gap_blocks)) - else: - all_blocks = set(range(self.task.n_blocks)) - new_blocks = all_blocks - union_blocks - if not new_blocks: - return - b = random.choice(tuple(new_blocks)) - - self.task.set_block_idx(b) - self.task.evaluate(prompts=[chosen_inc], predictor=self.predictor) - - def _select_survivors(self) -> None: - """Prune population via Pareto logic to enforce size constraints.""" - while len(self.incumbents) + len(self.non_incumbents) > self.population_size: - if len(self.non_incumbents) > 0: - # 1. Check Heterogeneity (Fairness Check) - blocks_map = self.task.get_evaluated_blocks(self.non_incumbents) - block_sets = list(blocks_map.values()) - - first_set = block_sets[0] - # Are all challengers evaluated on the exact same set of blocks? - is_homogeneous = all(s == first_set for s in block_sets) - - if not is_homogeneous: - # CASE A: Heterogeneous (Unfair comparison). - # Prune the prompt with the FEWEST evaluations (least reliable/least invested). - counts = [len(s) for s in block_sets] - min_count = min(counts) - - # Find all indices with the minimum count (handle ties randomly) - candidates = [i for i, c in enumerate(counts) if c == min_count] - victim_idx = random.choice(candidates) - - self.non_incumbents.pop(victim_idx) - continue - - # CASE B: Homogeneous (Fair comparison). - # Use NDS + Crowding Distance. - - # Get objective vectors for all challengers (safe because blocks are identical) - res = self.task.evaluate(self.non_incumbents, self.predictor, eval_strategy="evaluated") - vecs = self._get_objective_vectors(res) - - # Perform Non-Dominated Sort - fronts = self._non_dominated_sort(vecs) - - # Select the worst front (the last one) - worst_front_indices = fronts[-1] - - # Multiple candidates in worst front -> Prune by Crowding Distance - # We want to keep diversity (high CD), so we remove low CD. - worst_front_vecs = vecs[worst_front_indices] - dists = self._calculate_crowding_distance(worst_front_vecs) - - # Find index relative to the worst front list - min_dist = np.min(dists) - tied_indices = np.where(dists == min_dist)[0] - local_worst_idx = np.random.choice(tied_indices) - victim_idx = worst_front_indices[local_worst_idx] - - self.non_incumbents.pop(victim_idx) - continue - - # --- PRUNE FROM INCUMBENTS --- - # Fallback: If we only have incumbents, remove the least unique one. - res = self.task.evaluate(self.incumbents, self.predictor, eval_strategy="evaluated") - vecs = self._get_objective_vectors(res) - dists = self._calculate_crowding_distance(vecs) - - # Remove the one with the smallest crowding distance - victim_idx = int(np.argmin(dists)) - self.incumbents.pop(victim_idx) - - def _get_common_blocks(self, prompts: List[Prompt]) -> set[int]: - """Get the set of block indices that have been evaluated by all given prompts.""" - per_prompt = self.task.get_evaluated_blocks(prompts) # Dict[prompt -> Collection[int]] - block_sets = list(per_prompt.values()) - - if not block_sets: - return set() - - # Some task implementations may return lists instead of sets; normalize for typing and correctness. - common = set.intersection(*(set(s) for s in block_sets)) - return common - - def _select_parent_from_pool(self, selection_pool: List[Prompt]) -> Prompt: - """Tournament-pick a parent, preferring incumbents and using crowding for ties.""" - p1, p2 = random.sample(selection_pool, 2) - - if p1 in self.incumbents and p2 in self.incumbents: - return self._pick_incumbent_by_crowding(p1, p2) - if p1 in self.incumbents: - return p1 - if p2 in self.incumbents: - return p2 - - # both are non-incumbents - blocks_map = self.task.get_evaluated_blocks([p1, p2]) - blocks1 = blocks_map[p1] - blocks2 = blocks_map[p2] - - if blocks1 == blocks2: # both evaluated on same blocks - # Use full NDS + crowding on all non-incumbents that share this block set - blocks_map = self.task.get_evaluated_blocks(self.non_incumbents) - same_block = [p for p in self.non_incumbents if blocks_map[p] == blocks1] - - if len(same_block) >= 2: - self.task.set_block_idx(list(sorted(blocks1))) - res = self.task.evaluate(same_block, self.predictor) - vecs = self._get_objective_vectors(res) - - fronts = self._non_dominated_sort(vecs) - idx1 = same_block.index(p1) - idx2 = same_block.index(p2) - - ranks = {idx: rank for rank, front in enumerate(fronts) for idx in front} - r1 = ranks[idx1] - r2 = ranks[idx2] - - if r1 < r2: - return p1 - if r2 < r1: - return p2 - - front_indices = fronts[r1] - front_vecs = vecs[front_indices] - dists = self._calculate_crowding_distance(front_vecs) - - pos1 = front_indices.index(idx1) - pos2 = front_indices.index(idx2) - if dists[pos1] > dists[pos2]: - return p1 - if dists[pos2] > dists[pos1]: - return p2 - - # weaker dominance: larger eval set may dominate smaller on the smaller's blocks - elif blocks1.issubset(blocks2) and blocks1: - self.task.set_block_idx(list(sorted(blocks1))) - res = self.task.evaluate([p1, p2], self.predictor) - vecs = self._get_objective_vectors(res) - if self._is_dominated(vecs[0], vecs[1]): - return p2 - elif blocks2.issubset(blocks1) and blocks2: - self.task.set_block_idx(list(sorted(blocks2))) - res = self.task.evaluate([p1, p2], self.predictor) - vecs = self._get_objective_vectors(res) - if self._is_dominated(vecs[1], vecs[0]): - return p1 - - return random.choice((p1, p2)) - - - def _pick_incumbent_by_crowding(self, p1: Prompt, p2: Prompt) -> Prompt: - """Break incumbent ties using crowding distance over common evaluated blocks.""" - common_blocks = self._get_common_blocks(self.incumbents) - if common_blocks: - self.task.set_block_idx(common_blocks) - res = self.task.evaluate(self.incumbents, self.predictor) - inc_vectors = self._get_objective_vectors(res) - inc_distances = self._calculate_crowding_distance(inc_vectors) - - p1_idx = self.incumbents.index(p1) - p2_idx = self.incumbents.index(p2) - if inc_distances[p1_idx] > inc_distances[p2_idx]: - return p1 - if inc_distances[p2_idx] > inc_distances[p1_idx]: - return p2 - return random.choice((p1, p2)) - - - def _tournament_selection(self) -> Tuple[Prompt, Prompt]: - """Pick two distinct parents via tournament selection.""" - selection_pool = self.incumbents + self.non_incumbents - parent1 = self._select_parent_from_pool(selection_pool) - - parent2 = self._select_parent_from_pool(selection_pool) - while parent2 == parent1: - parent2 = self._select_parent_from_pool(selection_pool) - - return parent1, parent2 - - @staticmethod - def _non_dominated_sort(obj_vectors: np.ndarray) -> List[List[int]]: - """Perform fast non-dominated sorting (NSGA-II) in a vectorized manner.""" - n_solutions = obj_vectors.shape[0] - - greater = obj_vectors[:, None, :] > obj_vectors[None, :, :] - greater_equal = obj_vectors[:, None, :] >= obj_vectors[None, :, :] - dominates = np.all(greater_equal, axis=2) & np.any(greater, axis=2) - - domination_counts = dominates.sum(axis=0) - dominated_solutions = [list(np.where(dominates[i])[0]) for i in range(n_solutions)] - - fronts: List[List[int]] = [list(np.where(domination_counts == 0)[0])] - - current_front = 0 - while current_front < len(fronts) and len(fronts[current_front]) > 0: - next_front: List[int] = [] - for i in fronts[current_front]: - for dominated in dominated_solutions[i]: - domination_counts[dominated] -= 1 - if domination_counts[dominated] == 0: - next_front.append(dominated) - if len(next_front) > 0: - fronts.append(next_front) - current_front += 1 - - return fronts - - @staticmethod - def _is_dominated(vec1, vec2): - """Returns True if vec2 dominates vec1 in a maximize-all setting.""" - return np.all(vec2 >= vec1) and np.any(vec2 > vec1) - - @staticmethod - def _calculate_crowding_distance(obj_vectors: np.ndarray) -> np.ndarray: - """Calculate crowding distance for a set of solutions.""" - num_solutions, num_obj = obj_vectors.shape - if num_solutions <= 2: - return np.full(num_solutions, float("inf")) - - distances = np.zeros(num_solutions) - for i in range(num_obj): - sorted_indices = np.argsort(obj_vectors[:, i]) - distances[sorted_indices[0]] = float("inf") - distances[sorted_indices[-1]] = float("inf") - - f_min = obj_vectors[sorted_indices[0], i] - f_max = obj_vectors[sorted_indices[-1], i] - if f_max == f_min: - continue - - slice_indices = sorted_indices[1:-1] - next_vals = obj_vectors[sorted_indices[2:], i] - prev_vals = obj_vectors[sorted_indices[:-2], i] - distances[slice_indices] += (next_vals - prev_vals) / (f_max - f_min) - return distances diff --git a/promptolution/tasks/multi_objective_task.py b/promptolution/tasks/multi_objective_task.py index 1d557535..3e814712 100644 --- a/promptolution/tasks/multi_objective_task.py +++ b/promptolution/tasks/multi_objective_task.py @@ -8,7 +8,7 @@ from typing import Dict, List, Optional, Tuple -from promptolution.tasks.base_task import BaseTask, EvalResult, EvalStrategy, TaskType +from promptolution.tasks.base_task import BaseTask, EvalResult, EvalStrategy from promptolution.utils.prompt import Prompt @@ -127,17 +127,16 @@ def evaluate( # type: ignore 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, task.seq_cache, predictor + prompts_list, xs, ys, predictor ) # Record evaluated block for block strategies for prompt in prompts_list: - block_set = task.prompt_evaluated_blocks.setdefault(str(prompt), set()) + block_set = task.prompt_evaluated_blocks.setdefault(prompt, []) if isinstance(task.block_idx, list): - block_set.update(task.block_idx) + block_set.extend(task.block_idx) else: - block_set.add(task.block_idx) - + block_set.append(task.block_idx) per_task_results.append( EvalResult( scores=scores_array, @@ -156,7 +155,7 @@ def evaluate( # type: ignore # Mirror evaluated block bookkeeping using the first task for parity with BaseTask. first_task = self.tasks[0] first_result = per_task_results[0] - self.prompt_evaluated_blocks = {str(p): first_task.prompt_evaluated_blocks[str(p)] for p in prompts_list} + self.prompt_evaluated_blocks = {p: first_task.prompt_evaluated_blocks[p] for p in prompts_list} if self._scalarized_objective: return EvalResult( diff --git a/promptolution/tasks/reward_tasks.py b/promptolution/tasks/reward_tasks.py index 2da558d4..7bb79277 100644 --- a/promptolution/tasks/reward_tasks.py +++ b/promptolution/tasks/reward_tasks.py @@ -1,9 +1,10 @@ """Module for Reward tasks.""" +from collections import defaultdict + import numpy as np import pandas as pd -from collections import defaultdict from typing import TYPE_CHECKING, Callable, List, Optional diff --git a/promptolution/utils/capo_utils.py b/promptolution/utils/capo_utils.py index 9bda0c48..40ee707c 100644 --- a/promptolution/utils/capo_utils.py +++ b/promptolution/utils/capo_utils.py @@ -4,7 +4,7 @@ import random -from typing import List, Optional, Callable +from typing import Callable, List, Optional from promptolution.utils.formatting import extract_from_tag from promptolution.utils.prompt import Prompt diff --git a/promptolution/utils/prompt.py b/promptolution/utils/prompt.py index bc76de57..641e0dc7 100644 --- a/promptolution/utils/prompt.py +++ b/promptolution/utils/prompt.py @@ -59,10 +59,8 @@ def __eq__(self, other: object) -> bool: ) def __hash__(self) -> int: + """Hash function for use in sets and dict keys.""" return hash((self.instruction, tuple(self.few_shots), self.downstream_template)) - - - def sort_prompts_by_scores( diff --git a/tests/helpers/test_helpers.py b/tests/helpers/test_helpers.py index c77b2748..d90e1ea3 100644 --- a/tests/helpers/test_helpers.py +++ b/tests/helpers/test_helpers.py @@ -21,7 +21,6 @@ run_optimization, ) from promptolution.optimizers.capo import CAPO -from promptolution.optimizers.capoeira import Capoeira from promptolution.optimizers.evoprompt_de import EvoPromptDE from promptolution.optimizers.evoprompt_ga import EvoPromptGA from promptolution.optimizers.opro import OPRO @@ -407,10 +406,6 @@ def test_get_optimizer_variants(): assert isinstance(opt, CAPO) - opt2 = get_optimizer(pred, MockLLM(), task, optimizer="capoeira", config=cfg) - - assert isinstance(opt2, Capoeira) - opt3 = get_optimizer(pred, MockLLM(), task, optimizer="evopromptde", config=cfg) assert isinstance(opt3, EvoPromptDE) diff --git a/tests/optimizers/test_capoeira.py b/tests/optimizers/test_capoeira.py deleted file mode 100644 index 09955116..00000000 --- a/tests/optimizers/test_capoeira.py +++ /dev/null @@ -1,625 +0,0 @@ -from unittest.mock import patch - -import numpy as np -import pandas as pd - -from tests.mocks.mock_task import MockTask - -from promptolution.optimizers.capoeira import Capoeira -from promptolution.tasks.base_task import EvalResult -from promptolution.tasks.multi_objective_task import MultiObjectiveEvalResult, MultiObjectiveTask -from promptolution.utils.capo_utils import perform_crossover, perform_mutation -from promptolution.utils.prompt import Prompt -from promptolution.utils.templates import CAPO_CROSSOVER_TEMPLATE, CAPO_MUTATION_TEMPLATE - - -def test_capoeira_initialization(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): - optimizer = Capoeira( - predictor=mock_predictor, - task=mock_task, - meta_llm=mock_meta_llm, - initial_prompts=initial_prompts, - df_few_shots=mock_df, - ) - - assert optimizer.crossovers_per_iter == 4 - assert optimizer.population_size == len(initial_prompts) - assert isinstance(optimizer.df_few_shots, pd.DataFrame) - - -def test_capoeira_initialize_population(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): - optimizer = Capoeira( - predictor=mock_predictor, - task=mock_task, - meta_llm=mock_meta_llm, - initial_prompts=initial_prompts, - df_few_shots=mock_df, - ) - - with patch("random.randint", return_value=1): - optimizer._pre_optimization_loop() - population = optimizer.prompts - - assert len(population) == len(initial_prompts) - assert all(isinstance(p, Prompt) for p in population) - - -def test_capoeira_objective_vectors_and_sort(mock_meta_llm, mock_predictor, mock_task, mock_df): - optimizer = Capoeira( - predictor=mock_predictor, - task=mock_task, - meta_llm=mock_meta_llm, - initial_prompts=["short"], - df_few_shots=mock_df, - ) - - result = EvalResult( - scores=np.array([[0.4], [0.9]], dtype=float), - agg_scores=np.array([0.4, 0.9], dtype=float), - sequences=np.array([["s1"], ["s2"]], dtype=object), - input_tokens=np.array([[1.0], [1.0]], dtype=float), - output_tokens=np.array([[0.0], [0.0]], dtype=float), - agg_input_tokens=np.array([10.0, 8.0], dtype=float), - agg_output_tokens=np.array([0.0, 0.0], dtype=float), - ) - - vecs = optimizer._get_objective_vectors(result) - - assert vecs.shape == (2, 2) - assert np.allclose(vecs[:, 0], np.array([0.4, 0.9])) - assert np.allclose(vecs[:, 1], -np.array([10.0, 8.0])) - - fronts = optimizer._non_dominated_sort(vecs) - - assert fronts[0] == [1] - assert 0 in fronts[1] - - -def test_capoeira_meta_prompts(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): - optimizer = Capoeira( - predictor=mock_predictor, - task=mock_task, - meta_llm=mock_meta_llm, - initial_prompts=initial_prompts, - df_few_shots=mock_df, - crossovers_per_iter=2, - ) - - mother = Prompt("Instruction 1", ["Example 1"]) - father = Prompt("Instruction 2", ["Example 2"]) - perform_crossover([mother, father], optimizer=optimizer) - - full_task_desc = mock_task.task_description + "\n" + optimizer.predictor.extraction_description - expected_crossover = ( - CAPO_CROSSOVER_TEMPLATE.replace("", mother.instruction) - .replace("", father.instruction) - .replace("", full_task_desc) - ) - assert expected_crossover in mock_meta_llm.call_history[0]["prompts"] - - mock_meta_llm.reset() - parent = Prompt("Instruction 3", ["Example 3"]) - perform_mutation( - offsprings=[parent], - optimizer=optimizer, - ) - expected_mutation = CAPO_MUTATION_TEMPLATE.replace("", parent.instruction).replace( - "", full_task_desc - ) - assert expected_mutation in mock_meta_llm.call_history[0]["prompts"] - - -def test_capoeira_crowding_distance_edges(): - vecs = np.array([[1.0, 2.0], [3.0, 4.0]]) - dists = Capoeira._calculate_crowding_distance(vecs) - assert np.isinf(dists).all() - - -def test_capoeira_select_survivors_handles_heterogeneous_blocks(mock_meta_llm, mock_predictor): - def fake_evaluate(*_, **__): - return EvalResult( - scores=np.array([[0.5]], dtype=float), - agg_scores=np.array([0.5], dtype=float), - sequences=np.array([[""]], dtype=object), - input_tokens=np.array([[0.0]], dtype=float), - output_tokens=np.array([[0.0]], dtype=float), - agg_input_tokens=np.array([0.0], dtype=float), - agg_output_tokens=np.array([0.0], dtype=float), - ) - - task = MockTask( - eval_strategy="sequential_block", - n_blocks=2, - block_idx=0, - eval_blocks={}, - evaluate_fn=fake_evaluate, - ) - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["inc1", "inc2"], - df_few_shots=task.pop_datapoints(n=1), - ) - - c1, c2 = Prompt("c1"), Prompt("c2") - task.eval_blocks = {str(c1): {0}, str(c2): {0, 1}} - optimizer.incumbents = [Prompt("i1"), Prompt("i2")] - optimizer.non_incumbents = [c1, c2] - optimizer.population_size = 3 - - optimizer._select_survivors() - - assert len(optimizer.non_incumbents) == 1 - assert optimizer.non_incumbents[0].instruction == "c2" - - -def test_capoeira_select_survivors_homogeneous_prunes_lowest(mock_meta_llm, mock_predictor): - next_result: dict[str, EvalResult | None] = {"value": None} - - def fake_evaluate(prompts, *_, **__): - return next_result["value"] # type: ignore[return-value] - - task = MockTask( - eval_strategy="sequential_block", - n_blocks=2, - block_idx=0, - eval_blocks={}, - evaluate_fn=fake_evaluate, - ) - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["inc"], - df_few_shots=task.pop_datapoints(n=1), - ) - - c1, c2 = Prompt("c1"), Prompt("c2") - task.eval_blocks = {str(c1): {0}, str(c2): {0}} - - next_result["value"] = EvalResult( - scores=np.array([[0.1], [0.2]], dtype=float), - agg_scores=np.array([0.1, 0.2], dtype=float), - sequences=np.array([["s1"], ["s2"]], dtype=object), - input_tokens=np.array([[0.0], [0.0]], dtype=float), - output_tokens=np.array([[0.0], [0.0]], dtype=float), - agg_input_tokens=np.array([0.0, 0.0], dtype=float), - agg_output_tokens=np.array([0.0, 0.0], dtype=float), - ) - - optimizer.incumbents = [Prompt("inc")] # keeps population pressure - optimizer.non_incumbents = [c1, c2] - optimizer.population_size = 2 - - optimizer._select_survivors() - - assert len(optimizer.non_incumbents) == 1 - assert optimizer.non_incumbents[0].instruction == "c2" - - -def test_capoeira_select_survivors_prefers_lower_cost(mock_meta_llm, mock_predictor): - def fake_evaluate(prompts, *_, **__): - costs = np.array([1.0 if "cheap" in p.instruction else 5.0 for p in prompts], dtype=float) - return EvalResult( - scores=np.array([[0.4], [0.4]], dtype=float), - agg_scores=np.array([0.4, 0.4], dtype=float), - sequences=np.array([["s1"], ["s2"]], dtype=object), - input_tokens=costs.reshape(-1, 1), - output_tokens=np.zeros((len(prompts), 1)), - agg_input_tokens=costs, - agg_output_tokens=np.zeros(len(prompts)), - ) - - task = MockTask( - eval_strategy="sequential_block", - n_blocks=1, - block_idx=0, - eval_blocks={"cheap": {0}, "expensive": {0}}, - evaluate_fn=fake_evaluate, - ) - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["cheap", "expensive"], - df_few_shots=task.pop_datapoints(n=1), - ) - - optimizer.incumbents = [] - optimizer.non_incumbents = [Prompt("cheap"), Prompt("expensive")] - optimizer.population_size = 1 - - optimizer._select_survivors() - - assert len(optimizer.non_incumbents) == 1 - assert optimizer.non_incumbents[0].instruction == "cheap" - - -def test_capoeira_step_invokes_hooks(mock_meta_llm, mock_predictor, mock_df): - task = MockTask() - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["p1", "p2"], - df_few_shots=mock_df, - ) - - def fake_eval(prompts, *_, **__): - n = len(prompts) - return EvalResult( - scores=np.zeros((n, 1), dtype=float), - agg_scores=np.arange(n, dtype=float), - sequences=np.array([[""] for _ in range(n)], dtype=object), - input_tokens=np.ones((n, 1)), - output_tokens=np.zeros((n, 1)), - agg_input_tokens=np.ones(n), - agg_output_tokens=np.zeros(n), - ) - - optimizer.task.evaluate = fake_eval # type: ignore[assignment] - optimizer.incumbents = [Prompt("inc")] - optimizer.prompts = [Prompt("p1"), Prompt("p2")] - - with patch("promptolution.optimizers.capoeira.perform_crossover", return_value=[Prompt("c1")]), patch( - "promptolution.optimizers.capoeira.perform_mutation", return_value=[Prompt("m1")] - ), patch.object(optimizer, "_do_intensification") as do_int, patch.object( - optimizer, "_advance_one_incumbent" - ) as adv_inc, patch.object( - optimizer, "_select_survivors" - ) as sel: - optimizer._step() - - assert do_int.call_count == 1 - assert adv_inc.call_count == 1 - assert sel.call_count == 1 - - -def test_capoeira_do_intensification_updates_incumbents(mock_meta_llm, mock_predictor): - def fake_eval(prompts, *_, **__): - n = len(prompts) - scores = np.arange(1, n + 1, dtype=float).reshape(n, 1) - return EvalResult( - scores=scores, - agg_scores=scores.flatten(), - sequences=np.array([[""] for _ in range(n)], dtype=object), - input_tokens=np.ones((n, 1)), - output_tokens=np.zeros((n, 1)), - agg_input_tokens=np.ones(n), - agg_output_tokens=np.zeros(n), - ) - - task = MockTask(eval_strategy="sequential_block", n_blocks=2, block_idx=0, evaluate_fn=fake_eval) - challenger = Prompt("chal") - inc1, inc2 = Prompt("i1"), Prompt("i2") - task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {0}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["p"], - df_few_shots=task.pop_datapoints(n=1), - ) - optimizer.incumbents = [inc1, inc2] - update_mock = patch.object(optimizer, "_update_incumbent_front", autospec=True).start() - - with patch("random.choice", side_effect=lambda seq: seq[0]): - optimizer._do_intensification(challenger) - - patch.stopall() - assert challenger in optimizer.incumbents - update_mock.assert_called_once() - - -def test_capoeira_do_intensification_bootstrap_no_common_blocks(mock_meta_llm, mock_predictor): - def fake_eval(prompts, *_, **__): - n = len(prompts) - return EvalResult( - scores=np.zeros((n, 1)), - agg_scores=np.zeros(n), - sequences=np.array([[""] for _ in range(n)], dtype=object), - input_tokens=np.zeros((n, 1)), - output_tokens=np.zeros((n, 1)), - agg_input_tokens=np.zeros(n), - agg_output_tokens=np.zeros(n), - ) - - task = MockTask(eval_strategy="sequential_block", n_blocks=3, block_idx=0, evaluate_fn=fake_eval) - inc1, inc2, challenger = Prompt("i1"), Prompt("i2"), Prompt("chal") - task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {1}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["p"], - df_few_shots=task.pop_datapoints(n=1), - ) - optimizer.incumbents = [inc1, inc2] - with patch("random.randrange", return_value=2), patch.object( - optimizer, "_update_incumbent_front", autospec=True - ) as upd: - optimizer._do_intensification(challenger) - - assert task.block_idx == 2 - assert challenger in optimizer.incumbents - upd.assert_called_once_with(blocks={2}) - - -def test_capoeira_do_intensification_running_mean_path(monkeypatch, mock_meta_llm, mock_predictor): - task = MockTask(eval_strategy="sequential_block", n_blocks=2, block_idx=0) - inc1, inc2, challenger = Prompt("i1"), Prompt("i2"), Prompt("chal") - task.prompt_evaluated_blocks = {str(inc1): {0, 1}, str(inc2): {0, 1}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["p"], - df_few_shots=task.pop_datapoints(n=1), - ) - optimizer.incumbents = [inc1, inc2] - - vec1 = np.array([[0.1, -0.1], [0.2, -0.2], [0.15, -0.15]]) - vec2 = np.array([[0.2, -0.2], [0.3, -0.3], [0.25, -0.25]]) - - calls: list[tuple] = [] - - def fake_is_dom(_self, v1, v2): - calls.append((v1.copy(), v2.copy())) - return False - - monkeypatch.setattr(Capoeira, "_is_dominated", fake_is_dom) - - with patch.object(Capoeira, "_get_objective_vectors", side_effect=[vec1, vec2]), patch( - "random.choice", side_effect=lambda seq: list(seq)[0] - ), patch.object(optimizer, "_update_incumbent_front", autospec=True) as upd: - optimizer._do_intensification(challenger) - - # fold_vec path should call dominance check at least once - assert calls, "_is_dominated should be invoked when challenger_mean already set" - assert challenger in optimizer.incumbents - upd.assert_called_once() - - -def test_capoeira_do_intensification_dominated_challenger(monkeypatch, mock_meta_llm, mock_predictor): - task = MockTask(eval_strategy="sequential_block", n_blocks=1, block_idx=0) - inc1, inc2, challenger = Prompt("i1"), Prompt("i2"), Prompt("chal") - task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {0}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["p"], - df_few_shots=task.pop_datapoints(n=1), - ) - optimizer.incumbents = [inc1, inc2] - - dominated_vecs = np.array([[0.9, -0.1], [0.8, -0.1], [0.1, -0.1]]) - - with patch.object(Capoeira, "_get_objective_vectors", return_value=dominated_vecs), patch( - "random.choice", side_effect=lambda seq: list(seq)[0] - ): - optimizer._do_intensification(challenger) - - assert challenger in optimizer.non_incumbents - assert challenger not in optimizer.incumbents - - -def test_capoeira_update_incumbent_front_demotes(mock_meta_llm, mock_predictor): - def fake_eval(prompts, *_, **__): - scores = np.array([0.3, 0.1], dtype=float) - return EvalResult( - scores=scores.reshape(-1, 1), - agg_scores=scores, - sequences=np.array([["s1"], ["s2"]], dtype=object), - input_tokens=np.zeros((2, 1)), - output_tokens=np.zeros((2, 1)), - agg_input_tokens=np.zeros(2), - agg_output_tokens=np.zeros(2), - ) - - task = MockTask(eval_strategy="sequential_block", n_blocks=1, evaluate_fn=fake_eval) - inc1, inc2 = Prompt("best"), Prompt("worst") - task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {0}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["p"], - df_few_shots=task.pop_datapoints(n=1), - ) - optimizer.incumbents = [inc1, inc2] - - optimizer._update_incumbent_front() - - assert optimizer.incumbents == [inc1] - assert inc2 in optimizer.non_incumbents - - -def test_capoeira_advance_one_incumbent_no_gapblocks(mock_meta_llm, mock_predictor): - task = MockTask(eval_strategy="sequential_block", n_blocks=2, block_idx=0) - inc = Prompt("p1") - task.prompt_evaluated_blocks = {str(inc): {0, 1}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["p1"], - df_few_shots=task.pop_datapoints(n=1), - ) - optimizer.incumbents = [inc] - - called = {"evaluate": 0} - - def no_call(*args, **kwargs): - called["evaluate"] += 1 - raise AssertionError("evaluate should not be called when no new blocks") - - task.evaluate = no_call # type: ignore[assignment] - - optimizer._advance_one_incumbent() - - assert called["evaluate"] == 0 - - -def test_capoeira_get_closest_incumbent_returns_nearest(): - challenger = np.array([0.5, 0.5]) - incumbents = np.array([[0.0, 0.0], [0.6, 0.6]]) - res = Capoeira._get_closest_incumbent(None, challenger, incumbents) - assert np.allclose(res, incumbents[1]) - - -def test_capoeira_objective_vectors_multiobjective(mock_meta_llm, mock_predictor, mock_df): - t1 = MockTask(df=mock_df, n_subsamples=1, n_blocks=1) - t2 = MockTask(df=mock_df, n_subsamples=1, n_blocks=1) - multi_task = MultiObjectiveTask(tasks=[t1, t2]) - - optimizer = Capoeira( - predictor=mock_predictor, - task=multi_task, - meta_llm=mock_meta_llm, - initial_prompts=["p"], - df_few_shots=mock_df, - ) - - result = MultiObjectiveEvalResult( - scores=[np.array([[0.1], [0.2]]), np.array([[0.3], [0.4]])], - agg_scores=[np.array([0.1, 0.2]), np.array([0.3, 0.4])], - sequences=np.array([["s1"], ["s2"]], dtype=object), - input_tokens=np.array([[1.0], [2.0]]), - output_tokens=np.array([[0.0], [0.0]]), - agg_input_tokens=np.array([1.0, 2.0]), - agg_output_tokens=np.array([0.0, 0.0]), - ) - - vecs = optimizer._get_objective_vectors(result) - assert vecs.shape == (2, 3) - assert np.allclose(vecs[:, 0], [0.1, 0.2]) - assert np.allclose(vecs[:, 1], [0.3, 0.4]) - assert np.allclose(vecs[:, 2], -np.array([1.0, 2.0])) - - -def test_capoeira_advance_one_incumbent_chooses_gap(mock_meta_llm, mock_predictor): - def fake_eval(*_, **__): - return EvalResult( - scores=np.array([[0.0]]), - agg_scores=np.array([0.0]), - sequences=np.array([[""]], dtype=object), - input_tokens=np.array([[0.0]]), - output_tokens=np.array([[0.0]]), - agg_input_tokens=np.array([0.0]), - agg_output_tokens=np.array([0.0]), - ) - - task = MockTask(eval_strategy="sequential_block", n_blocks=3, block_idx=0, evaluate_fn=fake_eval) - p1, p2 = Prompt("p1"), Prompt("p2") - task.prompt_evaluated_blocks = {str(p1): {0}, str(p2): {0, 2}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["p1", "p2"], - df_few_shots=task.pop_datapoints(n=1), - ) - optimizer.incumbents = [p1, p2] - - with patch("random.choice", side_effect=lambda seq: list(seq)[0]): - optimizer._advance_one_incumbent() - - assert task.block_idx == 2 - - -def test_capoeira_select_survivors_heterogeneous_removes_lowest(mock_meta_llm, mock_predictor): - task = MockTask(eval_strategy="sequential_block", n_blocks=3) - c1, c2 = Prompt("c1"), Prompt("c2") - task.prompt_evaluated_blocks = {str(c1): {0}, str(c2): {1}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["inc"], - df_few_shots=task.pop_datapoints(n=1), - ) - optimizer.incumbents = [] - optimizer.non_incumbents = [c1, c2] - optimizer.population_size = 1 - - with patch("random.choice", side_effect=lambda seq: list(seq)[0]): - optimizer._select_survivors() - - assert len(optimizer.non_incumbents) == 1 - - -def test_capoeira_select_survivors_incumbent_only(mock_meta_llm, mock_predictor): - def fake_eval(prompts, *_, **__): - n = len(prompts) - vals = np.linspace(0.1, 0.2, n) - return EvalResult( - scores=np.tile(vals.reshape(n, 1), (1, 1)), - agg_scores=vals, - sequences=np.array([[""] for _ in range(n)], dtype=object), - input_tokens=np.ones((n, 1)), - output_tokens=np.zeros((n, 1)), - agg_input_tokens=np.ones(n), - agg_output_tokens=np.zeros(n), - ) - - task = MockTask(eval_strategy="sequential_block", n_blocks=2, evaluate_fn=fake_eval) - inc1, inc2 = Prompt("i1"), Prompt("i2") - task.prompt_evaluated_blocks = {str(inc1): {0}, str(inc2): {0}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["i1", "i2"], - df_few_shots=task.pop_datapoints(n=1), - ) - optimizer.incumbents = [inc1, inc2] - optimizer.non_incumbents = [] - optimizer.population_size = 1 - - optimizer._select_survivors() - - assert len(optimizer.incumbents) == 1 - - -def test_capoeira_get_common_blocks(mock_meta_llm, mock_predictor): - task = MockTask(eval_strategy="sequential_block", n_blocks=2) - p1, p2 = Prompt("p1"), Prompt("p2") - task.prompt_evaluated_blocks = {str(p1): {0, 1}, str(p2): {1}} - - optimizer = Capoeira( - predictor=mock_predictor, - task=task, - meta_llm=mock_meta_llm, - initial_prompts=["p1", "p2"], - df_few_shots=task.pop_datapoints(n=1), - ) - - common = optimizer._get_common_blocks([p1, p2]) - assert common == {1} - - -def test_capoeira_is_dominated_logic(): - assert Capoeira._is_dominated(np.array([0.1, 0.1]), np.array([0.2, 0.2])) - assert not Capoeira._is_dominated(np.array([0.3, 0.2]), np.array([0.3, 0.2])) - assert not Capoeira._is_dominated(np.array([0.4, 0.5]), np.array([0.3, 0.6])) - - -def test_capoeira_calculate_crowding_distance_three_points(): - vecs = np.array([[0.0, 0.0], [0.5, 0.5], [1.0, 1.0]]) - dists = Capoeira._calculate_crowding_distance(vecs) - assert np.isinf(dists[[0, -1]]).all() - assert dists[1] > 0 diff --git a/tests/tasks/test_base_task.py b/tests/tasks/test_base_task.py index 16e9fa06..4ed7529c 100644 --- a/tests/tasks/test_base_task.py +++ b/tests/tasks/test_base_task.py @@ -35,11 +35,9 @@ def test_subsample_and_block_controls(small_task): task.increment_block_idx() assert task.block_idx == 1 % task.n_blocks if task.n_blocks else 0 - task.set_block_idx([0, 1, 2]) - xs2, _ = task.subsample() + xs2, _ = task.subsample(block_idx=[0, 1, 2]) assert set(xs2) == set(task.xs) - task.set_block_idx(0) popped = task.pop_datapoints(n=1) assert len(popped) == 1 assert len(task.xs) == 2 @@ -121,7 +119,7 @@ def test_evaluate_with_block_list_updates_blocks(predictor, small_task): prompts = [Prompt("p1"), Prompt("p2")] task.evaluate(prompts, predictor) for p in prompts: - assert task.prompt_evaluated_blocks[str(p)] == {0, 1} + assert task.prompt_evaluated_blocks[p] == [0, 1] def test_task_config_applied(): diff --git a/tests/tasks/test_multi_objective_task.py b/tests/tasks/test_multi_objective_task.py index 7efcae7d..cfb4caad 100644 --- a/tests/tasks/test_multi_objective_task.py +++ b/tests/tasks/test_multi_objective_task.py @@ -40,7 +40,7 @@ def test_multi_objective_shares_block_and_caches(): assert len(t1.eval_cache) == len(t2.eval_cache) assert res.input_tokens.shape[0] == 1 - assert multi.prompt_evaluated_blocks[str(prompt)] == {1} + assert multi.prompt_evaluated_blocks[prompt] == [1] def test_multi_objective_requires_tasks(): @@ -73,7 +73,7 @@ def make_task(): assert np.allclose(multi_res.agg_scores[0], res1.agg_scores) assert np.allclose(multi_res.agg_scores[1], res2.agg_scores) assert multi_res.sequences.shape == res1.sequences.shape - assert multi.prompt_evaluated_blocks[str(prompt)] == {1} + assert multi.prompt_evaluated_blocks[prompt] == [1] class ConstantTask(BaseTask): diff --git a/tests/tasks/test_reward_tasks.py b/tests/tasks/test_reward_tasks.py index b6134573..eb37ab77 100644 --- a/tests/tasks/test_reward_tasks.py +++ b/tests/tasks/test_reward_tasks.py @@ -40,6 +40,7 @@ def test_reward_task_passes_reward_columns(): df = pd.DataFrame({"x": ["a", "b", "c"], "reward": [0.1, 0.2, 0.3]}) seen_rewards: list[float] = [] + def reward_fn(prediction: str, reward: float) -> float: seen_rewards.append(reward) return reward if prediction == "keep" else -1.0 diff --git a/tutorials/getting_started.ipynb b/tutorials/getting_started.ipynb index 7761168c..2c140f6d 100644 --- a/tutorials/getting_started.ipynb +++ b/tutorials/getting_started.ipynb @@ -163,7 +163,7 @@ "metadata": {}, "source": [ "Here's an explanation of each configuration parameter in the ExperimentConfig:\n", - "- `optimizer`: The algorithm used for prompt optimization. Currently we support \"capo\", \"capoeira\", \"evopromptga\", \"evopromptde\", and \"opro\". For this example, we use \"capo\" as it is capable of leveraging few-shot examples.\n", + "- `optimizer`: The algorithm used for prompt optimization. Currently we support \"capo\", \"evopromptga\", \"evopromptde\", and \"opro\". For this example, we use \"capo\" as it is capable of leveraging few-shot examples.\n", "- `task_description`: A string describing the task you're optimizing prompts for. This is used to provide the meta-llm with context about your task.\n", "- `prompts`: A list of initial prompt strings that will be used as the starting point for optimization.\n", "- `n_steps`: The number of optimization steps to run. Higher values allow more exploration and refinement but require more API calls and computational resources.\n", diff --git a/tutorials/reward_task_tutorial.ipynb b/tutorials/reward_task_tutorial.ipynb index 91f4af72..e0922408 100644 --- a/tutorials/reward_task_tutorial.ipynb +++ b/tutorials/reward_task_tutorial.ipynb @@ -201,7 +201,7 @@ "metadata": {}, "source": [ "Here's an explanation of each configuration parameter in the ExperimentConfig:\n", - "- `optimizer`: The algorithm used for prompt optimization. Currently we support \"capo\", \"capoeira\", \"evopromptga\", \"evopromptde\", and \"opro\". For this example, we use \"capo\" as it is capable of leveraging few-shot examples.\n", + "- `optimizer`: The algorithm used for prompt optimization. Currently we support \"capo\", \"evopromptga\", \"evopromptde\", and \"opro\". For this example, we use \"capo\" as it is capable of leveraging few-shot examples.\n", "- `task_description`: A string describing the task you're optimizing prompts for. This is used to provide the meta-llm with context about your task.\n", "- `prompts`: A list of initial prompt strings that will be used as the starting point for optimization.\n", "- `n_steps`: The number of optimization steps to run. Higher values allow more exploration and refinement but require more API calls and computational resources.\n", From c0f02be034ffbf0a4d3a9fa597f8e0b7abe2c493 Mon Sep 17 00:00:00 2001 From: finitearth Date: Wed, 14 Jan 2026 16:38:18 +0100 Subject: [PATCH 38/53] green test --- promptolution/tasks/base_task.py | 19 +++++++++++-------- tests/optimizers/test_capo.py | 22 ---------------------- 2 files changed, 11 insertions(+), 30 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 0e3f02ac..526896bb 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -97,7 +97,7 @@ def __init__( self.prompt_evaluated_blocks: Dict[Prompt, List[int]] = {} # prompt_str: set of evaluated block indices - def subsample(self, eval_strategy: Optional["EvalStrategy"] = None) -> Tuple[List[str], List[str]]: + def subsample(self, eval_strategy: Optional["EvalStrategy"] = None, block_idx: int | list[int] | None = None) -> Tuple[List[str], List[str]]: """Subsample the dataset based on the specified parameters. Args: @@ -106,6 +106,11 @@ def subsample(self, eval_strategy: Optional["EvalStrategy"] = None) -> Tuple[Lis Returns: Tuple[List[str], List[str]]: Subsampled input data and labels. """ + if block_idx is not None and isinstance(block_idx, int): + block_idx = [block_idx] + + if block_idx is not None: + return [self.xs[i] for i in block_idx], [self.ys[i] for i in block_idx] if eval_strategy is None: eval_strategy = self.eval_strategy @@ -181,7 +186,7 @@ def _collect_results_from_cache( datapoint_seqs: List[str] = [] for x, y in zip(xs, ys): cache_key = self._cache_key(prompt, x, str(y)) - datapoint_score = self.eval_cache.get(cache_key, np.nan) + 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) @@ -246,6 +251,7 @@ def evaluate( 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. @@ -254,7 +260,7 @@ def evaluate( """ prompts_list: List[Prompt] = [prompts] if isinstance(prompts, Prompt) else list(prompts) eval_strategy = eval_strategy or self.eval_strategy - xs, ys = self.subsample(eval_strategy=eval_strategy) + xs, ys = self.subsample(eval_strategy=eval_strategy, block_idx=block_idx) ( prompts_to_evaluate, xs_to_evaluate, @@ -365,15 +371,12 @@ def reset_block_idx(self) -> None: raise ValueError("Block reset is only valid for block subsampling.") self.block_idx = 0 - def set_block_idx(self, idx: Union[int, List[int]]) -> None: + 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.") - if isinstance(idx, list): - assert all(0 <= i < self.n_blocks for i in idx), "Block indices must be integers within valid range" - else: - assert isinstance(idx, int), "Block index must be an integer" + assert isinstance(idx, int), "Block index must be an integer" self.block_idx = idx diff --git a/tests/optimizers/test_capo.py b/tests/optimizers/test_capo.py index b21b1c6a..4d5a6c24 100644 --- a/tests/optimizers/test_capo.py +++ b/tests/optimizers/test_capo.py @@ -172,28 +172,6 @@ def test_mutate(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_ assert len(mutated) == 2 -def test_do_racing(mock_meta_llm, mock_predictor, initial_prompts, mock_df): - mock_task = MockTask(predetermined_scores=[0.89, 0.9] * 3) - optimizer = CAPO( - predictor=mock_predictor, - task=mock_task, - meta_llm=mock_meta_llm, - initial_prompts=initial_prompts, - df_few_shots=pd.concat([mock_df] * 5, ignore_index=True), - ) - optimizer._pre_optimization_loop() - survivors, scores = optimizer._do_racing( - [Prompt("good instruction", ["Example 1"]), Prompt("better instruction", ["Example 2"])], 1 - ) - assert len(survivors) == 1 - assert len(scores) == 1 - - assert "better instruction" in survivors[0].instruction - - assert mock_task.reset_block_idx.call_count == 2 - assert mock_task.increment_block_idx.call_count == 2 - - def test_capo_crossover_prompt(mock_meta_llm, mock_predictor, initial_prompts, mock_task, mock_df): """Test that when _crossover is called, the mock_meta_llm received a call with the correct meta prompt.""" optimizer = CAPO( From e9cd844d2d21c1e12f56a8486094e06d0a5e652b Mon Sep 17 00:00:00 2001 From: mo374z Date: Thu, 15 Jan 2026 22:51:22 +0100 Subject: [PATCH 39/53] change get evaluated blocks function to work also with one prompt --- promptolution/tasks/base_task.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 526896bb..d411fb7b 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -97,7 +97,9 @@ def __init__( 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: int | list[int] | None = None) -> Tuple[List[str], List[str]]: + def subsample( + self, eval_strategy: Optional["EvalStrategy"] = None, block_idx: int | list[int] | None = None + ) -> Tuple[List[str], List[str]]: """Subsample the dataset based on the specified parameters. Args: @@ -108,7 +110,7 @@ def subsample(self, eval_strategy: Optional["EvalStrategy"] = None, block_idx: i """ if block_idx is not None and isinstance(block_idx, int): block_idx = [block_idx] - + if block_idx is not None: return [self.xs[i] for i in block_idx], [self.ys[i] for i in block_idx] if eval_strategy is None: @@ -380,6 +382,7 @@ def set_block_idx(self, idx: int) -> None: self.block_idx = idx - def get_evaluated_blocks(self, prompts: List[Prompt]) -> Dict[Prompt, List[int]]: + def get_evaluated_blocks(self, prompts: Union[Prompt, List[Prompt]]) -> Dict[Prompt, List[int]]: """Return mapping of prompt string to evaluated block indices.""" - return {p: list(self.prompt_evaluated_blocks.get(p, [])) for p in prompts} \ No newline at end of file + 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} From 5e745328e9390740d1ecbb475eab948cea8e79f6 Mon Sep 17 00:00:00 2001 From: mo374z Date: Thu, 15 Jan 2026 23:12:02 +0100 Subject: [PATCH 40/53] allow for empty cache at key --- promptolution/tasks/base_task.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index d411fb7b..d1c77ea5 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -188,6 +188,8 @@ def _collect_results_from_cache( 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: + continue datapoint_score = self.eval_cache[cache_key] datapoint_scores.append(datapoint_score) datapoint_seqs.append(self.seq_cache.get(cache_key, "")) From a2101836557defc0e89025198c2ed2dfb91812e4 Mon Sep 17 00:00:00 2001 From: mo374z Date: Thu, 15 Jan 2026 23:21:22 +0100 Subject: [PATCH 41/53] some more cache issues --- promptolution/tasks/base_task.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index d1c77ea5..a51df284 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -220,6 +220,8 @@ def _compute_costs( seq_token_counts: List[float] = [] for x, y in zip(xs, ys): cache_key = self._cache_key(prompt, x, str(y)) + if cache_key not in self.seq_cache: + continue seq_text = self.seq_cache[cache_key] seq_token_counts.append(token_counter(seq_text)) From 745f7222439da5bde0ec5179d6fe3059ef52eee0 Mon Sep 17 00:00:00 2001 From: mo374z Date: Thu, 15 Jan 2026 23:47:53 +0100 Subject: [PATCH 42/53] change comput cost function --- promptolution/tasks/base_task.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index a51df284..c43dd3e7 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -213,20 +213,20 @@ def _compute_costs( per_prompt_inputs: List[np.ndarray] = [] per_prompt_outputs: List[np.ndarray] = [] - input_token_counts = np.array([token_counter(x) for x in xs], dtype=float) - for prompt in prompts: prompt_tokens = 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: continue seq_text = self.seq_cache[cache_key] seq_token_counts.append(token_counter(seq_text)) + input_token_counts.append(token_counter(prompt.construct_prompt() + " " + x)) - prompt_input_tokens = prompt_tokens + input_token_counts - output_token_counts = np.array(seq_token_counts, dtype=float) - input_token_counts + prompt_input_tokens = prompt_tokens + np.array(input_token_counts, dtype=float) + 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) From 05d5ebc5d9b76e04e4b6d18f9d0bfb559c719345 Mon Sep 17 00:00:00 2001 From: mo374z Date: Fri, 16 Jan 2026 01:28:09 +0100 Subject: [PATCH 43/53] some fixes --- promptolution/optimizers/capo.py | 2 +- promptolution/tasks/base_task.py | 10 ++++++---- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index 961cd9f6..2c215f6b 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -175,7 +175,7 @@ def _do_racing(self, candidates: List[Prompt], k: int) -> Tuple[List[Prompt], Li self.task.increment_block_idx() final_result = self.task.evaluate(candidates, self.predictor, eval_strategy="evaluated") - avg_scores = final_result.scores.tolist() + avg_scores = final_result.agg_scores.tolist() prompts, avg_scores = sort_prompts_by_scores(candidates, avg_scores, top_k=k) return prompts, avg_scores diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index c43dd3e7..3e3a2a0f 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -189,10 +189,12 @@ def _collect_results_from_cache( for x, y in zip(xs, ys): cache_key = self._cache_key(prompt, x, str(y)) if cache_key not in self.eval_cache: - continue - datapoint_score = self.eval_cache[cache_key] - datapoint_scores.append(datapoint_score) - datapoint_seqs.append(self.seq_cache.get(cache_key, "")) + 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) From 70ae29605487f8aadfb9c33e18ad1ffb7328f31f Mon Sep 17 00:00:00 2001 From: mo374z Date: Fri, 16 Jan 2026 01:46:37 +0100 Subject: [PATCH 44/53] fix compute costs function --- promptolution/tasks/base_task.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 3e3a2a0f..6277adca 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -222,6 +222,9 @@ def _compute_costs( 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)) @@ -236,8 +239,8 @@ def _compute_costs( inputs_array = np.vstack(per_prompt_inputs) outputs_array = np.vstack(per_prompt_outputs) - agg_input_tokens = inputs_array.mean(axis=1) - agg_output_tokens = outputs_array.mean(axis=1) + 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 From 276e5fa0ec964457a589667b61c96be9f0275200 Mon Sep 17 00:00:00 2001 From: mo374z Date: Fri, 16 Jan 2026 02:49:36 +0100 Subject: [PATCH 45/53] fix token count --- promptolution/tasks/base_task.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 6277adca..65ecafa0 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -216,7 +216,7 @@ def _compute_costs( per_prompt_outputs: List[np.ndarray] = [] for prompt in prompts: - prompt_tokens = token_counter(prompt.construct_prompt()) + prompt_token_count = token_counter(prompt.construct_prompt()) seq_token_counts: List[float] = [] input_token_counts = [] for x, y in zip(xs, ys): @@ -228,9 +228,9 @@ def _compute_costs( continue seq_text = self.seq_cache[cache_key] seq_token_counts.append(token_counter(seq_text)) - input_token_counts.append(token_counter(prompt.construct_prompt() + " " + x)) + input_token_counts.append(token_counter(x)) - prompt_input_tokens = prompt_tokens + np.array(input_token_counts, dtype=float) + 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)) From 83a6f9d885a86f8a96dd137fcd21cfbb407dad94 Mon Sep 17 00:00:00 2001 From: finitearth Date: Fri, 16 Jan 2026 15:02:31 +0100 Subject: [PATCH 46/53] fix tracking of blocks --- promptolution/tasks/base_task.py | 45 ++++++++++++--------- promptolution/tasks/multi_objective_task.py | 7 ---- 2 files changed, 27 insertions(+), 25 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 65ecafa0..1615fd39 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -88,7 +88,7 @@ def __init__( # If no y_column is provided, create a dummy y array self.ys = [""] * len(self.xs) - self.block_idx: int | list[int] = 0 + 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) @@ -98,18 +98,17 @@ def __init__( 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: int | list[int] | None = None + 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 and isinstance(block_idx, int): - block_idx = [block_idx] if block_idx is not None: return [self.xs[i] for i in block_idx], [self.ys[i] for i in block_idx] @@ -128,17 +127,9 @@ def subsample( 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": - if isinstance(self.block_idx, list): - block_indices: List[int] = [] - for block_id in self.block_idx: - start_idx = block_id * self.n_subsamples - end_idx = min((block_id + 1) * self.n_subsamples, len(self.xs)) - block_indices.extend(range(start_idx, end_idx)) - indices = np.array(sorted(set(block_indices)), dtype=int) - 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) + 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: @@ -268,9 +259,20 @@ def evaluate( 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, @@ -298,10 +300,17 @@ def evaluate( # Record evaluated block for block strategies for prompt in prompts_list: - if isinstance(self.block_idx, list): - self.prompt_evaluated_blocks.setdefault(prompt, []).extend(self.block_idx) - else: + if eval_strategy == "evaluated": + continue + elif block_idx is not None: + self.prompt_evaluated_blocks.setdefault(prompt, []).extend(block_idx) + elif eval_strategy in ["sequential_block", "random_block"]: self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx) + else: + 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 diff --git a/promptolution/tasks/multi_objective_task.py b/promptolution/tasks/multi_objective_task.py index 3e814712..b1163510 100644 --- a/promptolution/tasks/multi_objective_task.py +++ b/promptolution/tasks/multi_objective_task.py @@ -130,13 +130,6 @@ def evaluate( # type: ignore prompts_list, xs, ys, predictor ) - # Record evaluated block for block strategies - for prompt in prompts_list: - block_set = task.prompt_evaluated_blocks.setdefault(prompt, []) - if isinstance(task.block_idx, list): - block_set.extend(task.block_idx) - else: - block_set.append(task.block_idx) per_task_results.append( EvalResult( scores=scores_array, From 4bcda09264d6027613c2f8164c318c8ec21a3bed Mon Sep 17 00:00:00 2001 From: finitearth Date: Sun, 18 Jan 2026 17:49:39 +0100 Subject: [PATCH 47/53] fix block idx subsampling --- promptolution/tasks/base_task.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index 1615fd39..b060ea39 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -111,7 +111,14 @@ def subsample( """ if block_idx is not None: - return [self.xs[i] for i in block_idx], [self.ys[i] for i in block_idx] + indices = [] + 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 @@ -300,16 +307,15 @@ def evaluate( # Record evaluated block for block strategies for prompt in prompts_list: - if eval_strategy == "evaluated": - continue - elif block_idx is not None: + if block_idx is not None: self.prompt_evaluated_blocks.setdefault(prompt, []).extend(block_idx) elif eval_strategy in ["sequential_block", "random_block"]: self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx) - else: + 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( From ace326eb507b6e3a71bc94dbb48ac1332783da88 Mon Sep 17 00:00:00 2001 From: mo374z Date: Fri, 23 Jan 2026 12:25:14 +0100 Subject: [PATCH 48/53] allow for y_column in reward task --- promptolution/tasks/reward_tasks.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/promptolution/tasks/reward_tasks.py b/promptolution/tasks/reward_tasks.py index 7bb79277..cb4f922e 100644 --- a/promptolution/tasks/reward_tasks.py +++ b/promptolution/tasks/reward_tasks.py @@ -27,6 +27,7 @@ def __init__( 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, @@ -40,6 +41,7 @@ def __init__( 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. @@ -52,6 +54,7 @@ def __init__( super().__init__( df=df, x_column=x_column, + y_column=y_column, task_description=task_description, n_subsamples=n_subsamples, eval_strategy=eval_strategy, From 36ce80eb1bd455ad23a72cd23391c7d4bec3b5ee Mon Sep 17 00:00:00 2001 From: mo374z Date: Sun, 15 Feb 2026 17:54:24 +0100 Subject: [PATCH 49/53] formatting --- promptolution/tasks/base_task.py | 19 +++++++------------ tests/optimizers/test_capo.py | 2 -- 2 files changed, 7 insertions(+), 14 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index b060ea39..f759f829 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -109,16 +109,15 @@ def subsample( Returns: Tuple[List[str], List[str]]: Subsampled input data and labels. """ - if block_idx is not None: - indices = [] + 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 @@ -266,7 +265,7 @@ def evaluate( 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. @@ -276,10 +275,10 @@ def evaluate( """ 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, @@ -312,11 +311,7 @@ def evaluate( elif eval_strategy in ["sequential_block", "random_block"]: 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)) - ) - - + 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 diff --git a/tests/optimizers/test_capo.py b/tests/optimizers/test_capo.py index 4d5a6c24..63ce88c6 100644 --- a/tests/optimizers/test_capo.py +++ b/tests/optimizers/test_capo.py @@ -2,8 +2,6 @@ import pandas as pd -from tests.mocks.mock_task import MockTask - from promptolution.optimizers.capo import CAPO from promptolution.utils.capo_utils import build_few_shot_examples, perform_crossover, perform_mutation from promptolution.utils.prompt import Prompt From 6ecbfff7af11f311f6a27f55d131b4f0a62d6a48 Mon Sep 17 00:00:00 2001 From: mo374z Date: Sun, 15 Feb 2026 18:33:12 +0100 Subject: [PATCH 50/53] fixes in mo task and base task block idx handling --- promptolution/tasks/base_task.py | 24 +++++++++++++++------ promptolution/tasks/multi_objective_task.py | 15 ++++++++++--- tests/llms/test_vllm.py | 20 +++++++++++++---- 3 files changed, 46 insertions(+), 13 deletions(-) diff --git a/promptolution/tasks/base_task.py b/promptolution/tasks/base_task.py index f759f829..a733be4e 100644 --- a/promptolution/tasks/base_task.py +++ b/promptolution/tasks/base_task.py @@ -133,11 +133,19 @@ def subsample( 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": - 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] + # 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}'") @@ -309,7 +317,11 @@ def evaluate( if block_idx is not None: self.prompt_evaluated_blocks.setdefault(prompt, []).extend(block_idx) elif eval_strategy in ["sequential_block", "random_block"]: - self.prompt_evaluated_blocks.setdefault(prompt, []).append(self.block_idx) + # 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))) diff --git a/promptolution/tasks/multi_objective_task.py b/promptolution/tasks/multi_objective_task.py index b1163510..a4844cd2 100644 --- a/promptolution/tasks/multi_objective_task.py +++ b/promptolution/tasks/multi_objective_task.py @@ -145,10 +145,19 @@ def evaluate( # type: ignore stacked_scores = [r.scores for r in per_task_results] stacked_agg_scores = [r.agg_scores for r in per_task_results] - # Mirror evaluated block bookkeeping using the first task for parity with BaseTask. - first_task = self.tasks[0] + # 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] - self.prompt_evaluated_blocks = {p: first_task.prompt_evaluated_blocks[p] for p in prompts_list} if self._scalarized_objective: return EvalResult( diff --git a/tests/llms/test_vllm.py b/tests/llms/test_vllm.py index 6eef0310..84ae44f6 100644 --- a/tests/llms/test_vllm.py +++ b/tests/llms/test_vllm.py @@ -42,6 +42,9 @@ def mock_generate_side_effect(prompts_list, *args, **kwargs): # This is the most critical change. mock_from_pretrained.return_value = mock_tokenizer_instance + # 4. Make sure llm_instance.get_tokenizer() returns the mock tokenizer + mock_llm_instance.get_tokenizer.return_value = mock_tokenizer_instance + # --- Sampling Params Mock Setup --- mock_sampling_params_instance = MagicMock() mock_sampling_params.return_value = mock_sampling_params_instance @@ -87,13 +90,22 @@ def test_vllm_with_auto_batch_size(mock_vllm_dependencies): mock_vllm_dependencies["llm_instance"].llm_engine.model_executor.cache_config.block_size = 16 # Create VLLM instance with batch_size=None to trigger auto calculation - vllm_instance = VLLM(model_id="mock-model", batch_size=None, max_model_len=2048) + # With max_num_batched_tokens=16384 and max_model_len=2048: + # token_limited = 16384 // 2048 = 8 + # batch_size = min(max_num_seqs=10, token_limited=8) = 8 + vllm_instance = VLLM( + model_id="mock-model", + batch_size=None, + max_model_len=2048, + llm_kwargs={"max_num_seqs": 10, "max_num_batched_tokens": 16384}, + ) # Verify batch_size is greater than zero assert vllm_instance.batch_size > 0, "Batch size should be greater than zero" - # With num_gpu_blocks=1000, block_size=16, max_model_len=2048 - # batch_size = int((1000 * 16 / 2048) * 0.95) = int(7.8125 * 0.95) = int(7.42) = 7 - assert vllm_instance.batch_size == 7, f"Expected batch_size=7, got {vllm_instance.batch_size}" + # With max_num_batched_tokens=16384, max_model_len=2048, max_num_seqs=10 + # token_limited = 16384 // 2048 = 8 + # batch_size = min(10, 8) = 8 + assert vllm_instance.batch_size == 8, f"Expected batch_size=8, got {vllm_instance.batch_size}" # Test with a single prompt prompts = ["Test prompt"] From 3368665ce632c4070776e1c0c25a9ef0f4d66cd9 Mon Sep 17 00:00:00 2001 From: mo374z Date: Sun, 15 Feb 2026 18:34:53 +0100 Subject: [PATCH 51/53] add import since it is requried for test cases --- promptolution/llms/vllm.py | 1 + 1 file changed, 1 insertion(+) diff --git a/promptolution/llms/vllm.py b/promptolution/llms/vllm.py index 4fab6012..cd91a253 100644 --- a/promptolution/llms/vllm.py +++ b/promptolution/llms/vllm.py @@ -13,6 +13,7 @@ logger = get_logger(__name__) try: + from transformers import AutoTokenizer # noqa: F401 (import required for testing) from vllm import LLM from vllm.sampling_params import SamplingParams From a4136f9badf25014e5dd4406cf62ebe3bcfe8be5 Mon Sep 17 00:00:00 2001 From: mo374z Date: Sun, 15 Feb 2026 17:39:04 +0000 Subject: [PATCH 52/53] Update coverage badge in README [skip ci] --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 6090afd8..a0552b14 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@ -![Coverage](https://img.shields.io/badge/Coverage-90%25-brightgreen) +![Coverage](https://img.shields.io/badge/Coverage-95%25-brightgreen) [![CI](https://github.com/automl/promptolution/actions/workflows/ci.yml/badge.svg?branch=main)](https://github.com/automl/promptolution/actions/workflows/ci.yml) [![Docs](https://github.com/automl/promptolution/actions/workflows/docs.yml/badge.svg?branch=main)](https://github.com/automl/promptolution/actions/workflows/docs.yml) ![Code Style](https://img.shields.io/badge/Code%20Style-black-black) From 0b75d1e270a092cd000f9978761b756948deba8b Mon Sep 17 00:00:00 2001 From: mo374z Date: Sun, 15 Feb 2026 19:03:03 +0100 Subject: [PATCH 53/53] change wording --- promptolution/optimizers/capo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/promptolution/optimizers/capo.py b/promptolution/optimizers/capo.py index 2c215f6b..43ff5630 100644 --- a/promptolution/optimizers/capo.py +++ b/promptolution/optimizers/capo.py @@ -29,7 +29,7 @@ class CAPO(BaseOptimizer): """CAPO: Cost-Aware Prompt Optimization. - This class implements an evolutionary algorithm for optimizing prompts in large language models + This class implements an evolutionary algorithm for optimizing prompts in LLMs by incorporating racing techniques and multi-objective optimization. It uses crossover, mutation, and racing based on evaluation scores and statistical tests to improve efficiency while balancing performance with prompt length. It is adapted from the paper "CAPO: Cost-Aware Prompt Optimization" by Zehle et al., 2025.