diff --git a/README.md b/README.md index 21277e4..8de9381 100644 --- a/README.md +++ b/README.md @@ -29,12 +29,13 @@ CoolPrompt is a framework for automatic prompt creation and optimization. ## Core features -- **Optimize prompts** with our APO methods: - - HyPER / HyPER Light - - RE-GPS - - RIDER - - PromptCompressor - - *(legacy/deprecated)*: ReflectivePrompt, DistillPrompt +- **Optimize prompts** with our APO methods: + - HyPER / HyPER Light + - RE-GPS + - RIDER + - BRAVE + - PromptCompressor + - *(legacy/deprecated)*: ReflectivePrompt, DistillPrompt - **LLM-Agnostic Choice:** work with your custom llm (from open-sourced to proprietary) using [supported Langchain LLMs](https://python.langchain.com/docs/integrations/llms/) - **Develop own custom APO method in one library** - **Generate synthetic evaluation data** when no input dataset is provided @@ -67,8 +68,9 @@ Compared metrics: | `hyper_light` | None | Low | Medium | Low | | `hyper` | Required | Medium | High | Medium | | `regps` | Required | High | Very High | High | -| `rider` | Required | Very High | Very High | Very High | -| `compress` | None | Low | Medium | Low | +| `rider` | Required | Very High | Very High | Very High | +| `brave` | Required | High | Very High | Budget-controlled | +| `compress` | None | Low | Medium | Low | | `reflective` | Required | High | High | High | | `distill` | Required | High | High | High | @@ -103,9 +105,24 @@ print(prompt_tuner.final_prompt) # well-structured, and vividly descriptive essay on the theme of autumn... ``` -

- CoolPrompt full optimization demo -

+

+ CoolPrompt full optimization demo +

+ +Run the data-driven BRAVE optimizer by selecting it as the method: + +```python +final_prompt = prompt_tuner.run( + "Classify the sentiment of the text: {text}", + task="classification", + dataset=["Great product", "Very disappointing"], + target=["positive", "negative"], + method="brave", + problem_description="Classify product-review sentiment.", + max_steps=20, + initial_budget_tokens=50_000, +) +``` ## Examples diff --git a/coolprompt/data_generator/generator.py b/coolprompt/data_generator/generator.py index b7330bf..4bb4a6a 100644 --- a/coolprompt/data_generator/generator.py +++ b/coolprompt/data_generator/generator.py @@ -18,6 +18,7 @@ PROBLEM_DESCRIPTION_BASED_ON_EXAMPLES_TEMPLATE, GENERATION_CORNER_CASE_GENERATING_TEMPLATE, CLASSIFICATION_CORNER_CASE_GENERATING_TEMPLATE, + CLASSIFICATION_PROBLEM_DESCRIPTION_BASED_ON_EXAMPLES_TEMPLATE, ) from coolprompt.utils.enums import Task from coolprompt.utils.logging_config import logger @@ -85,21 +86,50 @@ def _examples_to_str(self, examples: List[Tuple[str, str]]) -> str: """ return "\n\n".join([f"Input: {inp}\nOutput: {out}" for (inp, out) in examples]) + @staticmethod + def _extract_labels(targets: List) -> List[str]: + seen = set() + labels = [] + for t in targets: + key = str(t) + if key not in seen: + seen.add(key) + labels.append(key) + return labels + def _generate_problem_description( - self, prompt: str, examples: Optional[List[Tuple[str, str]]] = None + self, + prompt: str, + examples: Optional[List[Tuple[str, str]]] = None, + task: Optional[Task] = None, + labels: Optional[List[str]] = None, ) -> str: """Generates problem description based on given user prompt Args: prompt (str): initial user prompt + examples (Optional[List[Tuple[str, str]]]): dataset examples + task (Optional[Task]): task type + labels (Optional[List[str]]): unique class labels for + classification tasks; extract with _extract_labels(all_targets) + before calling Returns: str: generated problem description """ if examples: - request = PROBLEM_DESCRIPTION_BASED_ON_EXAMPLES_TEMPLATE.format( - prompt=prompt, examples=self._examples_to_str(examples) - ) + if task == Task.CLASSIFICATION and labels: + template = CLASSIFICATION_PROBLEM_DESCRIPTION_BASED_ON_EXAMPLES_TEMPLATE + request = template.format( + prompt=prompt, + examples=self._examples_to_str(examples), + labels=", ".join(labels), + ) + else: + request = PROBLEM_DESCRIPTION_BASED_ON_EXAMPLES_TEMPLATE.format( + prompt=prompt, + examples=self._examples_to_str(examples), + ) else: request = PROBLEM_DESCRIPTION_TEMPLATE.format(prompt=prompt) @@ -196,7 +226,7 @@ def generate( "Problem description was not provided, " + "so it will be generated automatically" ) - problem_description = self._generate_problem_description(prompt) + problem_description = self._generate_problem_description(prompt, task=task) logger.info(f"Generated problem description: {problem_description}") if task == Task.CLASSIFICATION: diff --git a/coolprompt/evaluator/evaluator.py b/coolprompt/evaluator/evaluator.py index 6c993f4..ce1514f 100644 --- a/coolprompt/evaluator/evaluator.py +++ b/coolprompt/evaluator/evaluator.py @@ -6,7 +6,6 @@ from langchain_core.language_models.base import BaseLanguageModel from langchain_core.messages.ai import AIMessage -import numpy as np from coolprompt.evaluator.metrics import BaseMetric from coolprompt.utils.logging_config import logger from coolprompt.utils.enums import Task @@ -129,8 +128,9 @@ def evaluate( detailed_failures = [] if failed_examples and failed_examples > 0: parsed_answers = [self.metric.parse_output(a) for a in answers] - indices = np.argsort(score_per_task)[:failed_examples] - for i in indices: + bad_indices = [i for i, s in enumerate(score_per_task) if s < 1.0] + bad_indices.sort(key=lambda i: score_per_task[i]) + for i in bad_indices[:failed_examples]: detailed_failures.append( FailedExampleDetailed( instance=dataset[i], diff --git a/coolprompt/evaluator/metrics.py b/coolprompt/evaluator/metrics.py index 7610ffb..5f562af 100644 --- a/coolprompt/evaluator/metrics.py +++ b/coolprompt/evaluator/metrics.py @@ -159,7 +159,9 @@ def _extract_bad_examples( List[float]: List of float metrics (for each model answer). """ - indices = np.argsort(results)[:failed_examples] + bad_indices = [i for i, r in enumerate(results) if r < 1.0] + bad_indices.sort(key=lambda i: results[i]) + indices = bad_indices[:failed_examples] return [ { diff --git a/coolprompt/method_evaluation/method_evaluation.py b/coolprompt/method_evaluation/method_evaluation.py index 8e6e55c..6c803ad 100644 --- a/coolprompt/method_evaluation/method_evaluation.py +++ b/coolprompt/method_evaluation/method_evaluation.py @@ -4,6 +4,7 @@ from langchain_core.language_models import BaseLanguageModel from coolprompt.optimizer.autoprompting_method import AutoPromptingMethod +from coolprompt.optimizer.brave import BRAVEMethod from coolprompt.optimizer.distill_prompt import DistillMethod from coolprompt.optimizer.hyper.meta_prompt import HyPERLightMethod from coolprompt.optimizer.hyper.hyper import HyPERMethod @@ -21,6 +22,7 @@ "compress": CompressorMethod, "regps": ReGPSMethod, "rider": RIDERGenesisMethod, + "brave": BRAVEMethod, } @@ -37,7 +39,8 @@ def evaluate_method( Args: method: One of ``hyper_light``, ``hyper``, ``reflective`` / ``reflectiveprompt``, - ``distill``, ``compress``, ``regps``, ``rider`` (same names as in + ``distill``, ``compress``, ``regps``, ``rider``, ``brave`` + (same names as in ``PromptTuner`` / ``validate_method`` where applicable). model: LangChain language model used for optimization and evaluation. config: Benchmark configuration dict or path to a YAML file. diff --git a/coolprompt/optimizer/brave/README.md b/coolprompt/optimizer/brave/README.md new file mode 100644 index 0000000..2d9c11e --- /dev/null +++ b/coolprompt/optimizer/brave/README.md @@ -0,0 +1,31 @@ +## BRAVE optimizer + +BRAVE is a data-driven evolutionary prompt optimizer. It uses a contextual +controller to choose prompt-transformation operators while respecting a token +budget, and selects the final prompt on a validation split. + +Use it through the main CoolPrompt API: + +```python +from coolprompt import PromptTuner + +tuner = PromptTuner(target_model=model) +prompt = tuner.run( + start_prompt="Classify the sentiment of the input.", + task="classification", + dataset=train_inputs, + target=train_labels, + method="brave", + problem_description="Binary sentiment classification.", + max_steps=20, + initial_budget_tokens=50_000, +) +``` + +BRAVE configuration fields can be passed directly to `PromptTuner.run`, or as +a `BRAVEConfig` instance through the `config` keyword. The low-level +`brave(...)` function, `BRAVEEvoluter`, `BRAVEConfig`, and YAML configuration +loader are exported from `coolprompt.optimizer.brave`. + +Set `log_dir` to persist operation logs. Without it, optimization runs without +writing BRAVE-specific log files. diff --git a/coolprompt/optimizer/brave/__init__.py b/coolprompt/optimizer/brave/__init__.py new file mode 100644 index 0000000..935525c --- /dev/null +++ b/coolprompt/optimizer/brave/__init__.py @@ -0,0 +1,14 @@ +from coolprompt.optimizer.brave.evoluter import BRAVEEvoluter +from coolprompt.optimizer.brave.run import BRAVEMethod, brave +from coolprompt.optimizer.brave.utils import ( + BRAVEConfig, + load_brave_config_from_yaml, +) + +__all__ = [ + "brave", + "BRAVEMethod", + "BRAVEEvoluter", + "BRAVEConfig", + "load_brave_config_from_yaml", +] diff --git a/coolprompt/optimizer/brave/actions.py b/coolprompt/optimizer/brave/actions.py new file mode 100644 index 0000000..48e2131 --- /dev/null +++ b/coolprompt/optimizer/brave/actions.py @@ -0,0 +1,45 @@ +from dataclasses import dataclass, field +from typing import Any, Dict, List, Protocol + +from coolprompt.optimizer.brave.core_states import OptimizerState + + +@dataclass +class ActionResult: + """Describe the outcome and token cost of an optimizer action.""" + + action: str + delta_quality: float + cost_tokens: float + payload: Dict[str, Any] = field(default_factory=dict) + improved: bool = False + + +class ActionExecutor(Protocol): + """Executor interface for domain-specific implementation. + + You can implement this against your existing GRAPE pipeline. + """ + + def execute( + self, + action: str, + population: List[str], + state: OptimizerState, + train_data: Any, + val_data: Any, + ) -> ActionResult: + """Execute an action against the current optimizer context. + + Args: + action (str): name of the action to execute. + population (List[str]): current prompt population. + state (OptimizerState): current normalized optimizer state. + train_data (Any): training data available to the action. + val_data (Any): validation data available to the action. + + Returns: + ActionResult: measured action outcome and its payload. + """ + + pass diff --git a/coolprompt/optimizer/brave/batch_sampler.py b/coolprompt/optimizer/brave/batch_sampler.py new file mode 100644 index 0000000..9ab9aa9 --- /dev/null +++ b/coolprompt/optimizer/brave/batch_sampler.py @@ -0,0 +1,389 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import Dict, List, Sequence, Tuple + +import numpy as np + +from coolprompt.utils.enums import Task + + +@dataclass(frozen=True) +class GenerationFeatures: + """Store source and target lengths used for generation stratification.""" + + input_len: int + target_len: int + + +class StratifiedBatchSampler: + """Builds balanced train batches for every optimization epoch.""" + + def __init__( + self, + task: Task, + batch_size: int, + seed: int = 19, + generation_bins: int = 3, + ) -> None: + """Configure deterministic stratified batch sampling. + + Args: + task (Task): optimization task that determines stratification. + batch_size (int): maximum number of sampled examples. + seed (int): base seed combined with the epoch number. + generation_bins (int): number of quantile bins per length feature. + """ + + self.task = task + self.batch_size = max(int(batch_size), 1) + self.generation_bins = max(int(generation_bins), 2) + self._seed = int(seed) + + def _build_generation_features( + self, + dataset: List[str], + targets: List[str], + ) -> List[GenerationFeatures]: + """Build length features for paired generation examples. + + Args: + dataset (List[str]): source texts. + targets (List[str]): target texts paired with the sources. + + Returns: + List[GenerationFeatures]: source and target lengths per example. + """ + + features: List[GenerationFeatures] = [] + for source, target in zip(dataset, targets): + features.append( + GenerationFeatures( + input_len=len(source), + target_len=len(target), + ) + ) + return features + + def _quantile_edges(self, values: np.ndarray) -> np.ndarray: + """Return unique internal quantile boundaries for an array. + + Args: + values (np.ndarray): numeric values to divide into bins. + + Returns: + np.ndarray: sorted unique internal quantile boundaries. + """ + + if values.size == 0: + return np.array([], dtype=np.float64) + quantiles = np.linspace(0.0, 1.0, self.generation_bins + 1)[1:-1] + if quantiles.size == 0: + return np.array([], dtype=np.float64) + edges = np.quantile(values, quantiles) + return np.unique(edges.astype(np.float64)) + + @staticmethod + def _assign_bin(value: float, edges: np.ndarray) -> int: + """Map a value to the bin delimited by ``edges``. + + Args: + value (float): value to assign. + edges (np.ndarray): sorted bin boundaries. + + Returns: + int: zero-based bin index. + """ + + if edges.size == 0: + return 0 + return int(np.searchsorted(edges, value, side="right")) + + def _build_strata( + self, + dataset: List[str], + targets: List[str | int], + ) -> Dict[Tuple[str, ...], List[int]]: + """Group examples by label or source and target length bins. + + Args: + dataset (List[str]): source examples. + targets (List[str | int]): labels or generation targets. + + Returns: + Dict[Tuple[str, ...], List[int]]: stratum keys mapped to indices. + """ + + strata: Dict[Tuple[str, ...], List[int]] = {} + if self.task == Task.CLASSIFICATION: + for idx, target in enumerate(targets): + key = (str(target),) + strata.setdefault(key, []).append(idx) + return strata + + features = self._build_generation_features(dataset, targets) + input_lengths = np.array([f.input_len for f in features], dtype=np.float64) + target_lengths = np.array([f.target_len for f in features], dtype=np.float64) + input_edges = self._quantile_edges(input_lengths) + target_edges = self._quantile_edges(target_lengths) + + for idx, feats in enumerate(features): + key = ( + str(self._assign_bin(feats.input_len, input_edges)), + str(self._assign_bin(feats.target_len, target_edges)), + ) + strata.setdefault(key, []).append(idx) + return strata + + def _compute_quotas( + self, + strata_sizes: Dict[Tuple[str, ...], int], + total_size: int, + ) -> Dict[Tuple[str, ...], int]: + """Allocate batch slots proportionally across strata. + + Args: + strata_sizes (Dict[Tuple[str, ...], int]): examples per stratum. + total_size (int): total number of available examples. + + Returns: + Dict[Tuple[str, ...], int]: number of slots per stratum. + """ + + if total_size <= 0: + return {} + target_size = min(self.batch_size, total_size) + expected = { + key: target_size * (size / total_size) for key, size in strata_sizes.items() + } + base = {key: int(np.floor(value)) for key, value in expected.items()} + assigned = sum(base.values()) + remainder = target_size - assigned + + if remainder > 0: + ranked = sorted( + expected.items(), key=lambda item: item[1] - base[item[0]], reverse=True + ) + for key, _ in ranked[:remainder]: + base[key] += 1 + return base + + def sample( + self, + dataset: Sequence[str], + targets: Sequence[str | int], + epoch: int, + ) -> List[int]: + """Sample deterministic stratified indices for an epoch. + + Args: + dataset (Sequence[str]): source examples. + targets (Sequence[str | int]): labels or generation targets. + epoch (int): epoch used to derive the random seed. + + Returns: + List[int]: selected dataset indices. + """ + + total_size = len(dataset) + if total_size == 0: + return [] + if total_size <= self.batch_size: + return list(range(total_size)) + + rng = np.random.default_rng(self._seed + int(epoch)) + strata = self._build_strata(dataset, targets) + strata_sizes = {k: len(v) for k, v in strata.items()} + quotas = self._compute_quotas(strata_sizes, total_size) + + selected: List[int] = [] + selected_set = set() + leftovers: List[int] = [] + for key, indices in strata.items(): + quota = quotas.get(key, 0) + if quota <= 0: + leftovers.extend(indices) + continue + shuffled = list(indices) + rng.shuffle(shuffled) + take = min(quota, len(shuffled)) + picked = shuffled[:take] + selected.extend(picked) + selected_set.update(picked) + leftovers.extend(shuffled[take:]) + + if len(selected) < self.batch_size: + remain = [idx for idx in leftovers if idx not in selected_set] + rng.shuffle(remain) + need = self.batch_size - len(selected) + selected.extend(remain[:need]) + + if len(selected) < self.batch_size: + need = self.batch_size - len(selected) + fallback = rng.choice(total_size, size=need, replace=True).tolist() + selected.extend(int(x) for x in fallback) + + rng.shuffle(selected) + return selected[: self.batch_size] + + +class CurriculumStratifiedBatchSampler(StratifiedBatchSampler): + """Stratified sampler that gradually up-weights hard examples. + + During warmup_steps: pure stratified sampling (alpha=0). + After warmup: within each stratum, example weights blend uniform + and difficulty-based as alpha linearly grows to max_alpha. + """ + + def __init__( + self, + task: Task, + batch_size: int, + total_steps: int, + seed: int = 19, + generation_bins: int = 3, + warmup_steps: int = 20, + max_alpha: float = 0.6, + ) -> None: + """Configure curriculum sampling and its difficulty schedule. + + Args: + task (Task): optimization task that determines stratification. + batch_size (int): maximum number of sampled examples. + total_steps (int): number of curriculum steps. + seed (int): base random seed. + generation_bins (int): number of quantile bins per length feature. + warmup_steps (int): steps that use uniform stratum sampling. + max_alpha (float): maximum weight assigned to difficulty. + """ + + super().__init__(task, batch_size, seed, generation_bins) + self.total_steps = max(int(total_steps), 1) + self.warmup_steps = max(int(warmup_steps), 0) + self.max_alpha = float(np.clip(max_alpha, 0.0, 1.0)) + self._error_counts: Dict[int, int] = {} + self._eval_counts: Dict[int, int] = {} + + def update_difficulties( + self, + batch_indices: List[int], + failed_indices: List[int], + ) -> None: + """Record per-example outcomes after an evaluation step. + + Args: + batch_indices (List[int]): global indices included in the batch. + failed_indices (List[int]): batch indices where evaluation failed. + """ + failed_set = set(failed_indices) + for idx in batch_indices: + self._eval_counts[idx] = self._eval_counts.get(idx, 0) + 1 + if idx in failed_set: + self._error_counts[idx] = self._error_counts.get(idx, 0) + 1 + + def _curriculum_alpha(self, epoch: int) -> float: + """Return the difficulty weight for the requested epoch. + + Args: + epoch (int): current curriculum epoch. + + Returns: + float: difficulty mixture weight in ``[0, max_alpha]``. + """ + + if epoch <= self.warmup_steps: + return 0.0 + ramp_len = max(self.total_steps - self.warmup_steps, 1) + return self.max_alpha * min((epoch - self.warmup_steps) / ramp_len, 1.0) + + def _difficulty(self, idx: int) -> float: + """Return the observed difficulty of an example. + + Args: + idx (int): dataset index. + + Returns: + float: empirical error rate, or ``0.5`` when unseen. + """ + + evals = self._eval_counts.get(idx, 0) + if evals == 0: + return 0.5 # neutral prior for unseen examples + return self._error_counts.get(idx, 0) / evals + + def sample( + self, + dataset: Sequence[str], + targets: Sequence[str | int], + epoch: int, + ) -> List[int]: + """Sample stratified indices with curriculum difficulty weighting. + + Args: + dataset (Sequence[str]): source examples. + targets (Sequence[str | int]): labels or generation targets. + epoch (int): current curriculum epoch. + + Returns: + List[int]: selected dataset indices. + """ + + alpha = self._curriculum_alpha(epoch) + if alpha == 0.0: + return super().sample(dataset, targets, epoch) + + total_size = len(dataset) + if total_size == 0: + return [] + if total_size <= self.batch_size: + return list(range(total_size)) + + rng = np.random.default_rng(self._seed + int(epoch)) + strata = self._build_strata(dataset, targets) + strata_sizes = {k: len(v) for k, v in strata.items()} + quotas = self._compute_quotas(strata_sizes, total_size) + + selected: List[int] = [] + leftovers: List[int] = [] + + for key, indices in strata.items(): + quota = quotas.get(key, 0) + if quota <= 0: + leftovers.extend(indices) + continue + + take = min(quota, len(indices)) + if take >= len(indices): + selected.extend(indices) + continue + + # Within-stratum blend: uniform + difficulty + raw = np.array( + [(1.0 - alpha) + alpha * self._difficulty(i) for i in indices], + dtype=np.float64, + ) + raw = np.clip(raw, 1e-12, None) + p = raw / raw.sum() + + picked_local = rng.choice( + len(indices), size=take, replace=False, p=p + ).tolist() + picked_set = set(picked_local) + selected.extend(indices[i] for i in picked_local) + leftovers.extend( + indices[i] for i in range(len(indices)) if i not in picked_set + ) + + if len(selected) < self.batch_size: + selected_set = set(selected) + remain = [idx for idx in leftovers if idx not in selected_set] + rng.shuffle(remain) + selected.extend(remain[: self.batch_size - len(selected)]) + + if len(selected) < self.batch_size: + need = self.batch_size - len(selected) + fallback = rng.choice(total_size, size=need, replace=True).tolist() + selected.extend(int(x) for x in fallback) + + rng.shuffle(selected) + return selected[: self.batch_size] diff --git a/coolprompt/optimizer/brave/bayesian_sampling.py b/coolprompt/optimizer/brave/bayesian_sampling.py new file mode 100644 index 0000000..6e6b253 --- /dev/null +++ b/coolprompt/optimizer/brave/bayesian_sampling.py @@ -0,0 +1,311 @@ +import math +from typing import Dict, List + +import numpy as np + +from coolprompt.optimizer.brave.core_states import OptimizerState + + +class StateFeaturizer: + """Maps OptimizerState to a dense feature vector.""" + + def transform(self, s: OptimizerState) -> np.ndarray: + """Convert an optimizer state into contextual bandit features. + + Args: + s (OptimizerState): normalized optimizer state. + + Returns: + np.ndarray: dense feature vector for the controller. + """ + + x = np.array( + [ + s.val_quality, + s.quality_slope, + s.stagnation, + s.useless_ops_ratio, + s.remaining_budget_ratio, + s.epoch_progress, + # progress under budget pressure + s.stagnation * s.remaining_budget_ratio, + s.population_diversity, + # stagnation is most dangerous when population has converged + s.stagnation * (1.0 - s.population_diversity), + ], + dtype=np.float64, + ) + return x + + @property + def dim(self) -> int: + """Return the number of features emitted by :meth:`transform`. + + Returns: + int: feature-vector dimensionality. + """ + + return 9 + + +class BayesianLinearTS: + """Bayesian linear regression with Thompson sampling.""" + + def __init__(self, dim: int, alpha: float = 1.0, sigma2: float = 1.0) -> None: + """Initialize the precision matrix and response vector. + + Args: + dim (int): number of regression features. + alpha (float): isotropic prior precision. + sigma2 (float): observation-noise variance. + """ + + self.dim = dim + self.alpha = alpha + self.sigma2 = sigma2 + self.A = alpha * np.eye(dim) + self.b = np.zeros(dim) + + def update(self, x: np.ndarray, y: float) -> None: + """Update posterior statistics with one observation. + + Args: + x (np.ndarray): observation feature vector. + y (float): observed scalar response. + """ + + self.A += np.outer(x, x) / self.sigma2 + self.b += (x * y) / self.sigma2 + + def sample_theta(self, rng: np.random.Generator) -> np.ndarray: + """Draw regression coefficients from the current posterior. + + Args: + rng (np.random.Generator): random-number generator to use. + + Returns: + np.ndarray: sampled coefficient vector. + """ + + # Numerical guard + A_inv = np.linalg.pinv(self.A) + mu = A_inv @ self.b + cov = self.sigma2 * A_inv + return rng.multivariate_normal(mu, cov) + + def posterior_mean(self) -> np.ndarray: + """Return the posterior mean of the regression coefficients. + + Returns: + np.ndarray: posterior coefficient mean. + """ + + A_inv = np.linalg.pinv(self.A) + return A_inv @ self.b + + def predictive_mean(self, x: np.ndarray) -> float: + """Predict the expected response for a feature vector. + + Args: + x (np.ndarray): feature vector. + + Returns: + float: posterior predictive mean. + """ + + return float(np.dot(self.posterior_mean(), x)) + + def predictive_std(self, x: np.ndarray) -> float: + """Estimate posterior predictive uncertainty for a feature vector. + + Args: + x (np.ndarray): feature vector. + + Returns: + float: posterior predictive standard deviation. + """ + + A_inv = np.linalg.pinv(self.A) + var = float(np.dot(x, A_inv @ x)) * self.sigma2 + return float(math.sqrt(max(var, 1e-12))) + + +class OnlineActionMLP: + """Tiny shared-trunk neural contextual bandit (numpy-only). + + Heads predict: + - benefit + - cost + - improvement logit (for P(improvement > 0)) + """ + + def __init__( + self, + actions: List[str], + input_dim: int, + hidden_dim: int = 32, + learning_rate: float = 5e-3, + seed: int = 123, + ) -> None: + """Initialize the shared trunk and action-specific heads. + + Args: + actions (List[str]): supported action names. + input_dim (int): input feature dimensionality. + hidden_dim (int): width of the shared hidden layer. + learning_rate (float): stochastic-gradient learning rate. + seed (int): parameter-initialization seed. + """ + + self.actions = actions + self.a2i = {a: i for i, a in enumerate(actions)} + self.input_dim = input_dim + self.hidden_dim = hidden_dim + self.lr = learning_rate + self.rng = np.random.default_rng(seed) + + # Shared trunk + self.W1 = self.rng.normal(0.0, 0.1, size=(input_dim, hidden_dim)) + self.b1 = np.zeros(hidden_dim) + + # Per-action heads + n = len(actions) + self.W_benefit = self.rng.normal(0.0, 0.1, size=(n, hidden_dim)) + self.b_benefit = np.zeros(n) + self.W_cost = self.rng.normal(0.0, 0.1, size=(n, hidden_dim)) + self.b_cost = np.zeros(n) + self.W_impr = self.rng.normal(0.0, 0.1, size=(n, hidden_dim)) + self.b_impr = np.zeros(n) + + @staticmethod + def _relu(z: np.ndarray) -> np.ndarray: + """Apply the rectified linear activation element-wise. + + Args: + z (np.ndarray): activation input. + + Returns: + np.ndarray: rectified values. + """ + + return np.maximum(z, 0.0) + + @staticmethod + def _sigmoid(z: float) -> float: + """Compute a numerically bounded logistic sigmoid. + + Args: + z (float): sigmoid logit. + + Returns: + float: probability in ``(0, 1)``. + """ + + z = float(np.clip(z, -20.0, 20.0)) + return 1.0 / (1.0 + math.exp(-z)) + + def _forward_hidden(self, x: np.ndarray) -> np.ndarray: + """Project an input through the shared hidden layer. + + Args: + x (np.ndarray): input feature vector. + + Returns: + np.ndarray: hidden representation. + """ + + return self._relu(x @ self.W1 + self.b1) + + def predict(self, action: str, x: np.ndarray) -> Dict[str, float]: + """Predict benefit, cost, and improvement probability for an action. + + Args: + action (str): action whose heads should be evaluated. + x (np.ndarray): contextual feature vector. + + Returns: + Dict[str, float]: benefit, positive cost, and improvement + probability. + """ + + idx = self.a2i[action] + h = self._forward_hidden(x) + benefit = float(np.dot(self.W_benefit[idx], h) + self.b_benefit[idx]) + # ensure positive-ish cost + raw_cost = float(np.dot(self.W_cost[idx], h) + self.b_cost[idx]) + cost = float(np.log1p(math.exp(np.clip(raw_cost, -20.0, 20.0))) + 1e-6) + impr_logit = float(np.dot(self.W_impr[idx], h) + self.b_impr[idx]) + impr_prob = self._sigmoid(impr_logit) + return {"benefit": benefit, "cost": cost, "impr_prob": impr_prob} + + def update( + self, + action: str, + x: np.ndarray, + target_benefit: float, + target_cost: float, + target_impr: float, + ) -> None: + """Train the selected action heads and shared trunk with one sample. + + Args: + action (str): action whose heads should be updated. + x (np.ndarray): contextual feature vector. + target_benefit (float): observed quality benefit. + target_cost (float): observed token cost. + target_impr (float): binary improvement target. + """ + + idx = self.a2i[action] + h_pre = x @ self.W1 + self.b1 + h = self._relu(h_pre) + + # forward + pred_b = float(np.dot(self.W_benefit[idx], h) + self.b_benefit[idx]) + raw_cost = float(np.dot(self.W_cost[idx], h) + self.b_cost[idx]) + pred_c = float(np.log1p(math.exp(np.clip(raw_cost, -20.0, 20.0))) + 1e-6) + pred_l = float(np.dot(self.W_impr[idx], h) + self.b_impr[idx]) + pred_p = self._sigmoid(pred_l) + + # losses: + # benefit, cost -> mse + # improvement -> logistic BCE + db = pred_b - float(target_benefit) + dc = pred_c - float(target_cost) + dp = pred_p - float(target_impr) + + # gradients wrt head outputs + # cost head uses softplus(raw_cost), + # d pred_c / d raw_cost = sigmoid(raw_cost) + dsoftplus = self._sigmoid(raw_cost) + draw_cost = dc * dsoftplus + + # head grads + gWb = db * h + gbb = db + gWc = draw_cost * h + gbc = draw_cost + gWi = dp * h + gbi = dp + + # backprop to hidden + gh = ( + db * self.W_benefit[idx] + + draw_cost * self.W_cost[idx] + + dp * self.W_impr[idx] + ) + gh = gh * (h_pre > 0.0).astype(float) + + # trunk grads + gW1 = np.outer(x, gh) + gb1 = gh + + # SGD updates (only selected action heads) + self.W_benefit[idx] -= self.lr * gWb + self.b_benefit[idx] -= self.lr * gbb + self.W_cost[idx] -= self.lr * gWc + self.b_cost[idx] -= self.lr * gbc + self.W_impr[idx] -= self.lr * gWi + self.b_impr[idx] -= self.lr * gbi + self.W1 -= self.lr * gW1 + self.b1 -= self.lr * gb1 diff --git a/coolprompt/optimizer/brave/controller.py b/coolprompt/optimizer/brave/controller.py new file mode 100644 index 0000000..e82f64a --- /dev/null +++ b/coolprompt/optimizer/brave/controller.py @@ -0,0 +1,349 @@ +import math +from typing import Any, Dict, List, Optional, Tuple + +import numpy as np + +from coolprompt.optimizer.brave.bayesian_sampling import ( + BayesianLinearTS, + OnlineActionMLP, +) + + +class EVCController: + """Action selector maximizing expected epistemic value per token. + + Score(action) = Benefit / Cost + Benefit ~ ΔQuality + λ_c * ΔCoverage - λ_d * ΔDrift + """ + + def __init__( + self, + actions: List[str], + feature_dim: int, + min_cost_eps: float = 1e-6, + max_action_budget_share: float = 0.35, + uncertainty_penalty_beta: float = 0.35, + neural_weight: float = 0.5, + alpha_roi_ema: float = 0.1, + improve_prob_weight: float = 0.6, + kill_switch_min_trials: int = 10, + kill_switch_roi_threshold: float = -0.0002, + kill_switch_base_cooldown: int = 5, + kill_switch_scaling_factor: float = 20.0, + use_neural_bandit: bool = True, + neural_hidden_dim: int = 32, + neural_learning_rate: float = 5e-3, + seed: int = 42, + ) -> None: + """Initialize action models, constraints, and kill-switch state. + + Args: + actions (List[str]): action names available to the controller. + feature_dim (int): contextual feature-vector dimensionality. + min_cost_eps (float): lower bound used in cost divisions. + max_action_budget_share (float): maximum remaining-budget fraction + that a single action may consume. + uncertainty_penalty_beta (float): strength of the cost uncertainty + penalty. + neural_weight (float): weight of neural versus linear predictions. + alpha_roi_ema (float): smoothing factor for realized ROI. + improve_prob_weight (float): influence of improvement probability. + kill_switch_min_trials (int): observations required before + disabling an action. + kill_switch_roi_threshold (float): ROI threshold for disabling. + kill_switch_base_cooldown (int): minimum disabled duration. + kill_switch_scaling_factor (float): usefulness-based cooldown + scale. + use_neural_bandit (bool): whether to blend neural predictions. + neural_hidden_dim (int): neural bandit hidden-layer width. + neural_learning_rate (float): neural bandit learning rate. + seed (int): Thompson-sampling and neural initialization seed. + """ + + self.actions = actions + self.min_cost_eps = min_cost_eps + self.max_action_budget_share = max_action_budget_share + self.uncertainty_penalty_beta = uncertainty_penalty_beta + self.neural_weight = min(max(neural_weight, 0.0), 1.0) + self.alpha_roi_ema = alpha_roi_ema + self.improve_prob_weight = improve_prob_weight + self.kill_switch_min_trials = kill_switch_min_trials + self.kill_switch_roi_threshold = kill_switch_roi_threshold + self.kill_switch_base_cooldown = kill_switch_base_cooldown + self.kill_switch_scaling_factor = kill_switch_scaling_factor + self.use_neural_bandit = use_neural_bandit + self.rng = np.random.default_rng(seed) + + self.benefit_models: Dict[str, BayesianLinearTS] = { + a: BayesianLinearTS(feature_dim, alpha=1.0, sigma2=1.0) for a in actions + } + self.cost_models: Dict[str, BayesianLinearTS] = { + a: BayesianLinearTS(feature_dim, alpha=1.0, sigma2=1.0) for a in actions + } + self.improvement_models: Dict[str, BayesianLinearTS] = { + a: BayesianLinearTS(feature_dim, alpha=1.0, sigma2=1.0) for a in actions + } + self.neural_bandit = ( + OnlineActionMLP( + actions=actions, + input_dim=feature_dim, + hidden_dim=neural_hidden_dim, + learning_rate=neural_learning_rate, + seed=seed + 17, + ) + if use_neural_bandit + else None + ) + self.action_stats: Dict[str, Dict[str, float]] = { + a: { + "trials": 0.0, + "success_count": 0.0, + "ema_roi": 0.0, + "last_selected_step": -1.0, + "disabled_until_step": -1.0, + } + for a in actions + } + self.global_step: int = 0 + + def _sample_benefit_cost(self, action: str, x: np.ndarray) -> Tuple[float, float]: + """Sample an action's benefit and cost from its posteriors. + + Args: + action (str): action to sample. + x (np.ndarray): contextual feature vector. + + Returns: + Tuple[float, float]: sampled benefit and positive cost. + """ + + theta_b = self.benefit_models[action].sample_theta(self.rng) + theta_c = self.cost_models[action].sample_theta(self.rng) + + benefit = float(np.dot(theta_b, x)) + cost = float(np.dot(theta_c, x)) + cost = max(cost, self.min_cost_eps) + return benefit, cost + + @staticmethod + def _sigmoid(z: float) -> float: + """Compute a numerically bounded logistic sigmoid. + + Args: + z (float): sigmoid logit. + + Returns: + float: probability in ``(0, 1)``. + """ + + z = float(np.clip(z, -20.0, 20.0)) + return 1.0 / (1.0 + math.exp(-z)) + + def _calculate_usefulness(self, action: str) -> float: + """Calculate an action's empirical usefulness. + + Args: + action (str): action whose successes should be inspected. + + Returns: + float: success rate in ``[0, 1]``; ``1`` before any trials. + """ + st = self.action_stats[action] + trials = st["trials"] + + if trials < 1: + return 1.0 # No data yet, assume useful + + success_rate = st["success_count"] / trials + return float(success_rate) + + def _calculate_adaptive_cooldown(self, action: str) -> int: + """Calculate cooldown duration from operator usefulness. + + Args: + action (str): action to place on cooldown. + + Returns: + int: number of controller steps to disable the action. + """ + usefulness = self._calculate_usefulness(action) + lack_of_usefulness = 1.0 - usefulness + if lack_of_usefulness < 0.5: + lack_of_usefulness = self._sigmoid(15 * (lack_of_usefulness - 0.5)) + else: + lack_of_usefulness = self._sigmoid(4 * (lack_of_usefulness - 0.5)) + + cooldown = int( + self.kill_switch_base_cooldown + + self.kill_switch_scaling_factor * lack_of_usefulness + ) + + return cooldown + + def _should_disable(self, action: str) -> bool: + """Return whether an action is blocked by its kill switch. + + Args: + action (str): action to inspect. + + Returns: + bool: whether the action should be skipped at the current step. + """ + + st = self.action_stats[action] + if self.global_step < int(st["disabled_until_step"]): + return True + + if st["trials"] < self.kill_switch_min_trials: + return False + # Check if just exiting disability period + # (give second chance by resetting ema_roi) + if st["disabled_until_step"] > 0 and self.global_step == int( + st["disabled_until_step"] + ): + st["ema_roi"] = 0.0 # Reset to give fresh evaluation + st["disabled_until_step"] = -1.0 + return False + + if st["ema_roi"] < self.kill_switch_roi_threshold: + st["disabled_until_step"] = float( + self.global_step + self._calculate_adaptive_cooldown(action) + ) + return True + return False + + def select_action( + self, + x: np.ndarray, + remaining_budget_tokens: float, + candidate_actions: Optional[List[str]] = None, + ) -> Tuple[str, Dict[str, float]]: + """Select the eligible action with the highest value per cost. + + Args: + x (np.ndarray): contextual feature vector. + remaining_budget_tokens (float): tokens available for future + actions. + candidate_actions (Optional[List[str]]): optional action subset. + + Returns: + Tuple[str, Dict[str, float]]: selected action and scores by action. + Returns ``(None, None)`` when no action is eligible. + """ + + self.global_step += 1 + best_action: Optional[str] = None + best_score = -math.inf + scores: Dict[str, float] = {} + if candidate_actions is not None: + action_pool = candidate_actions + else: + action_pool = self.actions + + for action in action_pool: + if self._should_disable(action): + continue + + benefit_hat, cost_hat = self._sample_benefit_cost(action, x) + impr_hat_linear = self._sigmoid( + self.improvement_models[action].predictive_mean(x) + ) + b_unc = self.benefit_models[action].predictive_std(x) + c_unc = self.cost_models[action].predictive_std(x) + uncertainty = 0.5 * (b_unc + c_unc) + + if self.neural_bandit is not None: + nn = self.neural_bandit.predict(action, x) + benefit_hat = ( + 1.0 - self.neural_weight + ) * benefit_hat + self.neural_weight * nn["benefit"] + cost_hat = ( + 1.0 - self.neural_weight + ) * cost_hat + self.neural_weight * nn["cost"] + impr_prob = ( + 1.0 - self.neural_weight + ) * impr_hat_linear + self.neural_weight * nn["impr_prob"] + else: + impr_prob = impr_hat_linear + + if cost_hat > remaining_budget_tokens: + continue + if cost_hat > self.max_action_budget_share * max( + remaining_budget_tokens, 1.0 + ): + continue + + effective_cost = cost_hat * ( + 1.0 + self.uncertainty_penalty_beta * uncertainty + ) + score = benefit_hat * (1.0 + self.improve_prob_weight * impr_prob) + score = score / max(effective_cost, self.min_cost_eps) + scores[action] = score + if score > best_score: + best_score = score + best_action = action + + if best_action is None: + return None, None + + self.action_stats[best_action]["last_selected_step"] = float(self.global_step) + return best_action, scores + + def update( + self, + action: str, + x_before: np.ndarray, + delta_quality: float, + actual_cost_tokens: float, + improved: bool = False, + ) -> None: + """Update action models and statistics from an observed outcome. + + Args: + action (str): action that produced the outcome. + x_before (np.ndarray): context observed before execution. + delta_quality (float): measured quality change. + actual_cost_tokens (float): measured token cost. + improved (bool): whether the operation was considered useful. + """ + + benefit = float(delta_quality) + cost = max(float(actual_cost_tokens), self.min_cost_eps) + improvement = 1.0 if float(delta_quality) > 0.0 else 0.0 + + self.benefit_models[action].update(x_before, benefit) + self.cost_models[action].update(x_before, cost) + self.improvement_models[action].update(x_before, improvement) + if self.neural_bandit is not None: + self.neural_bandit.update( + action=action, + x=x_before, + target_benefit=benefit, + target_cost=cost, + target_impr=improvement, + ) + + st = self.action_stats[action] + st["trials"] += 1.0 + if improved: + st["success_count"] += 1.0 + realized_roi = benefit / max(cost, self.min_cost_eps) + # EMA for kill-switch + st["ema_roi"] = (1.0 - self.alpha_roi_ema) * st[ + "ema_roi" + ] + self.alpha_roi_ema * realized_roi + + def diagnostics(self) -> Dict[str, Any]: + """Return a serializable snapshot of controller state. + + Returns: + Dict[str, Any]: step, action statistics, and model settings. + """ + + return { + "global_step": self.global_step, + "action_stats": self.action_stats, + "uncertainty_penalty_beta": self.uncertainty_penalty_beta, + "neural_weight": self.neural_weight, + "use_neural_bandit": self.use_neural_bandit, + } diff --git a/coolprompt/optimizer/brave/core_states.py b/coolprompt/optimizer/brave/core_states.py new file mode 100644 index 0000000..2366f03 --- /dev/null +++ b/coolprompt/optimizer/brave/core_states.py @@ -0,0 +1,29 @@ +from dataclasses import dataclass, field +from typing import List + + +@dataclass +class OptimizerState: + """Compact normalized state for the controller. + + All scalar fields should be in [0, 1] whenever possible. + """ + + val_quality: float = 0.0 + quality_slope: float = 0.0 + stagnation: float = 0.0 + useless_ops_ratio: float = 0.0 + remaining_budget_ratio: float = 1.0 + epoch_progress: float = 0.0 + population_diversity: float = 0.5 + + +@dataclass +class ReflectionRecord: + """Record a reflection together with its score and creation step.""" + + text: str + source_action: str + utility_score: float = 0.0 + contradiction_score: float = 0.0 + tags: List[str] = field(default_factory=list) diff --git a/coolprompt/optimizer/brave/evoluter.py b/coolprompt/optimizer/brave/evoluter.py new file mode 100644 index 0000000..43d61e9 --- /dev/null +++ b/coolprompt/optimizer/brave/evoluter.py @@ -0,0 +1,1212 @@ +from __future__ import annotations + +import json +import numpy as np +from pathlib import Path +import random +from time import sleep +from typing import Dict, List, Mapping, Optional, Any, Tuple +from langchain_core.language_models.base import BaseLanguageModel +from langchain_core.messages.ai import AIMessage + +from coolprompt.evaluator import Evaluator +from coolprompt.language_model import create_chat_model +from coolprompt.optimizer.brave.actions import ActionResult +from coolprompt.optimizer.brave.batch_sampler import ( + StratifiedBatchSampler, + CurriculumStratifiedBatchSampler, +) +from coolprompt.optimizer.brave.bayesian_sampling import StateFeaturizer +from coolprompt.optimizer.brave.controller import EVCController +from coolprompt.optimizer.brave.core_states import OptimizerState +from coolprompt.optimizer.brave.operation_logger import OperationLogger +from coolprompt.optimizer.brave.operators import ( + Operator, + PopulationInitializationOperator, + ParaphraseInitializationOperator, + CrossoverOperator, + ElitistMutationOperator, + CompressorOperator, + GradientStepOperator, + HypeOperator, + LongTermMutationOperator, + ParaphrasingByPDOperator, + ZeroOrderMutationOperator, + CreativeRoleAndStyleMutationOperator, + CreativeZeroOrderMutationOperator, + HardFewShotExamplesOperator, + BiggerPopulationInitializationOperator, +) +from coolprompt.optimizer.brave.population_diversity import PopulationDiversityManager +from coolprompt.optimizer.brave.utils import ( + BRAVEConfig, + OptimizationLog, + reranking_population, +) +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.utils.utils import get_dataset_split + + +class BRAVEEvoluter: + """BRAVE (skeleton implementation). + + This class wires together: + - EVC controller + - drift-aware memory + - budget-constrained loop + """ + + ACTIONS = { + "crossover", + "elitist_mutation", + "compression", + "gradient_step", + "hype", + "long_term_mutation", + "paraphrase", + "zero_order", + "creative_role_and_style", + "creative_zero_order", + "few_shot_mutation", + } + + def __init__( + self, + model: BaseLanguageModel, + evaluator: Evaluator, + config: Optional[BRAVEConfig] = None, + seed: int = 19, + verbose: bool = True, + log_dir: Optional[str] = None, + ) -> None: + """Initialize BRAVE's controller, operators, and runtime state. + + Args: + model (BaseLanguageModel): language model used for prompt + generation. + evaluator (Evaluator): evaluator used to score prompts. + config (Optional[BRAVEConfig]): optimizer configuration. + seed (int): random seed for sampling and controller models. + verbose (bool): whether operation logging is enabled. + log_dir (Optional[str]): directory for operation logs. + + Raises: + ValueError: if the configured population initializer is + unsupported. + """ + + self.log_dir = log_dir + self.verbose = verbose + self.logger = None + if self.verbose and self.log_dir: + self.logger = OperationLogger(log_dir=log_dir) + + self.model = create_chat_model(model) + self.model.reset_stats() + + self.evaluator = evaluator + self.cfg = config or BRAVEConfig() + self.featurizer = StateFeaturizer() + self.cost_tracker_stats = {} + self.seed = seed + self.train_batch_sampler: Optional[StratifiedBatchSampler] = None + self.current_train_batch_data: Optional[List[str]] = None + self.current_train_batch_targets: Optional[List[Any]] = None + + self.actions = self.cfg.actions + if self.cfg.actions == "all": + self.actions = list(self.ACTIONS) + + self.controller = EVCController( + actions=self.actions, + max_action_budget_share=self.cfg.max_action_budget_share, + alpha_roi_ema=self.cfg.alpha_roi_ema, + feature_dim=self.featurizer.dim, + uncertainty_penalty_beta=self.cfg.uncertainty_penalty_beta, + neural_weight=self.cfg.neural_weight, + improve_prob_weight=self.cfg.improve_prob_weight, + kill_switch_min_trials=self.cfg.kill_switch_min_trials, + kill_switch_roi_threshold=self.cfg.kill_switch_roi_threshold, + kill_switch_base_cooldown=self.cfg.kill_switch_base_cooldown, + kill_switch_scaling_factor=self.cfg.kill_switch_scaling_factor, + use_neural_bandit=self.cfg.use_neural_bandit, + neural_hidden_dim=self.cfg.neural_hidden_dim, + neural_learning_rate=self.cfg.neural_learning_rate, + seed=seed, + ) + + self.diversity_manager = PopulationDiversityManager( + similarity_threshold=self.cfg.diversity_similarity_threshold, + max_per_cluster=self.cfg.diversity_max_per_cluster, + auto_threshold=self.cfg.diversity_auto_threshold, + target_cluster_count=self.cfg.population_size, + use_hierarchical=self.cfg.diversity_use_hierarchical, + use_bert=self.cfg.diversity_use_bert, + bert_weight=self.cfg.diversity_bert_weight, + duplicate_threshold=self.cfg.diversity_duplicate_threshold, + ) + + self.long_term_reflection = "" + self.short_term_reflections = [] + + if self.cfg.population_initializer == "paraphrase": + self.population_initializer = ParaphraseInitializationOperator( + logger=self.logger + ) + elif self.cfg.population_initializer == "brave": + self.population_initializer = PopulationInitializationOperator( + logger=self.logger + ) + elif self.cfg.population_initializer == "bigger": + self.population_initializer = BiggerPopulationInitializationOperator( + logger=self.logger + ) + else: + raise ValueError( + "Unsupported population initializer: " + self.cfg.population_initializer + ) + + self.crossover_operator = CrossoverOperator(self.logger) + self.elitist_mutation_operator = ElitistMutationOperator(self.logger) + self.compressor_operator = CompressorOperator( + model=self.model, logger=self.logger + ) + self.gradient_step_operator = GradientStepOperator(self.logger) + self.hype_operator = HypeOperator(model=self.model, logger=self.logger) + self.long_term_mutation_operator = LongTermMutationOperator(self.logger) + self.paraphrasing_operator = ParaphrasingByPDOperator(self.logger) + self.zero_order_operator = ZeroOrderMutationOperator(self.logger) + self.creative_role_and_style_operator = CreativeRoleAndStyleMutationOperator( + self.logger + ) + self.creative_zero_order_mutation_operator = CreativeZeroOrderMutationOperator( + self.logger + ) + + self.population: List[Prompt] = [] + self.logs: List[OptimizationLog] = [] + self.initial_budget: float = self.cfg.initial_budget_tokens + self.artifacts: Dict[str, bool] = self._init_artifacts() + + def _calculate_costs(self, new_tracker_stats: Dict[str, float]) -> Dict[str, float]: + """Calculate tracker-stat deltas and store cumulative values. + + Args: + new_tracker_stats (Dict[str, float]): latest cumulative model + statistics. + + Returns: + Dict[str, float]: change in every supplied statistic. + """ + + delta = { + metric: new_stat - self.cost_tracker_stats.get(metric, 0.0) + for metric, new_stat in new_tracker_stats.items() + } + self.cost_tracker_stats = new_tracker_stats + return delta + + @staticmethod + def _init_artifacts() -> Dict[str, bool]: + """Create the initial action-artifact availability flags. + + Returns: + Dict[str, bool]: artifact names mapped to initial availability. + """ + + return { + "has_eval": False, + "has_failures": False, + "has_gradients": False, + "has_short_term": False, + "has_offspring": False, + "has_memory_update": False, + "has_best_prompt": True, + } + + def _update_artifacts( + self, action: str, result: ActionResult, improved: bool + ) -> None: + """Update action-artifact flags from an execution result. + + Args: + action (str): action that was executed. + result (ActionResult): action outcome and payload. + improved (bool): whether the action improved the population. + """ + + if action == "crossover": + self.artifacts["has_offspring"] = True + if action == "mutation": + self.artifacts["has_offspring"] = True + if improved: + self.artifacts["has_best_prompt"] = True + + payload_artifacts = result.payload.get("artifacts") + if isinstance(payload_artifacts, Mapping): + for k, v in payload_artifacts.items(): + if k in self.artifacts: + self.artifacts[k] = bool(v) + + def _compute_state( + self, + best_quality: float, + recent_quality: List[float], + useless_ops_count: int, + steps_done: int, + remaining_budget: float, + population_diversity: float = 0.5, + ) -> OptimizerState: + """Build controller state from quality, progress, and budget data. + + Args: + best_quality (float): best score observed so far. + recent_quality (List[float]): recent best-score history. + useless_ops_count (int): number of recent non-improving operations. + steps_done (int): completed optimization steps. + remaining_budget (float): unspent token budget. + population_diversity (float): current diversity estimate. + + Returns: + OptimizerState: normalized contextual state for action selection. + """ + + slope = 0.0 + if len(recent_quality) >= 2: + slope = max(min(recent_quality[-1] - recent_quality[-2], 1.0), -1.0) + progress = min(steps_done / max(self.cfg.max_steps, 1), 1.0) + useless_ratio = useless_ops_count / max(steps_done, 1) + stagnation = 1.0 if abs(slope) < self.cfg.min_improvement else 0.0 + + return OptimizerState( + val_quality=float(np.clip(best_quality, 0.0, 1.0)), + quality_slope=float(np.clip((slope + 1.0) / 2.0, 0.0, 1.0)), + stagnation=float(stagnation), + useless_ops_ratio=float(np.clip(useless_ratio, 0.0, 1.0)), + remaining_budget_ratio=float( + np.clip(remaining_budget / self.cfg.initial_budget_tokens, 0.0, 1.0) + ), + epoch_progress=float(progress), + population_diversity=float(np.clip(population_diversity, 0.0, 1.0)), + ) + + def _update_elitist(self, new_prompt: Prompt) -> float: + """Update the elite prompt and compute population-aware gain. + + Args: + new_prompt (Prompt): evaluated candidate prompt. + + Returns: + float: quality gain adjusted by population mean and minimum. + """ + + mean_score = float(np.mean([p.score for p in self.population])) + min_score = float(np.min([p.score for p in self.population])) + delta_mean = new_prompt.score - mean_score + delta_min = new_prompt.score - min_score + delta = ( + new_prompt.score + - self.best_quality + + self.cfg.lambda_mean_quality * delta_mean + + self.cfg.lambda_min_quality * delta_min + ) + if new_prompt.score > self.best_quality: + self.elitist = new_prompt + self.best_quality = new_prompt.score + elif ( + new_prompt.score == self.best_quality + and self.elitist.origin == PromptOrigin.MANUAL + ): + self.elitist = new_prompt + return delta + + def _update_val_elitist(self, new_prompt: Prompt) -> None: + """Promote a prompt when it improves the best validation score. + + Args: + new_prompt (Prompt): candidate prompt to validate. + """ + + vs = self._evaluate_val_cached(new_prompt) + if vs > self.best_val_quality: + self.best_val_quality = vs + self.val_elitist = new_prompt + + def _rescore_population(self) -> None: + """Re-evaluate and rank the current population on training data.""" + + for prompt in self.population: + self._evaluate(prompt, split="train") + self.population = reranking_population(self.population) + self.elitist = self.population[0] + self.best_quality = self.elitist.score + + def _add_new_prompt(self, prompt: Prompt, step: int) -> bool: + """Insert a prompt and enforce population limits. + + Args: + prompt (Prompt): evaluated prompt to add. + step (int): current optimization step used for logging. + + Returns: + bool: whether the prompt remains in the bounded population. + """ + + self.population.append(prompt) + self.population = reranking_population(self.population) + if len(self.population) > self.cfg.population_size: + if self.cfg.population_clusterization: + self.population = self.diversity_manager.maintain_diversity( + self.population, self.cfg.population_size + ) + filter_report = self.diversity_manager.get_filter_report() + if filter_report and self.logger is not None: + self.logger.log_diversity_filter( + step=step, filter_report=filter_report + ) + else: + self.population = self.population[: self.cfg.population_size] + return prompt in self.population + + def _crossover(self, iteration: int) -> ActionResult: + """Execute crossover and package its outcome for the controller. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: crossover quality change, cost, and offspring data. + """ + + (offspring, short_term_reflection) = self.crossover_operator.run( + iteration=iteration, + population=self.population, + problem_description=self.problem_description, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + improved = self._add_new_prompt(offspring, iteration) + delta_quality = self._update_elitist(offspring) + + self.short_term_reflections.append(short_term_reflection) + if len(self.short_term_reflections) > self.cfg.population_size: + self.short_term_reflections = self.short_term_reflections[1:] + + costs = self._calculate_costs(self.model.get_stats()) + return ActionResult( + action="crossover", + delta_quality=delta_quality, + cost_tokens=costs["total_tokens"], + improved=improved, + ) + + def _elitist_mutation(self, iteration: int) -> ActionResult: + """Execute an elite mutation and update reflection memory. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: mutation quality change, cost, and prompt data. + """ + + prompt_to_mutate = self.elitist + if random.random() < self.cfg.random_mutation_probability: + prompt_to_mutate = np.random.choice(self.population) + mutated, new_long_term_reflection = self.elitist_mutation_operator.run( + iteration=iteration, + elitist=prompt_to_mutate, + problem_description=self.problem_description, + long_term_reflection=self.long_term_reflection, + short_term_reflections=self.short_term_reflections, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + improved = self._add_new_prompt(mutated, iteration) + + delta_quality = self._update_elitist(mutated) + self.long_term_reflection = new_long_term_reflection + costs = self._calculate_costs(self.model.get_stats()) + return ActionResult( + action="mutation", + delta_quality=delta_quality, + cost_tokens=costs["total_tokens"], + improved=improved, + ) + + def _basic_mutation( + self, iteration: int, action_name: str, mutation_operator: Operator, **kwargs + ) -> ActionResult: + """Run a mutation operator and package its measured outcome. + + Args: + iteration (int): current optimization iteration. + action_name (str): controller-facing action name. + mutation_operator (Operator): operator to execute. + **kwargs (Any): additional arguments forwarded to ``run``. + + Returns: + ActionResult: mutation quality change, cost, and prompt data. + """ + + prompt_to_mutate = np.random.choice(self.population) + mutated = mutation_operator.run( + iteration=iteration, prompt=prompt_to_mutate, **kwargs + ) + improved = self._add_new_prompt(mutated, iteration) + + delta_quality = self._update_elitist(mutated) + costs = self._calculate_costs(self.model.get_stats()) + return ActionResult( + action=action_name, + delta_quality=delta_quality, + cost_tokens=costs["total_tokens"], + improved=improved, + ) + + def _compression(self, iteration: int) -> ActionResult: + """Compress a length-weighted population member. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: compression quality change, cost, and prompt data. + """ + + lengths = np.array([len(p.text) for p in self.population], dtype=float) + weights = lengths / lengths.sum() + ind = np.random.choice(len(self.population), p=weights) + prompt_to_mutate = self.population[ind] + mutated = self.compressor_operator.run( + iteration=iteration, prompt=prompt_to_mutate, evaluate_fn=self._evaluate + ) + improved = self._add_new_prompt(mutated, iteration) + delta_quality = self._update_elitist(mutated) + costs = self._calculate_costs(self.model.get_stats()) + return ActionResult( + action="long_compression", + delta_quality=delta_quality, + cost_tokens=costs["total_tokens"], + improved=improved, + ) + + def _gradient_step(self, iteration: int) -> ActionResult: + """Apply a feedback-derived textual-gradient mutation. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: textual-gradient action outcome. + """ + + return self._basic_mutation( + iteration=iteration, + action_name="gradient_step", + mutation_operator=self.gradient_step_operator, + problem_description=self.problem_description, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + + def _hype(self, iteration: int) -> ActionResult: + """Apply the HYPE optimizer to a sampled prompt. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: HYPE action outcome. + """ + + return self._basic_mutation( + iteration=iteration, + action_name="hype", + mutation_operator=self.hype_operator, + problem_description=self.problem_description, + evaluate_fn=self._evaluate, + ) + + def _long_term_mutation(self, iteration: int) -> ActionResult: + """Mutate a prompt using accumulated long-term reflection. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: long-term mutation outcome. + """ + + return self._basic_mutation( + iteration=iteration, + action_name="long_term_mutation", + mutation_operator=self.long_term_mutation_operator, + problem_description=self.problem_description, + long_term_reflection=self.long_term_reflection, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + + def _paraphrasing(self, iteration: int) -> ActionResult: + """Apply a problem-aware paraphrasing mutation. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: paraphrasing action outcome. + """ + + return self._basic_mutation( + iteration=iteration, + action_name="paraphrase", + mutation_operator=self.paraphrasing_operator, + problem_description=self.problem_description, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + + def _zero_order_mutation(self, iteration: int) -> ActionResult: + """Generate a zero-order mutation from the problem description. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: zero-order mutation outcome. + """ + + return self._basic_mutation( + iteration=iteration, + action_name="zero_order", + mutation_operator=self.zero_order_operator, + problem_description=self.problem_description, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + + def _creative_role_and_style_mutation(self, iteration: int) -> ActionResult: + """Apply a creative role-and-style mutation. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: role-and-style mutation outcome. + """ + + return self._basic_mutation( + iteration=iteration, + action_name="creative_role_and_style", + mutation_operator=self.creative_role_and_style_operator, + problem_description=self.problem_description, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + + def _creative_zero_order_mutation(self, iteration: int) -> ActionResult: + """Apply a creative zero-order mutation. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: creative zero-order mutation outcome. + """ + + return self._basic_mutation( + iteration=iteration, + action_name="creative_zero_order", + mutation_operator=self.creative_zero_order_mutation_operator, + problem_description=self.problem_description, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + + def _few_shots_mutation(self, iteration: int) -> ActionResult: + """Mutate a prompt by inserting a difficult few-shot example. + + Args: + iteration (int): current optimization iteration. + + Returns: + ActionResult: few-shot mutation outcome. + """ + + return self._basic_mutation( + iteration=iteration, + action_name="few_shot_mutation", + mutation_operator=self.few_shots_mutation_operator, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + + def _execute_action(self, action: str, iteration: int) -> ActionResult: + """Dispatch a named BRAVE action to its implementation. + + Args: + action (str): configured action name. + iteration (int): current optimization iteration. + + Returns: + ActionResult: result returned by the action implementation. + + Raises: + ValueError: if ``action`` is not supported. + """ + + match action: + case "crossover": + return self._crossover(iteration) + case "elitist_mutation": + return self._elitist_mutation(iteration) + case "compression": + return self._compression(iteration) + case "gradient_step": + return self._gradient_step(iteration) + case "hype": + return self._hype(iteration) + case "long_term_mutation": + return self._long_term_mutation(iteration) + case "paraphrase": + return self._paraphrasing(iteration) + case "zero_order": + return self._zero_order_mutation(iteration) + case "creative_role_and_style": + return self._creative_role_and_style_mutation(iteration) + case "creative_zero_order": + return self._creative_zero_order_mutation(iteration) + case "few_shot_mutation": + return self._few_shots_mutation(iteration) + case _: + raise ValueError(f"Unsupported action: {action}") + + def _init_train_batch_sampler(self) -> None: + """Initialize the configured stratified training-batch sampler.""" + + self.train_batch_sampler = None + self.current_train_batch_data = None + self.current_train_batch_targets = None + self.current_train_batch_indices = None + + batch_size = int(self.cfg.train_batch_size) + if not self.cfg.use_stratified_train_batches: + return + if batch_size <= 0: + return + if len(self.train_data) <= batch_size: + return + + if self.cfg.use_curriculum_batches: + self.train_batch_sampler = CurriculumStratifiedBatchSampler( + task=self.evaluator.task, + batch_size=batch_size, + total_steps=self.cfg.max_steps, + seed=self.seed, + generation_bins=self.cfg.generation_strata_bins, + warmup_steps=self.cfg.curriculum_warmup_steps, + max_alpha=self.cfg.curriculum_max_alpha, + ) + else: + self.train_batch_sampler = StratifiedBatchSampler( + task=self.evaluator.task, + batch_size=batch_size, + seed=self.seed, + generation_bins=self.cfg.generation_strata_bins, + ) + + def _refresh_train_batch(self, epoch: int) -> None: + """Select and cache the training subset for an epoch. + + Args: + epoch (int): current optimization epoch. + """ + + if self.train_batch_sampler is None: + self.current_train_batch_data = None + self.current_train_batch_targets = None + self.current_train_batch_indices = None + return + + batch_indices = self.train_batch_sampler.sample( + dataset=self.train_data, + targets=self.train_targets, + epoch=epoch, + ) + self.current_train_batch_indices = batch_indices + self.current_train_batch_data = [self.train_data[i] for i in batch_indices] + self.current_train_batch_targets = [ + self.train_targets[i] for i in batch_indices + ] + + def _get_train_eval_data(self) -> Tuple[List[str], List[Any]]: + """Return the active training batch or complete training split. + + Returns: + Tuple[List[str], List[Any]]: evaluation inputs and targets. + """ + + if ( + self.current_train_batch_data is not None + and self.current_train_batch_targets is not None + ): + return (self.current_train_batch_data, self.current_train_batch_targets) + return self.train_data, self.train_targets + + def _evaluate_val_cached(self, prompt: Prompt) -> float: + """Evaluate prompt on val set, using cached val_score if available. + + Does NOT overwrite prompt.score — train score is preserved. + + Args: + prompt (Prompt): prompt to evaluate or read from cache. + + Returns: + float: cached or newly computed validation score; ``0.0`` when + evaluation fails. + """ + if prompt.val_score is not None: + return prompt.val_score + try: + score, _ = self.evaluator.evaluate( + prompt=prompt.text, + dataset=self.val_data, + targets=self.val_targets, + failed_examples=self.cfg.bad_examples_num, + ) + except Exception: + score = 0.0 + prompt.set_val_score(float(score)) + return prompt.val_score + + def _evaluate(self, prompt: Prompt, split="train") -> None: + """Evaluates given prompt on self.dataset and records the score. + + Args: + prompt (Prompt): a prompt to evaluate. + split (str, optional): Which split of dataset to use. + Defaults to 'train'. + """ + if split == "val": + self._evaluate_val_cached(prompt) + return + dataset, targets = self._get_train_eval_data() + + try: + score, bad_examples = self.evaluator.evaluate( + prompt=prompt.text, + dataset=dataset, + targets=targets, + failed_examples=self.cfg.bad_examples_num, + ) + except Exception: + score = 0 + bad_examples = [] + + prompt.set_score(score) + prompt.set_bad_examples(bad_examples) + + if ( + split == "train" + and isinstance(self.train_batch_sampler, CurriculumStratifiedBatchSampler) + and self.current_train_batch_indices is not None + ): + bad_inputs = {ex["input"] for ex in bad_examples} + failed_global = [ + self.current_train_batch_indices[i] + for i, text in enumerate(dataset) + if text in bad_inputs + ] + self.train_batch_sampler.update_difficulties( + self.current_train_batch_indices, + failed_global, + ) + + def _llm_query(self, requests: List[str]) -> List[str]: + """Provides api to query requests to the model. + + Args: + requests (List[str]): string requests. + + Returns: + List[str]: model answers. + """ + + requests = [request.replace('"', "'") for request in requests] + + answers = None + for _ in range(5): + try: + answers = self.model.batch(requests) + break + except Exception as e: + print(e) + sleep(60) + + if answers is None: + return [""] * len(requests) + + answers = [a.content if isinstance(a, AIMessage) else a for a in answers] + + return answers + + def optimize( + self, + initial_prompt: str, + problem_description: str, + train_data: List[str], + train_targets: List[str], + val_data: List[str], + val_targets: List[str], + ) -> Dict[str, Any]: + """Optimize an initial prompt under token and step budgets. + + Args: + initial_prompt (str): seed prompt for population initialization. + problem_description (str): description of the target task. + train_data (List[str]): training inputs. + train_targets (List[str]): expected training outputs. + val_data (List[str]): validation inputs. + val_targets (List[str]): expected validation outputs. + + Returns: + Dict[str, Any]: best prompts, logs, controller diagnostics, and + efficiency statistics. + """ + + np.random.seed(self.seed) + random.seed(self.seed) + + self.artifacts = self._init_artifacts() + remaining_budget = self.cfg.initial_budget_tokens + self.initial_budget = self.cfg.initial_budget_tokens + + self.train_data = train_data + self.train_targets = train_targets + self.val_data = val_data + self.val_targets = val_targets + self.problem_description = problem_description + + needs_few_shot_data = ( + "few_shot_mutation" in self.actions + or "hard_few_shot_mutation" in self.actions + ) + if needs_few_shot_data: + cnt = self.cfg.few_shot_examples_from_data_cnt + ratio = cnt * 1.0 / len(self.train_data) + max_num = self.cfg.few_shot_examples_max_num + (self.train_data, examples_data, self.train_targets, examples_targets) = ( + get_dataset_split( + dataset=self.train_data, + target=self.train_targets, + validation_size=ratio, + train_as_test=False, + random_state=self.seed, + ) + ) + data_sample = list(zip(examples_data, examples_targets)) + if "few_shot_mutation" in self.actions: + self.few_shots_mutation_operator = HardFewShotExamplesOperator( + max_few_shot_examples_num=max_num, + data_sample=data_sample, + logger=self.logger, + ) + + self._init_train_batch_sampler() + self._refresh_train_batch(epoch=0) + + recent_quality: List[float] = [0.0] + no_improve_steps = 0 + useless_ops_count = 0 + spent_tokens = 0.0 + self.val_elitist: Optional[Prompt] = None + self.best_val_quality: float = float("-inf") + + self.population = self.population_initializer.run( + initial_prompt=initial_prompt, + population_size=self.cfg.initial_population_size, + problem_description=problem_description, + model=self.model, + llm_query_fn=self._llm_query, + evaluate_fn=self._evaluate, + ) + self.elitist = self.population[0] + self.best_quality = self.elitist.score + + spent = self._calculate_costs(self.model.get_stats())["total_tokens"] + remaining_budget -= spent + spent_tokens += spent + + for step in range(1, self.cfg.max_steps + 1): + if remaining_budget <= 0: + break + self._refresh_train_batch(epoch=step) + + if ( + self.cfg.rescore_steps > 0 + and step % self.cfg.rescore_steps == 0 + and self.train_batch_sampler is not None + ): + self._rescore_population() + rescore_cost = self._calculate_costs(self.model.get_stats())[ + "total_tokens" + ] + remaining_budget -= rescore_cost + spent_tokens += rescore_cost + + diversity = self.diversity_manager.compute_diversity(self.population) + state = self._compute_state( + best_quality=self.best_quality, + recent_quality=recent_quality, + useless_ops_count=useless_ops_count, + steps_done=step, + remaining_budget=remaining_budget, + population_diversity=diversity, + ) + x = self.featurizer.transform(state) + + action, score_dict = self.controller.select_action( + x=x, + remaining_budget_tokens=remaining_budget, + candidate_actions=self.actions, + ) + if action is None: + break + + action_score = float(score_dict.get(action, 0.0)) + is_fallback = "fallback" in score_dict + controller_diag = { + "action_score": action_score, + "fallback": 1.0 if "fallback" in score_dict else 0.0, + "ema_roi": float(self.controller.action_stats[action]["ema_roi"]), + "trials": float(self.controller.action_stats[action]["trials"]), + } + + if self.logger is not None: + self.logger.log_controller_state( + iteration=step, + selected_action=action, + action_scores=score_dict, + is_fallback=is_fallback, + action_stats=self.controller.action_stats, + global_step=self.controller.global_step, + ) + + result = self._execute_action(action=action, iteration=step) + + # Update controller with realized outcome + self.controller.update( + action=action, + x_before=x, + delta_quality=result.delta_quality, + actual_cost_tokens=result.cost_tokens, + improved=result.improved, + ) + + # Apply token budget + remaining_budget -= result.cost_tokens + spent_tokens += result.cost_tokens + + # Update pseudo quality + if result.improved: + no_improve_steps = 0 + else: + no_improve_steps += 1 + + useful_operation = result.delta_quality > 0.0 + if not useful_operation: + useless_ops_count += 1 + + self._update_artifacts( + action=action, result=result, improved=result.improved + ) + + recent_quality.append(self.best_quality) + value_per_token = result.delta_quality + value_per_token /= max(result.cost_tokens, 1e-6) + self.logs.append( + OptimizationLog( + step=step, + action=action, + score=action_score, + delta_quality=result.delta_quality, + cost_tokens=result.cost_tokens, + cumulative_spent=spent_tokens, + value_per_token=value_per_token, + useful_operation=useful_operation, + controller_diag=controller_diag, + remaining_budget=max(remaining_budget, 0.0), + best_quality=self.best_quality, + ) + ) + + if step % 5 == 0 and self.logger is not None: + self.logger.log_population(step, self.population) + + if ( + self.cfg.val_checkpoint_steps > 0 + and step % self.cfg.val_checkpoint_steps == 0 + ): + for candidate in self.population[: self.cfg.val_checkpoint_topk]: + self._update_val_elitist(candidate) + + if no_improve_steps >= self.cfg.patience_steps: + break + + if self.cfg.early_stop and self.best_quality == 1.0: + break + + if self.logger is not None: + self.logger.log_population(-2, self.population) + + for prompt in self.population: + prompt.set_score(self._evaluate_val_cached(prompt)) + self._update_val_elitist(prompt) + self.population.append(self.val_elitist) + self.population = list( + sorted(self.population, key=lambda prompt: prompt.val_score, reverse=True) + ) + if self.logger is not None: + self.logger.log_population(-3, self.population) + + summary = { + "best_prompt": self.elitist.text, + "best_quality": self.best_quality, + "best_val_prompt": ( + self.val_elitist.text if self.val_elitist else self.elitist.text + ), + "best_val_quality": self.best_val_quality if self.val_elitist else -1, + "remaining_budget_tokens": max(remaining_budget, 0.0), + "steps_done": len(self.logs), + "logs": self.logs, + "efficiency": self._build_efficiency_summary(), + "controller": self.controller.diagnostics(), + "artifacts": self.artifacts, + } + if self.log_dir and self.verbose: + self.export_logs_jsonl(f"{self.log_dir}/all_iterations_log.jsonl") + self.export_summary_json(f"{self.log_dir}/summary_log.json", summary) + return summary + + def _quality_at_budget_fraction(self, fraction: float) -> float: + """Return the best quality reached within a budget fraction. + + Args: + fraction (float): fraction of the initial budget to inspect. + + Returns: + float: best quality reached before the spending threshold. + """ + + target_spend = self.initial_budget * fraction + best = 0.0 + for row in self.logs: + if row.cumulative_spent <= target_spend: + best = max(best, row.best_quality) + return best + + def _build_efficiency_summary(self) -> Dict[str, Any]: + """Summarize quality gains, useful actions, and token efficiency. + + Returns: + Dict[str, Any]: aggregate spending and quality statistics. + """ + + if not self.logs: + return { + "spent_tokens": 0.0, + "value_per_1k_tokens": 0.0, + "useful_ops_ratio": 0.0, + "quality_at_budget": {"25%": 0.0, "50%": 0.0, "75%": 0.0, "100%": 0.0}, + } + + spent_tokens = max(self.logs[-1].cumulative_spent, 1e-6) + best_quality = self.logs[-1].best_quality + useful_ops = sum(1 for x in self.logs if x.useful_operation) + useful_ratio = useful_ops / max(len(self.logs), 1) + return { + "spent_tokens": spent_tokens, + "value_per_1k_tokens": (best_quality / spent_tokens) * 1000.0, + "useful_ops_ratio": useful_ratio, + "quality_at_budget": { + "25%": self._quality_at_budget_fraction(0.25), + "50%": self._quality_at_budget_fraction(0.50), + "75%": self._quality_at_budget_fraction(0.75), + "100%": self._quality_at_budget_fraction(1.00), + }, + } + + def export_logs_jsonl(self, path: str) -> None: + """Write per-step optimization logs as JSON Lines. + + Args: + path (str): destination JSONL path; parent directories are created. + """ + + out_path = Path(path) + out_path.parent.mkdir(parents=True, exist_ok=True) + with out_path.open("w", encoding="utf-8") as f: + for row in self.logs: + f.write( + json.dumps( + { + "step": row.step, + "action": row.action, + "score": row.score, + "delta_quality": row.delta_quality, + "cost_tokens": row.cost_tokens, + "cumulative_spent": row.cumulative_spent, + "value_per_token": row.value_per_token, + "useful_operation": row.useful_operation, + "controller_diag": row.controller_diag, + "remaining_budget": row.remaining_budget, + "best_quality": row.best_quality, + }, + ensure_ascii=True, + ) + + "\n" + ) + + def export_summary_json( + self, path: str, summary: Optional[Dict[str, Any]] = None + ) -> None: + """Write an optimization summary to a JSON file. + + Args: + path (str): destination JSON path; parent directories are created. + summary (Optional[Dict[str, Any]]): payload to write, or ``None`` + to derive a summary from current logs. + """ + + out_path = Path(path) + out_path.parent.mkdir(parents=True, exist_ok=True) + payload = ( + summary + if summary is not None + else { + "best_quality": self.logs[-1].best_quality if self.logs else 0.0, + "steps_done": len(self.logs), + "efficiency": self._build_efficiency_summary(), + } + ) + serializable = dict(payload) + if "logs" in serializable: + serializable["logs"] = [ + { + "step": row.step, + "action": row.action, + "score": row.score, + "delta_quality": row.delta_quality, + "cost_tokens": row.cost_tokens, + "cumulative_spent": row.cumulative_spent, + "value_per_token": row.value_per_token, + "useful_operation": row.useful_operation, + "controller_diag": row.controller_diag, + "remaining_budget": row.remaining_budget, + "best_quality": row.best_quality, + } + for row in serializable["logs"] + ] + with out_path.open("w", encoding="utf-8") as f: + json.dump(serializable, f, ensure_ascii=True, indent=2) diff --git a/coolprompt/optimizer/brave/operation_logger.py b/coolprompt/optimizer/brave/operation_logger.py new file mode 100644 index 0000000..63eda59 --- /dev/null +++ b/coolprompt/optimizer/brave/operation_logger.py @@ -0,0 +1,460 @@ +from dataclasses import dataclass, asdict +from pathlib import Path +from typing import List, Dict, Any, Tuple +from datetime import datetime +import yaml + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt + + +@dataclass +class ElitistMutationLog: + """Represent one logged mutation of the elite prompt.""" + + iteration: int + timestamp: str + elitist_prompt: str + prev_score: float + mutated_prompt: str + mutated_score: float + new_long_term_reflection: str + short_term_reflections: List[str] + + +@dataclass +class GradientStepLog: + """Represent one logged textual-gradient step.""" + + iteration: int + timestamp: str + prompt: str + prev_score: float + mutated_prompt: str + mutated_score: float + textual_gradient: str + + +@dataclass +class MutationLog: + """Represent one generic prompt mutation log entry.""" + + iteration: int + timestamp: str + prompt: str + prev_score: float + mutated_prompt: str + mutated_score: float + + +@dataclass +class CreativeRoleStyleMutationLog: + """Represent one creative role-and-style mutation.""" + + iteration: int + timestamp: str + prompt: str + prev_score: float + mutated_prompt: str + mutated_score: float + style: str + role: str + + +@dataclass +class FewShotExamplesMutationLog: + """Represent one mutation of a prompt's few-shot examples.""" + + iteration: int + timestamp: str + prompt: str + prev_score: float + mutated_prompt: str + mutated_score: float + added_few_shot: List[str] + removed_few_shot: List[str] + + +@dataclass +class CrossoverLog: + """Represent one crossover and its parent feedback.""" + + iteration: int + timestamp: str + parent1_prompt: str + parent1_score: float + parent2_prompt: str + parent2_score: float + offspring_prompt: str + offsprint_score: float + parent1_textual_gradient: str + parent2_textual_gradient: str + short_term_reflection: str + + +@dataclass +class PopulationLog: + """Represent a population snapshot at an optimization iteration.""" + + iteration: int + timestamp: str + population: List[Dict[str, Any]] + + +@dataclass +class ControllerStateLog: + """Represent a controller action-selection snapshot.""" + + iteration: int + timestamp: str + selected_action: str + action_scores: dict + is_fallback: bool + action_stats: dict + global_step: int + + +class OperationLogger: + """Persist BRAVE operation diagnostics as YAML files.""" + + def __init__(self, log_dir: str = "operation_logs"): + """Create a logger that writes beneath ``log_dir``. + + Args: + log_dir (str): directory in which YAML logs are stored. + """ + + self.log_dir = Path(log_dir) + self.log_dir.mkdir(parents=True, exist_ok=True) + + def _append_logs(self, filename: str, key: str, log_entry: Any) -> None: + """Append a dataclass log entry beneath a YAML document key. + + Args: + filename (str): YAML file to update. + key (str): top-level list key in the YAML document. + log_entry (Any): dataclass instance to serialize and append. + """ + + existing_logs = [] + if filename.exists(): + with open(filename, "r", encoding="utf-8") as f: + data = yaml.safe_load(f) or {} + existing_logs = data.get(key, []) + + existing_logs.append(asdict(log_entry)) + with open(filename, "w", encoding="utf-8") as f: + yaml.dump({key: existing_logs}, f, allow_unicode=True, sort_keys=False) + + def log_elitist_mutation( + self, + iteration: int, + elitist_prompt: str, + prev_score: float, + mutated_prompt: str, + mutated_score: float, + new_long_term_reflection: str, + short_term_reflections: List[str], + ) -> None: + """Log an elite mutation to a YAML file. + + Args: + iteration (int): optimization iteration. + elitist_prompt (str): elite prompt before mutation. + prev_score (float): elite score before mutation. + mutated_prompt (str): generated prompt text. + mutated_score (float): generated prompt score. + new_long_term_reflection (str): updated long-term reflection. + short_term_reflections (List[str]): reflections used by mutation. + """ + log_entry = ElitistMutationLog( + iteration=iteration, + timestamp=datetime.now().isoformat(), + elitist_prompt=elitist_prompt, + prev_score=prev_score, + mutated_prompt=mutated_prompt, + mutated_score=mutated_score, + new_long_term_reflection=new_long_term_reflection, + short_term_reflections=short_term_reflections, + ) + + log_file = self.log_dir / "elitist_mutations.yaml" + self._append_logs(log_file, "mutations", log_entry) + + def log_mutation( + self, + iteration: int, + prompt: str, + prev_score: float, + mutated_prompt: str, + mutated_score: float, + file_name: str = "mutations", + ) -> None: + """Log a generic mutation to a YAML file. + + Args: + iteration (int): optimization iteration. + prompt (str): source prompt text. + prev_score (float): source prompt score. + mutated_prompt (str): generated prompt text. + mutated_score (float): generated prompt score. + file_name (str): output filename without extension. + """ + log_entry = MutationLog( + iteration=iteration, + timestamp=datetime.now().isoformat(), + prompt=prompt, + prev_score=prev_score, + mutated_prompt=mutated_prompt, + mutated_score=mutated_score, + ) + + log_file = self.log_dir / f"{file_name}.yaml" + self._append_logs(log_file, "mutations", log_entry) + + def log_gradient_step( + self, + iteration: int, + prompt: str, + prev_score: float, + mutated_prompt: str, + mutated_score: float, + textual_gradient: str, + ) -> None: + """Log a textual-gradient mutation. + + Args: + iteration (int): optimization iteration. + prompt (str): source prompt text. + prev_score (float): source prompt score. + mutated_prompt (str): generated prompt text. + mutated_score (float): generated prompt score. + textual_gradient (str): feedback applied by the mutation. + """ + + log_entry = GradientStepLog( + iteration=iteration, + timestamp=datetime.now().isoformat(), + prompt=prompt, + prev_score=prev_score, + mutated_prompt=mutated_prompt, + mutated_score=mutated_score, + textual_gradient=textual_gradient, + ) + + log_file = self.log_dir / "gradient_steps.yaml" + self._append_logs(log_file, "mutations", log_entry) + + def log_creative_role_style_mutation( + self, + iteration: int, + prompt: str, + prev_score: float, + mutated_prompt: str, + mutated_score: float, + style: str, + role: str, + ) -> None: + """Log a creative role-and-style mutation. + + Args: + iteration (int): optimization iteration. + prompt (str): source prompt text. + prev_score (float): source prompt score. + mutated_prompt (str): generated prompt text. + mutated_score (float): generated prompt score. + style (str): generated writing style. + role (str): generated model role. + """ + + log_entry = CreativeRoleStyleMutationLog( + iteration=iteration, + timestamp=datetime.now().isoformat(), + prompt=prompt, + prev_score=prev_score, + mutated_prompt=mutated_prompt, + mutated_score=mutated_score, + style=style, + role=role, + ) + + log_file = self.log_dir / "creative_role_style_mutations.yaml" + self._append_logs(log_file, "mutations", log_entry) + + def log_few_shot_mutation( + self, + iteration: int, + prompt: str, + prev_score: float, + mutated_prompt: str, + mutated_score: float, + added_few_shot: Tuple[str, str], + removed_few_shot: Tuple[str, str], + file_name: str = "few_shot_mutations", + ) -> None: + """Log an addition or replacement of a few-shot example. + + Args: + iteration (int): optimization iteration. + prompt (str): source prompt text. + prev_score (float): source prompt score. + mutated_prompt (str): generated prompt text. + mutated_score (float): generated prompt score. + added_few_shot (Tuple[str, str]): inserted input-output example. + removed_few_shot (Tuple[str, str]): replaced example, if any. + file_name (str): output filename without extension. + """ + + log_entry = FewShotExamplesMutationLog( + iteration=iteration, + timestamp=datetime.now().isoformat(), + prompt=prompt, + prev_score=prev_score, + mutated_prompt=mutated_prompt, + mutated_score=mutated_score, + added_few_shot=list(added_few_shot), + removed_few_shot=list(removed_few_shot), + ) + + log_file = self.log_dir / f"{file_name}.yaml" + self._append_logs(log_file, "mutations", log_entry) + + def log_crossover( + self, + iteration: int, + parent1_prompt: str, + parent1_score: float, + parent2_prompt: str, + parent2_score: float, + parent1_textual_gradient: str, + parent2_textual_gradient: str, + offspring_prompt: str, + offspring_score: float, + short_term_reflection: str, + ) -> None: + """Log a crossover operation to a YAML file. + + Args: + iteration (int): optimization iteration. + parent1_prompt (str): first parent text. + parent1_score (float): first parent score. + parent2_prompt (str): second parent text. + parent2_score (float): second parent score. + parent1_textual_gradient (str): first parent feedback. + parent2_textual_gradient (str): second parent feedback. + offspring_prompt (str): generated offspring text. + offspring_score (float): generated offspring score. + short_term_reflection (str): crossover reflection. + """ + log_entry = CrossoverLog( + iteration=iteration, + timestamp=datetime.now().isoformat(), + parent1_prompt=parent1_prompt, + parent1_score=parent1_score, + parent2_prompt=parent2_prompt, + parent2_score=parent2_score, + parent1_textual_gradient=parent1_textual_gradient, + parent2_textual_gradient=parent2_textual_gradient, + offspring_prompt=offspring_prompt, + offsprint_score=offspring_score, + short_term_reflection=short_term_reflection, + ) + + log_file = self.log_dir / "crossovers.yaml" + self._append_logs(log_file, "crossovers", log_entry) + + def log_population(self, iteration: int, population: List[Prompt]) -> None: + """Write a complete population snapshot for an iteration. + + Args: + iteration (int): optimization iteration. + population (List[Prompt]): evaluated prompts to serialize. + """ + + population = [p.to_dict() for p in population] + log_entry = PopulationLog( + iteration=iteration, + timestamp=datetime.now().isoformat(), + population=population, + ) + + log_file = self.log_dir / f"{iteration}_population.yaml" + with open(log_file, "w", encoding="utf-8") as f: + yaml.dump(asdict(log_entry), f, allow_unicode=True, sort_keys=False) + + def log_controller_state( + self, + iteration: int, + selected_action: str, + action_scores: dict, + is_fallback: bool, + action_stats: dict, + global_step: int, + ) -> None: + """Log controller state and action selection. + + Args: + iteration (int): optimization iteration. + selected_action (str): selected action name. + action_scores (dict): controller scores by action. + is_fallback (bool): whether fallback selection was used. + action_stats (dict): current per-action statistics. + global_step (int): controller step counter. + """ + log_entry = ControllerStateLog( + iteration=iteration, + timestamp=datetime.now().isoformat(), + selected_action=selected_action, + action_scores=action_scores, + is_fallback=is_fallback, + action_stats=action_stats, + global_step=global_step, + ) + + log_file = self.log_dir / "controller_state.yaml" + existing_logs = [] + if log_file.exists(): + with open(log_file, "r", encoding="utf-8") as f: + data = yaml.safe_load(f) or {} + existing_logs = data.get("controller_states", []) + + existing_logs.append(asdict(log_entry)) + with open(log_file, "w", encoding="utf-8") as f: + yaml.dump( + {"controller_states": existing_logs}, + f, + allow_unicode=True, + sort_keys=False, + ) + + def log_diversity_filter(self, step: int, filter_report: dict) -> None: + """Log a population-diversity filtering report. + + Args: + step (int): optimization step. + filter_report (dict): thresholds, clusters, and removed indices. + """ + log_file = self.log_dir / "diversity_filter.yaml" + existing_logs = [] + if log_file.exists(): + with open(log_file, "r", encoding="utf-8") as f: + data = yaml.safe_load(f) or {} + existing_logs = data.get("diversity_filters", []) + + log_entry = { + "step": step, + "timestamp": datetime.now().isoformat(), + "duplicate_threshold": filter_report.get("threshold"), + "num_clusters": filter_report.get("num_clusters"), + "num_removed": filter_report.get("num_removed"), + "removed_indices": filter_report.get("removed_indices"), + "deduplication_removed": filter_report.get("deduplication_removed"), + } + + existing_logs.append(log_entry) + with open(log_file, "w", encoding="utf-8") as f: + yaml.dump( + {"diversity_filters": existing_logs}, + f, + allow_unicode=True, + sort_keys=False, + ) diff --git a/coolprompt/optimizer/brave/operators/__init__.py b/coolprompt/optimizer/brave/operators/__init__.py new file mode 100644 index 0000000..3697c62 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/__init__.py @@ -0,0 +1,53 @@ +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.operators.bigger_initializer import ( + BiggerPopulationInitializationOperator, +) +from coolprompt.optimizer.brave.operators.compressor import CompressorOperator +from coolprompt.optimizer.brave.operators.creative_role_and_style import ( + CreativeRoleAndStyleMutationOperator, +) +from coolprompt.optimizer.brave.operators.creative_zero_order import ( + CreativeZeroOrderMutationOperator, +) +from coolprompt.optimizer.brave.operators.crossover import CrossoverOperator +from coolprompt.optimizer.brave.operators.elitist_mutation import ( + ElitistMutationOperator, +) +from coolprompt.optimizer.brave.operators.few_shot_examples import ( + FewShotExamplesOperator, +) +from coolprompt.optimizer.brave.operators.hard_few_shot_examples import ( + HardFewShotExamplesOperator, +) +from coolprompt.optimizer.brave.operators.gradient_step import GradientStepOperator +from coolprompt.optimizer.brave.operators.hype import HypeOperator +from coolprompt.optimizer.brave.operators.initializer import ( + PopulationInitializationOperator, +) +from coolprompt.optimizer.brave.operators.long_term_mutation import ( + LongTermMutationOperator, +) +from coolprompt.optimizer.brave.operators.paraphrase_initializer import ( + ParaphraseInitializationOperator, +) +from coolprompt.optimizer.brave.operators.paraphrasing import ParaphrasingByPDOperator +from coolprompt.optimizer.brave.operators.zero_order import ZeroOrderMutationOperator + +__all__ = [ + "Operator", + "CrossoverOperator", + "ElitistMutationOperator", + "PopulationInitializationOperator", + "CompressorOperator", + "GradientStepOperator", + "HypeOperator", + "LongTermMutationOperator", + "ParaphrasingByPDOperator", + "ZeroOrderMutationOperator", + "CreativeRoleAndStyleMutationOperator", + "CreativeZeroOrderMutationOperator", + "FewShotExamplesOperator", + "HardFewShotExamplesOperator", + "ParaphraseInitializationOperator", + "BiggerPopulationInitializationOperator", +] diff --git a/coolprompt/optimizer/brave/operators/basic_operator.py b/coolprompt/optimizer/brave/operators/basic_operator.py new file mode 100644 index 0000000..5a6c101 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/basic_operator.py @@ -0,0 +1,31 @@ +from abc import ABC, abstractmethod +from typing import Optional, Any +from coolprompt.optimizer.brave.operation_logger import OperationLogger + + +class Operator(ABC): + """Base interface for BRAVE prompt-transformation operators.""" + + def __init__(self, logger: Optional[OperationLogger] = None) -> None: + """Store an optional operation logger. + + Args: + logger (Optional[OperationLogger]): logger for operator + diagnostics. + """ + + self.logger = logger + + @abstractmethod + def run(self, *args: Any, **kwargs: Any) -> Any: + """Run the operator and return its generated prompt or result. + + Args: + *args (Any): positional operator inputs. + **kwargs (Any): keyword operator inputs. + + Returns: + Any: operator-specific result. + """ + + pass diff --git a/coolprompt/optimizer/brave/operators/bigger_initializer.py b/coolprompt/optimizer/brave/operators/bigger_initializer.py new file mode 100644 index 0000000..70965c0 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/bigger_initializer.py @@ -0,0 +1,114 @@ +from typing import List, Callable +from langchain_core.language_models.base import BaseLanguageModel + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.operators.creative_role_and_style import ( + CreativeRoleAndStyleMutationOperator, +) +from coolprompt.optimizer.brave.operators.creative_zero_order import ( + CreativeZeroOrderMutationOperator, +) +from coolprompt.optimizer.brave.operators.gradient_step import GradientStepOperator +from coolprompt.optimizer.brave.operators.hype import HypeOperator +from coolprompt.optimizer.brave.prompt_templates import PROMPT_BY_DESCRIPTION_TEMPLATE +from coolprompt.optimizer.brave.utils import reranking_population, PROMPT_TAGS +from coolprompt.utils.parsing import extract_answer + + +class BiggerPopulationInitializationOperator(Operator): + """Initialize a large prompt population through several strategies.""" + + def run( + self, + initial_prompt: str, + problem_description: str, + population_size: int, # just for interface + model: BaseLanguageModel, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> List[Prompt]: + """Generate, evaluate, and return an expanded initial population. + + Args: + initial_prompt (str): seed prompt included in the population. + problem_description (str): description of the target task. + population_size (int): target size retained for interface parity. + model (BaseLanguageModel): model used by delegated operators. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + List[Prompt]: evaluated initial prompt candidates. + """ + + prompt_by_description_template = PROMPT_BY_DESCRIPTION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description + ) + prompt_by_pd = extract_answer( + answer=llm_query_fn([prompt_by_description_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + prompt_by_pd = Prompt(prompt_by_pd, origin=PromptOrigin.BY_PD) + + initial_prompt = Prompt(initial_prompt, PromptOrigin.MANUAL) + evaluate_fn(initial_prompt, "train") + + hype_operator = HypeOperator(model, logger=self.logger) + hyped_prompt = hype_operator.run( + iteration=-1, + prompt=initial_prompt, + problem_description=problem_description, + evaluate_fn=evaluate_fn, + ) + + gradient_step_operator = GradientStepOperator(self.logger) + gradient_step_prompt = gradient_step_operator.run( + iteration=-1, + prompt=initial_prompt, + problem_description=problem_description, + llm_query_fn=llm_query_fn, + evaluate_fn=evaluate_fn, + ) + + creative_zero_order_operator = CreativeZeroOrderMutationOperator( + logger=self.logger + ) + creative_zero_order_prompt = creative_zero_order_operator.run( + iteration=-1, + prompt=initial_prompt, + problem_description=problem_description, + llm_query_fn=llm_query_fn, + evaluate_fn=evaluate_fn, + ) + + creative_ras_operator = CreativeRoleAndStyleMutationOperator(logger=self.logger) + creative_roled_and_styled_prompt = creative_ras_operator.run( + iteration=-1, + prompt=initial_prompt, + problem_description=problem_description, + llm_query_fn=llm_query_fn, + evaluate_fn=evaluate_fn, + ) + + population = [ + prompt_by_pd, + hyped_prompt, + gradient_step_prompt, + creative_zero_order_prompt, + creative_roled_and_styled_prompt, + ] + + for prompt in population: + evaluate_fn(prompt, "train") + + population.append(initial_prompt) + + population = reranking_population(population) + + if self.logger is not None: + self.logger.log_population(iteration=0, population=population) + + return population diff --git a/coolprompt/optimizer/brave/operators/compressor.py b/coolprompt/optimizer/brave/operators/compressor.py new file mode 100644 index 0000000..be317a5 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/compressor.py @@ -0,0 +1,59 @@ +from typing import Callable +from langchain_core.language_models.base import BaseLanguageModel + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.prompt_compressor.compressor import PromptCompressor + + +class CompressorOperator(Operator): + """Compress a prompt while preserving its intended behavior.""" + + def __init__(self, model: BaseLanguageModel, **kwargs) -> None: + """Create a prompt compressor backed by the supplied model. + + Args: + model (BaseLanguageModel): model used for compression. + **kwargs (Any): base-operator arguments such as ``logger``. + """ + + super().__init__(**kwargs) + self.compressor = PromptCompressor(model) + + def run( + self, + iteration: int, + prompt: Prompt, + evaluate_fn: Callable[[Prompt, str], None], + ) -> Prompt: + """Compress and evaluate a prompt. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): prompt to compress. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Prompt: evaluated compressed prompt, or a zero-scored failure + placeholder when compression raises an exception. + """ + + try: + compressed = self.compressor.compress(prompt.text) + compressed = Prompt(compressed, origin=PromptOrigin.COMPRESSED) + evaluate_fn(compressed, "train") + except Exception: + compressed = Prompt("failed to compress", origin=PromptOrigin.COMPRESSED) + compressed.set_score(0) + + if self.logger is not None: + self.logger.log_mutation( + iteration=iteration, + prompt=prompt.text, + prev_score=prompt.score, + mutated_prompt=compressed.text, + mutated_score=compressed.score, + file_name="compressions", + ) + + return compressed diff --git a/coolprompt/optimizer/brave/operators/creative_role_and_style.py b/coolprompt/optimizer/brave/operators/creative_role_and_style.py new file mode 100644 index 0000000..a853ea6 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/creative_role_and_style.py @@ -0,0 +1,78 @@ +from typing import List, Tuple, Callable + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import ( + CREATIVE_STYLE_AND_ROLE_TEMPLATE, + CREATIVE_ZERO_ORDER_MUTATION_TEMPLATE, +) +from coolprompt.optimizer.brave.utils import PROMPT_TAGS, STYLE_TAGS, ROLE_TAGS +from coolprompt.utils.parsing import extract_answer + + +class CreativeRoleAndStyleMutationOperator(Operator): + """Mutate a prompt by inventing and applying a role and writing style.""" + + def run( + self, + iteration: int, + prompt: Prompt, + problem_description: str, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Tuple[Prompt, str]: + """Generate and evaluate a role-and-style mutation. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): source prompt. + problem_description (str): description of the target task. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Tuple[Prompt, str]: evaluated mutation and generated style text. + """ + + style_and_role_template = CREATIVE_STYLE_AND_ROLE_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description + ) + print(style_and_role_template) + model_answer = llm_query_fn([style_and_role_template])[0] + print(model_answer) + style = extract_answer( + answer=model_answer, tags=STYLE_TAGS, format_mismatch_label="" + ) + role = extract_answer( + answer=model_answer, tags=ROLE_TAGS, format_mismatch_label="" + ) + + mutation_template = CREATIVE_ZERO_ORDER_MUTATION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + STYLE=style, + ROLE=role, + PROMPT=prompt.text, + ) + mutated_offspring = extract_answer( + answer=llm_query_fn([mutation_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + mutated_offspring = Prompt( + mutated_offspring, origin=PromptOrigin.CREATIVE_IN_STYLE_OF + ) + evaluate_fn(mutated_offspring, "train") + + if self.logger is not None: + self.logger.log_creative_role_style_mutation( + iteration=iteration, + prompt=prompt.text, + prev_score=prompt.score, + mutated_prompt=mutated_offspring.text, + mutated_score=mutated_offspring.score, + style=style, + role=role, + ) + + return mutated_offspring diff --git a/coolprompt/optimizer/brave/operators/creative_zero_order.py b/coolprompt/optimizer/brave/operators/creative_zero_order.py new file mode 100644 index 0000000..cd0dd0c --- /dev/null +++ b/coolprompt/optimizer/brave/operators/creative_zero_order.py @@ -0,0 +1,58 @@ +from typing import List, Callable + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import ( + CREATIVE_ZERO_ORDER_MUTATION_TEMPLATE, +) +from coolprompt.optimizer.brave.utils import PROMPT_TAGS +from coolprompt.utils.parsing import extract_answer + + +class CreativeZeroOrderMutationOperator(Operator): + """Apply a creative zero-order mutation without evaluation feedback.""" + + def run( + self, + iteration: int, + prompt: Prompt, # won't be used, but needed for the interface + problem_description: str, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Prompt: + """Generate and evaluate a creative zero-order mutation. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): interface-compatible source prompt; not read. + problem_description (str): description of the target task. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Prompt: evaluated generated prompt. + """ + + generating_template = CREATIVE_ZERO_ORDER_MUTATION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + ) + generated = extract_answer( + answer=llm_query_fn([generating_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + generated = Prompt(generated, origin=PromptOrigin.CREATIVE_ZERO_ORDER_PD) + evaluate_fn(generated, "train") + + if self.logger is not None: + self.logger.log_mutation( + iteration=iteration, + prompt="", + prev_score=-1.0, + mutated_prompt=generated.text, + mutated_score=generated.score, + file_name="creative_zero_orders", + ) + + return generated diff --git a/coolprompt/optimizer/brave/operators/crossover.py b/coolprompt/optimizer/brave/operators/crossover.py new file mode 100644 index 0000000..ff5eb32 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/crossover.py @@ -0,0 +1,156 @@ +from typing import List, Tuple, Callable +import numpy as np + +from coolprompt.optimizer.reflective_prompt.prompt import ( + BadExample, + Prompt, + PromptOrigin, +) +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import ( + TEXTUAL_GRADIENT_TEMPLATE, + SHORT_TERM_REFLECTION_TEMPLATE, + CROSSOVER_TEMPLATE, +) +from coolprompt.optimizer.brave.utils import PROMPT_TAGS, HINT_TAGS, FEEDBACK_TAGS +from coolprompt.utils.parsing import extract_answer + + +class CrossoverOperator(Operator): + """Combine two parent prompts using feedback-driven reflection.""" + + def _make_bad_examples(self, bad_examples: List[BadExample]) -> str: + """Converts an array of bad examples into string format + + Args: + bad_examples (List[BadExample]): list of bad examples. + + Returns: + str: string representation of bad examples + """ + + return "\n\n".join( + [ + "\n".join( + ( + f"Input: {example.input}", + f"Model Output: {example.output}", + f"Correct Output: {example.correct}", + ) + ) + for example in bad_examples + ] + ) + + def _gen_textual_gradient( + self, + prompt: Prompt, + problem_description: str, + llm_query_fn: Callable[[List[str]], List[str]], + ) -> str: + """Generate a textual gradient for a prompt. + + Args: + prompt (Prompt): prompt to critique. + problem_description (str): description of the target task. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + + Returns: + str: cached or newly generated textual gradient. + """ + + if prompt.gradient is not None: + return prompt.gradient + + request = TEXTUAL_GRADIENT_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + PROMPT=prompt.text, + EXAMPLES=self._make_bad_examples(prompt.bad_examples), + ) + gradient = extract_answer( + answer=llm_query_fn([request])[0], + tags=FEEDBACK_TAGS, + format_mismatch_label="", + ) + prompt.gradient = gradient + return gradient + + def run( + self, + iteration: int, + population: List[Prompt], + problem_description: str, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Tuple[Prompt, str]: + """Cross two prompts and evaluate their offspring. + + Args: + iteration (int): optimization iteration used for logging. + population (List[Prompt]): parent prompt population. + problem_description (str): description of the target task. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Tuple[Prompt, str]: evaluated offspring and short-term reflection. + """ + + scores = np.array([prompt.score for prompt in population]) + probas = (scores + 1e-5) / np.sum(scores + 1e-5) + parents = np.random.choice(population, size=2, replace=False, p=probas) + + parents = [ + ( + parent, + self._gen_textual_gradient( + parent, problem_description, llm_query_fn=llm_query_fn + ), + ) + for parent in parents + ] + + short_term_template = SHORT_TERM_REFLECTION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + PROMPT1=parents[0][0].text, + FEEDBACK1=parents[0][1], + PROMPT2=parents[1][0].text, + FEEDBACK2=parents[1][1], + ) + short_term_reflection = extract_answer( + answer=llm_query_fn([short_term_template])[0], + tags=HINT_TAGS, + format_mismatch_label="", + ) + + crossover_template = CROSSOVER_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + PARENT1=parents[0][0].text, + PARENT2=parents[1][0].text, + SHORT_TERM_REFLECTION=short_term_reflection, + ) + offspring = extract_answer( + answer=llm_query_fn([crossover_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + offspring = Prompt(offspring, origin=PromptOrigin.CROSSOVER) + evaluate_fn(offspring, "train") + + if self.logger is not None: + self.logger.log_crossover( + iteration=iteration, + parent1_prompt=parents[0][0].text, + parent1_score=parents[0][0].score, + parent2_prompt=parents[1][0].text, + parent2_score=parents[1][0].score, + parent1_textual_gradient=parents[0][1], + parent2_textual_gradient=parents[1][1], + offspring_prompt=offspring.text, + offspring_score=offspring.score, + short_term_reflection=short_term_reflection, + ) + + return offspring, short_term_reflection diff --git a/coolprompt/optimizer/brave/operators/elitist_mutation.py b/coolprompt/optimizer/brave/operators/elitist_mutation.py new file mode 100644 index 0000000..74ea89f --- /dev/null +++ b/coolprompt/optimizer/brave/operators/elitist_mutation.py @@ -0,0 +1,85 @@ +from typing import List, Tuple, Callable + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import ( + LONG_TERM_REFLECTION_TEMPLATE, + LONG_TERM_REFLECTION_UPDATE_TEMPLATE, + ELITIST_MUTATION_TEMPLATE, +) +from coolprompt.optimizer.brave.utils import PROMPT_TAGS, HINT_TAGS +from coolprompt.utils.parsing import extract_answer + + +class ElitistMutationOperator(Operator): + """Mutate the elite prompt using short- and long-term reflections.""" + + def run( + self, + iteration: int, + elitist: Prompt, + problem_description: str, + long_term_reflection: str, + short_term_reflections: List[str], + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Tuple[Prompt, str]: + """Generate an elite mutation and update long-term reflection. + + Args: + iteration (int): optimization iteration used for logging. + elitist (Prompt): elite prompt to mutate. + problem_description (str): description of the target task. + long_term_reflection (str): accumulated reflection text. + short_term_reflections (List[str]): recent crossover reflections. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Tuple[Prompt, str]: evaluated mutation and new long-term + reflection. + """ + + if long_term_reflection == "": + long_term_template = LONG_TERM_REFLECTION_TEMPLATE.format( + SHORT_TERM_REFLECTIONS="/n".join(short_term_reflections) + ) + else: + long_term_template = LONG_TERM_REFLECTION_UPDATE_TEMPLATE.format( + SHORT_TERM_REFLECTIONS="/n".join(short_term_reflections), + LONG_TERM_REFLECTION=long_term_reflection, + ) + new_long_term_reflection = extract_answer( + answer=llm_query_fn([long_term_template])[0], + tags=HINT_TAGS, + format_mismatch_label="", + ) + + mutation_template = ELITIST_MUTATION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + ELITIST_PROMPT=elitist.text, + LONG_TERM_REFLECTION=new_long_term_reflection, + ) + mutated_offspring = extract_answer( + answer=llm_query_fn([mutation_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + mutated_offspring = Prompt( + mutated_offspring, origin=PromptOrigin.ELITIST_MUTATION + ) + evaluate_fn(mutated_offspring, "train") + + if self.logger is not None: + self.logger.log_elitist_mutation( + iteration=iteration, + elitist_prompt=elitist.text, + prev_score=elitist.score, + mutated_prompt=mutated_offspring.text, + mutated_score=mutated_offspring.score, + new_long_term_reflection=new_long_term_reflection, + short_term_reflections=short_term_reflections, + ) + + return mutated_offspring, new_long_term_reflection diff --git a/coolprompt/optimizer/brave/operators/few_shot_examples.py b/coolprompt/optimizer/brave/operators/few_shot_examples.py new file mode 100644 index 0000000..421f654 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/few_shot_examples.py @@ -0,0 +1,142 @@ +from typing import List, Tuple, Callable +import numpy as np + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import ( + FEW_SHOT_EXAMPLES_REMOVING_TEMPLATE, + FEW_SHOT_EXAMPLES_INCORPORATING_TEMPLATE, +) +from coolprompt.optimizer.brave.utils import PROMPT_TAGS +from coolprompt.utils.parsing import extract_answer + + +class FewShotExamplesOperator(Operator): + """Mutate prompts by adding or replacing embedded few-shot examples.""" + + def __init__( + self, + max_few_shot_examples_num: int, + data_sample: List[Tuple[str, str]], + **kwargs, + ) -> None: + """Store the candidate example pool and maximum example count. + + Args: + max_few_shot_examples_num (int): maximum examples in a prompt. + data_sample (List[Tuple[str, str]]): candidate input-output pairs. + **kwargs (Any): base-operator arguments such as ``logger``. + """ + + super().__init__(**kwargs) + self.max_few_shot_examples_num = max_few_shot_examples_num + self.examples = data_sample + + def _filter_possible_examples( + self, prompt_few_shots: List[Tuple[str, str]] + ) -> List[Tuple[str, str]]: + """Return examples not already attached to the prompt. + + Args: + prompt_few_shots (List[Tuple[str, str]]): current prompt examples. + + Returns: + List[Tuple[str, str]]: examples eligible for insertion. + """ + + return [example for example in self.examples if example not in prompt_few_shots] + + def _prepare_examples(self, examples: List[Tuple[str, str]]) -> str: + """Format examples for insertion into an LLM request. + + Args: + examples (List[Tuple[str, str]]): input-output pairs. + + Returns: + str: examples separated by blank lines. + """ + + return "\n\n".join([f"Input: {inp}\nOutput: {out}" for inp, out in examples]) + + def run( + self, + iteration: int, + prompt: Prompt, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Prompt: + """Insert a sampled example, rewrite the prompt, and evaluate it. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): prompt whose examples should change. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Prompt: evaluated rewritten prompt, or a zero-scored failure + placeholder when rewriting fails. + """ + + possible_examples = self._filter_possible_examples( + prompt_few_shots=prompt.few_shot_examples + ) + ind = np.random.choice(len(possible_examples)) + example_to_add = possible_examples[ind] + + original_few_shots = list(prompt.few_shot_examples) + removed = ("", "") + if len(prompt.few_shot_examples) == self.max_few_shot_examples_num: + ind = np.random.choice(len(prompt.few_shot_examples)) + removed = prompt.few_shot_examples[ind] + prompt.few_shot_examples[ind] = example_to_add + else: + prompt.add_few_shot_example(example_to_add) + + removing_template = FEW_SHOT_EXAMPLES_REMOVING_TEMPLATE.format( + PROMPT=prompt.text + ) + prompt_without_few_shots = extract_answer( + answer=llm_query_fn([removing_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + + few_shot_template = FEW_SHOT_EXAMPLES_INCORPORATING_TEMPLATE.format( + PROMPT=prompt_without_few_shots, + EXAMPLES=self._prepare_examples(prompt.few_shot_examples), + ) + try: + prompt_with_few_shots = extract_answer( + answer=llm_query_fn([few_shot_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + except Exception: + prompt_with_few_shots = None + + if prompt_with_few_shots: + mutated_offspring = Prompt( + prompt_with_few_shots, origin=PromptOrigin.FEW_SHOT + ) + evaluate_fn(mutated_offspring, "train") + else: + prompt.few_shot_examples = original_few_shots + mutated_offspring = Prompt( + "FAILED TO PRODUCE", origin=PromptOrigin.FEW_SHOT + ) + mutated_offspring.set_score(0) + + if self.logger is not None: + self.logger.log_few_shot_mutation( + iteration=iteration, + prompt=prompt.text, + prev_score=prompt.score, + mutated_prompt=mutated_offspring.text, + mutated_score=mutated_offspring.score, + added_few_shot=example_to_add, + removed_few_shot=removed, + ) + + return mutated_offspring diff --git a/coolprompt/optimizer/brave/operators/gradient_step.py b/coolprompt/optimizer/brave/operators/gradient_step.py new file mode 100644 index 0000000..0b15bbc --- /dev/null +++ b/coolprompt/optimizer/brave/operators/gradient_step.py @@ -0,0 +1,128 @@ +from typing import List, Callable + +from coolprompt.optimizer.reflective_prompt.prompt import ( + Prompt, + PromptOrigin, + BadExample, +) +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import ( + TEXTUAL_GRADIENT_TEMPLATE, + GRADIENT_STEP_TEMPLATE, +) +from coolprompt.optimizer.brave.utils import PROMPT_TAGS, FEEDBACK_TAGS +from coolprompt.utils.parsing import extract_answer + + +class GradientStepOperator(Operator): + """Improve a prompt using a textual gradient from failed examples.""" + + def _make_bad_examples(self, bad_examples: List[BadExample]) -> str: + """Converts an array of bad examples into string format + + Args: + bad_examples (List[BadExample]): list of bad examples. + + Returns: + str: string representation of bad examples + """ + + return "\n\n".join( + [ + "\n".join( + ( + f"Input: {example.input}", + f"Model Output: {example.output}", + f"Correct Output: {example.correct}", + ) + ) + for example in bad_examples + ] + ) + + def _gen_textual_gradient( + self, + prompt: Prompt, + problem_description: str, + llm_query_fn: Callable[[List[str]], List[str]], + ) -> str: + """Generate a textual gradient for a prompt. + + Args: + prompt (Prompt): prompt to critique. + problem_description (str): description of the target task. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + + Returns: + str: cached or newly generated textual gradient. + """ + + if prompt.gradient is not None: + return prompt.gradient + + request = TEXTUAL_GRADIENT_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + PROMPT=prompt.text, + EXAMPLES=self._make_bad_examples(prompt.bad_examples), + ) + gradient = extract_answer( + answer=llm_query_fn([request])[0], + tags=FEEDBACK_TAGS, + format_mismatch_label="", + ) + prompt.gradient = gradient + return gradient + + def run( + self, + iteration: int, + prompt: Prompt, + problem_description: str, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Prompt: + """Generate feedback, apply it, and evaluate the mutation. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): prompt to improve. + problem_description (str): description of the target task. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Prompt: evaluated textual-gradient mutation. + """ + + gradient = self._gen_textual_gradient( + prompt=prompt, + problem_description=problem_description, + llm_query_fn=llm_query_fn, + ) + + gradient_step_template = GRADIENT_STEP_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + PROMPT=prompt.text, + TEXTUAL_GRADIENT=gradient, + ) + gradiented = extract_answer( + answer=llm_query_fn([gradient_step_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + gradiented = Prompt(gradiented, origin=PromptOrigin.GRADIENT_STEP) + evaluate_fn(gradiented, "train") + + if self.logger is not None: + self.logger.log_gradient_step( + iteration=iteration, + prompt=prompt.text, + prev_score=prompt.score, + mutated_prompt=gradiented.text, + mutated_score=gradiented.score, + textual_gradient=gradient, + ) + + return gradiented diff --git a/coolprompt/optimizer/brave/operators/hard_few_shot_examples.py b/coolprompt/optimizer/brave/operators/hard_few_shot_examples.py new file mode 100644 index 0000000..cf4a31e --- /dev/null +++ b/coolprompt/optimizer/brave/operators/hard_few_shot_examples.py @@ -0,0 +1,169 @@ +from typing import List, Tuple, Callable + +import numpy as np +from sklearn.feature_extraction.text import TfidfVectorizer +from sklearn.metrics.pairwise import cosine_similarity + +from coolprompt.optimizer.reflective_prompt.prompt import ( + BadExample, + Prompt, + PromptOrigin, +) +from coolprompt.optimizer.brave.operators.few_shot_examples import ( + FewShotExamplesOperator, +) +from coolprompt.optimizer.brave.prompt_templates import ( + FEW_SHOT_EXAMPLES_REMOVING_TEMPLATE, + FEW_SHOT_EXAMPLES_INCORPORATING_TEMPLATE, +) +from coolprompt.optimizer.brave.utils import PROMPT_TAGS +from coolprompt.utils.parsing import extract_answer + + +class HardFewShotExamplesOperator(FewShotExamplesOperator): + """Few-shot operator that selects examples closest to the prompt's + current failure cases instead of sampling uniformly at random. + + When no bad_examples are available (e.g. early in optimization), + falls back to uniform random selection from the parent class. + """ + + def _select_example( + self, + possible_examples: List[Tuple[str, str]], + bad_examples: List[BadExample], + ) -> Tuple[str, str]: + """Pick the candidate whose input is most similar to bad_examples. + + Similarity is TF-IDF cosine, averaged over all bad_example inputs. + Falls back to random if bad_examples is empty or TF-IDF fails. + + Args: + possible_examples (List[Tuple[str, str]]): insertion candidates. + bad_examples (List[BadExample]): current prompt failures. + + Returns: + Tuple[str, str]: selected input-output example. + """ + if not bad_examples or len(possible_examples) == 1: + return possible_examples[np.random.choice(len(possible_examples))] + + candidate_inputs = [ex[0] for ex in possible_examples] + bad_inputs = [be.input for be in bad_examples] + + try: + vectorizer = TfidfVectorizer( + analyzer="word", + ngram_range=(1, 2), + lowercase=True, + min_df=1, + ) + all_texts = candidate_inputs + bad_inputs + tfidf = vectorizer.fit_transform(all_texts) + n = len(candidate_inputs) + sim = cosine_similarity(tfidf[:n], tfidf[n:]) # (n_cands, n_bad) + scores = sim.mean(axis=1) + best_idx = int(np.argmax(scores)) + except Exception: + best_idx = np.random.choice(len(possible_examples)) + + return possible_examples[best_idx] + + def run( + self, + iteration: int, + prompt: Prompt, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Prompt: + """Insert a hard example, rewrite the prompt, and evaluate it. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): prompt whose examples should change. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Prompt: evaluated rewritten prompt, or a zero-scored failure + placeholder when rewriting fails. + """ + + possible_examples = self._filter_possible_examples( + prompt_few_shots=prompt.few_shot_examples + ) + example_to_add = self._select_example(possible_examples, prompt.bad_examples) + + original_few_shots = list(prompt.few_shot_examples) + removed = ("", "") + if len(prompt.few_shot_examples) == self.max_few_shot_examples_num: + ind = np.random.choice(len(prompt.few_shot_examples)) + removed = prompt.few_shot_examples[ind] + prompt.few_shot_examples[ind] = example_to_add + else: + prompt.add_few_shot_example(example_to_add) + + removing_template = FEW_SHOT_EXAMPLES_REMOVING_TEMPLATE.format( + PROMPT=prompt.text + ) + answer_after_removal = llm_query_fn([removing_template])[0] + prompt_without_few_shots = extract_answer( + answer=answer_after_removal, + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + + few_shot_template = FEW_SHOT_EXAMPLES_INCORPORATING_TEMPLATE.format( + PROMPT=prompt_without_few_shots, + EXAMPLES=self._prepare_examples(prompt.few_shot_examples), + ) + try: + answer_with_few_shot = llm_query_fn([few_shot_template])[0] + prompt_with_few_shots = extract_answer( + answer=answer_with_few_shot, + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + except Exception: + prompt_with_few_shots = None + + if prompt_with_few_shots: + mutated_offspring = Prompt( + prompt_with_few_shots, + origin=PromptOrigin.FEW_SHOT, + ) + evaluate_fn(mutated_offspring, "train") + else: + prompt.few_shot_examples = original_few_shots + mutated_offspring = Prompt( + "FAILED TO PRODUCE", + origin=PromptOrigin.FEW_SHOT, + ) + mutated_offspring.set_score(0) + + if self.logger is not None: + self.logger.log_few_shot_mutation( + iteration=iteration, + prompt=prompt.text, + prev_score=prompt.score, + mutated_prompt=mutated_offspring.text, + mutated_score=mutated_offspring.score, + added_few_shot=answer_after_removal, + removed_few_shot=answer_with_few_shot, + file_name="failed_few_shot_mutations", + ) + + if self.logger is not None: + self.logger.log_few_shot_mutation( + iteration=iteration, + prompt=prompt.text, + prev_score=prompt.score, + mutated_prompt=mutated_offspring.text, + mutated_score=mutated_offspring.score, + added_few_shot=example_to_add, + removed_few_shot=removed, + file_name="hard_few_shot_mutations", + ) + + return mutated_offspring diff --git a/coolprompt/optimizer/brave/operators/hype.py b/coolprompt/optimizer/brave/operators/hype.py new file mode 100644 index 0000000..d63e1b8 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/hype.py @@ -0,0 +1,58 @@ +from typing import Callable +from langchain_core.language_models.base import BaseLanguageModel + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.hyper.meta_prompt import MetaPromptOptimizer + + +class HypeOperator(Operator): + """Apply the HYPE prompt optimization strategy.""" + + def __init__(self, model: BaseLanguageModel, **kwargs) -> None: + """Create a HYPE optimizer backed by the supplied model. + + Args: + model (BaseLanguageModel): model used by HYPE. + **kwargs (Any): base-operator arguments such as ``logger``. + """ + + super().__init__(**kwargs) + self.hype = MetaPromptOptimizer(model) + + def run( + self, + iteration: int, + prompt: Prompt, + problem_description: str, + evaluate_fn: Callable[[Prompt, str], None], + ) -> Prompt: + """Optimize and evaluate a prompt with HYPE. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): prompt to optimize. + problem_description (str): description of the target task. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Prompt: evaluated HYPE mutation. + """ + + hyped = self.hype.optimize( + prompt.text, meta_info={"problem_description": problem_description} + ) + hyped = Prompt(hyped, origin=PromptOrigin.HYPE) + evaluate_fn(hyped, "train") + + if self.logger is not None: + self.logger.log_mutation( + iteration=iteration, + prompt=prompt.text, + prev_score=prompt.score, + mutated_prompt=hyped.text, + mutated_score=hyped.score, + file_name="hype", + ) + + return hyped diff --git a/coolprompt/optimizer/brave/operators/initializer.py b/coolprompt/optimizer/brave/operators/initializer.py new file mode 100644 index 0000000..33c0102 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/initializer.py @@ -0,0 +1,67 @@ +from typing import List, Callable +from langchain_core.language_models.base import BaseLanguageModel + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import PROMPT_BY_DESCRIPTION_TEMPLATE +from coolprompt.optimizer.brave.utils import reranking_population, PROMPT_TAGS +from coolprompt.optimizer.hyper.meta_prompt import MetaPromptOptimizer +from coolprompt.utils.parsing import extract_answer + + +class PopulationInitializationOperator(Operator): + """Create BRAVE's initial population from a problem description.""" + + def run( + self, + initial_prompt: str, + problem_description: str, + population_size: int, # just for interface + model: BaseLanguageModel, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> List[Prompt]: + """Generate and evaluate an initial population of diverse prompts. + + Args: + initial_prompt (str): seed prompt included in the population. + problem_description (str): description of the target task. + population_size (int): number of prompts to generate. + model (BaseLanguageModel): model retained for interface parity. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + List[Prompt]: evaluated initial population. + """ + + prompt_by_description_template = PROMPT_BY_DESCRIPTION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description + ) + prompt_by_pd = extract_answer( + answer=llm_query_fn([prompt_by_description_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + prompt_by_pd = Prompt(prompt_by_pd, origin=PromptOrigin.BY_PD) + + hype = MetaPromptOptimizer(model) + prompt_after_hype = hype.optimize( + prompt=initial_prompt, + meta_info={"problem_description": problem_description}, + ) + prompt_after_hype = Prompt(prompt_after_hype, origin=PromptOrigin.HYPE) + + initial_prompt = Prompt(initial_prompt, PromptOrigin.MANUAL) + population = [initial_prompt, prompt_after_hype, prompt_by_pd] + + for prompt in population: + evaluate_fn(prompt, "train") + + population = reranking_population(population) + + if self.logger is not None: + self.logger.log_population(iteration=0, population=population) + + return population diff --git a/coolprompt/optimizer/brave/operators/long_term_mutation.py b/coolprompt/optimizer/brave/operators/long_term_mutation.py new file mode 100644 index 0000000..65c8859 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/long_term_mutation.py @@ -0,0 +1,62 @@ +from typing import List, Tuple, Callable + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import ELITIST_MUTATION_TEMPLATE +from coolprompt.optimizer.brave.utils import PROMPT_TAGS +from coolprompt.utils.parsing import extract_answer + + +class LongTermMutationOperator(Operator): + """Mutate a prompt using the accumulated long-term reflection.""" + + def run( + self, + iteration: int, + prompt: Prompt, + problem_description: str, + long_term_reflection: str, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Tuple[Prompt, str]: + """Generate and evaluate a mutation guided by long-term memory. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): prompt to mutate. + problem_description (str): description of the target task. + long_term_reflection (str): accumulated reflection text. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Tuple[Prompt, str]: evaluated mutation and unchanged reflection. + """ + + mutation_template = ELITIST_MUTATION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + ELITIST_PROMPT=prompt.text, + LONG_TERM_REFLECTION=long_term_reflection, + ) + mutated_offspring = extract_answer( + answer=llm_query_fn([mutation_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + mutated_offspring = Prompt( + mutated_offspring, origin=PromptOrigin.LONG_TERM_MUTATION + ) + evaluate_fn(mutated_offspring, "train") + + if self.logger is not None: + self.logger.log_mutation( + iteration=iteration, + prompt=prompt.text, + prev_score=prompt.score, + mutated_prompt=mutated_offspring.text, + mutated_score=mutated_offspring.score, + file_name="long_term_mutations", + ) + + return mutated_offspring diff --git a/coolprompt/optimizer/brave/operators/paraphrase_initializer.py b/coolprompt/optimizer/brave/operators/paraphrase_initializer.py new file mode 100644 index 0000000..63b9ceb --- /dev/null +++ b/coolprompt/optimizer/brave/operators/paraphrase_initializer.py @@ -0,0 +1,59 @@ +from typing import List, Callable +from langchain_core.language_models.base import BaseLanguageModel + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import PROMPT_BY_DESCRIPTION_TEMPLATE +from coolprompt.optimizer.brave.utils import reranking_population, PROMPT_TAGS +from coolprompt.utils.parsing import extract_answer + + +class ParaphraseInitializationOperator(Operator): + """Initialize a population by paraphrasing a seed prompt.""" + + def run( + self, + initial_prompt: str, + population_size: int, + problem_description: str, + model: BaseLanguageModel, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> List[Prompt]: + """Generate, evaluate, and return paraphrases of a seed prompt. + + Args: + initial_prompt (str): prompt to paraphrase. + population_size (int): number of population members. + problem_description (str): description of the target task. + model (BaseLanguageModel): model retained for interface parity. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + List[Prompt]: evaluated paraphrase population. + """ + + prompt_by_description_template = PROMPT_BY_DESCRIPTION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description + ) + answers = llm_query_fn([prompt_by_description_template] * (population_size - 1)) + prompts = [ + extract_answer(answer=ans, tags=PROMPT_TAGS, format_mismatch_label="") + for ans in answers + ] + prompts = [Prompt(prompt, origin=PromptOrigin.BY_PD) for prompt in prompts] + + initial_prompt = Prompt(initial_prompt, PromptOrigin.MANUAL) + prompts.append(initial_prompt) + + for prompt in prompts: + evaluate_fn(prompt, "train") + + population = reranking_population(prompts) + + if self.logger is not None: + self.logger.log_population(iteration=0, population=population) + + return population diff --git a/coolprompt/optimizer/brave/operators/paraphrasing.py b/coolprompt/optimizer/brave/operators/paraphrasing.py new file mode 100644 index 0000000..32a0e49 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/paraphrasing.py @@ -0,0 +1,58 @@ +from typing import List, Callable + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import ( + PARAPHRASE_BY_DESCRIPTION_TEMPLATE, +) +from coolprompt.optimizer.brave.utils import PROMPT_TAGS +from coolprompt.utils.parsing import extract_answer + + +class ParaphrasingByPDOperator(Operator): + """Paraphrase a prompt in the context of its problem description.""" + + def run( + self, + iteration: int, + prompt: Prompt, + problem_description: str, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Prompt: + """Generate and evaluate a problem-aware paraphrase. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): prompt to paraphrase. + problem_description (str): description of the target task. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Prompt: evaluated paraphrased prompt. + """ + + paraphrasing_template = PARAPHRASE_BY_DESCRIPTION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, PROMPT=prompt.text + ) + paraphrased = extract_answer( + answer=llm_query_fn([paraphrasing_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + paraphrased = Prompt(paraphrased, origin=PromptOrigin.PARAPHRASED) + evaluate_fn(paraphrased, "train") + + if self.logger is not None: + self.logger.log_mutation( + iteration=iteration, + prompt=prompt.text, + prev_score=prompt.score, + mutated_prompt=paraphrased.text, + mutated_score=paraphrased.score, + file_name="paraphrases", + ) + + return paraphrased diff --git a/coolprompt/optimizer/brave/operators/zero_order.py b/coolprompt/optimizer/brave/operators/zero_order.py new file mode 100644 index 0000000..381c565 --- /dev/null +++ b/coolprompt/optimizer/brave/operators/zero_order.py @@ -0,0 +1,56 @@ +from typing import List, Callable + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt, PromptOrigin +from coolprompt.optimizer.brave.operators.basic_operator import Operator +from coolprompt.optimizer.brave.prompt_templates import PROMPT_BY_DESCRIPTION_TEMPLATE +from coolprompt.optimizer.brave.utils import PROMPT_TAGS +from coolprompt.utils.parsing import extract_answer + + +class ZeroOrderMutationOperator(Operator): + """Mutate a prompt without using evaluation-derived gradients.""" + + def run( + self, + iteration: int, + prompt: Prompt, # won't be used, but needed for the interface + problem_description: str, + llm_query_fn: Callable[[List[str]], List[str]], + evaluate_fn: Callable[[Prompt, str], None], + ) -> Prompt: + """Generate and evaluate a zero-order prompt mutation. + + Args: + iteration (int): optimization iteration used for logging. + prompt (Prompt): interface-compatible source prompt; not read. + problem_description (str): description of the target task. + llm_query_fn (Callable[[List[str]], List[str]]): batched LLM + callback. + evaluate_fn (Callable[[Prompt, str], None]): prompt evaluator. + + Returns: + Prompt: evaluated generated prompt. + """ + + generating_template = PROMPT_BY_DESCRIPTION_TEMPLATE.format( + PROBLEM_DESCRIPTION=problem_description, + ) + generated = extract_answer( + answer=llm_query_fn([generating_template])[0], + tags=PROMPT_TAGS, + format_mismatch_label="", + ) + generated = Prompt(generated, origin=PromptOrigin.BY_PD) + evaluate_fn(generated, "train") + + if self.logger is not None: + self.logger.log_mutation( + iteration=iteration, + prompt="", + prev_score=-1.0, + mutated_prompt=generated.text, + mutated_score=generated.score, + file_name="zero_orders", + ) + + return generated diff --git a/coolprompt/optimizer/brave/population_diversity.py b/coolprompt/optimizer/brave/population_diversity.py new file mode 100644 index 0000000..9fa6d42 --- /dev/null +++ b/coolprompt/optimizer/brave/population_diversity.py @@ -0,0 +1,433 @@ +from typing import List, Tuple +import numpy as np +from sklearn.feature_extraction.text import TfidfVectorizer +from sklearn.metrics.pairwise import cosine_similarity +from scipy.cluster.hierarchy import linkage, fcluster +from sentence_transformers import SentenceTransformer + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt + + +class BERTSimilarityComputer: + """Compute semantic similarity using BERT embeddings""" + + def __init__( + self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2" + ) -> None: + """Initialize a model for semantic similarity. + + Args: + model_name (str): SentenceTransformers model identifier. + """ + self.model = SentenceTransformer(model_name) + self.available = True + + def compute_similarity(self, texts: List[str]) -> np.ndarray: + """Compute a semantic similarity matrix using BERT. + + Args: + texts (List[str]): texts to compare. + + Returns: + np.ndarray: pairwise cosine similarities, or ``None`` when fewer + than two texts are supplied or encoding fails. + """ + if not self.available or len(texts) < 2: + return None + + try: + embeddings = self.model.encode(texts, convert_to_numpy=True) + similarity = cosine_similarity(embeddings) + return similarity + except Exception: + return None + + +class PopulationDiversityManager: + """Manages population diversity by removing similar prompts""" + + def __init__( + self, + similarity_threshold: float = 0.85, + max_per_cluster: int = 2, + auto_threshold: bool = True, + target_cluster_count: int = None, + use_hierarchical: bool = True, + use_bert: bool = True, + bert_weight: float = 0.6, + duplicate_threshold: float = 0.95, + ): + """Initialize population-diversity controls. + + Args: + similarity_threshold (float): cosine-similarity clustering cutoff. + max_per_cluster (int): maximum prompts retained per cluster. + auto_threshold (bool): whether to adapt the clustering threshold. + target_cluster_count (int): desired number of clusters. + use_hierarchical (bool): whether to use hierarchical clustering. + use_bert (bool): whether to blend semantic BERT similarity. + bert_weight (float): BERT weight in the hybrid similarity. + duplicate_threshold (float): cutoff for near-duplicate removal. + """ + self.similarity_threshold = similarity_threshold + self.max_per_cluster = max_per_cluster + self.auto_threshold = auto_threshold + self.target_cluster_count = target_cluster_count + self.use_hierarchical = use_hierarchical + self.use_bert = use_bert + self.bert_weight = bert_weight + self.duplicate_threshold = duplicate_threshold + + # TF-IDF for fast syntactic similarity + self.vectorizer = TfidfVectorizer( + analyzer="word", + ngram_range=(1, 2), + lowercase=True, + max_features=500, + min_df=1, + ) + + # BERT for semantic similarity + self.bert_computer = None + if use_bert: + self.bert_computer = BERTSimilarityComputer() + + self.last_filter_report = {} + + def _compute_similarity_matrix(self, prompts: List[Prompt]) -> np.ndarray: + """Compute a hybrid TF-IDF and BERT similarity matrix. + + Args: + prompts (List[Prompt]): prompts to compare. + + Returns: + np.ndarray: square pairwise-similarity matrix. + """ + texts = [p.text for p in prompts] + + if len(texts) < 2: + return np.array([[1.0]]) + + # Compute TF-IDF similarity + try: + tfidf_matrix = self.vectorizer.fit_transform(texts) + tfidf_sim = cosine_similarity(tfidf_matrix) + except ValueError: + tfidf_sim = np.ones((len(texts), len(texts))) + + # Compute BERT similarity if available + if ( + self.use_bert + and self.bert_computer is not None + and self.bert_computer.available + ): + bert_sim = self.bert_computer.compute_similarity(texts) + if bert_sim is not None: + # Hybrid similarity: weighted average + similarity = ( + 1 - self.bert_weight + ) * tfidf_sim + self.bert_weight * bert_sim + return similarity + + return tfidf_sim + + def _adaptive_threshold( + self, similarity_matrix: np.ndarray, population_size: int + ) -> float: + """Choose a threshold near the target number of clusters. + + Args: + similarity_matrix (np.ndarray): pairwise similarities. + population_size (int): number of prompts represented by the matrix. + + Returns: + float: adapted threshold, or the configured fixed threshold. + """ + if population_size < 2 or self.target_cluster_count is None: + return self.similarity_threshold + + # For hierarchical clustering + if self.use_hierarchical: + distance_matrix = 1 - similarity_matrix + # Make distance matrix for clustering + upper_triangle = distance_matrix[np.triu_indices_from(distance_matrix, k=1)] + if len(upper_triangle) == 0: + return self.similarity_threshold + + # Binary search for threshold + low, high = 0.0, 1.0 + best_threshold = self.similarity_threshold + + for _ in range(15): + mid = (low + high) / 2 + # Convert similarity threshold to distance threshold + distance_threshold = 1 - mid + + try: + if len(upper_triangle) > 0: + Z = linkage(upper_triangle, method="complete") + clusters = fcluster(Z, distance_threshold, criterion="distance") + num_clusters = len(np.unique(clusters)) + else: + num_clusters = len(similarity_matrix) + except Exception: + num_clusters = len(similarity_matrix) + + if num_clusters < self.target_cluster_count: + high = mid # Need more clusters, raise threshold + else: + low = mid # Have enough clusters, lower threshold + + if abs(num_clusters - self.target_cluster_count) <= 1: + best_threshold = mid + break + + return best_threshold + else: + return self.similarity_threshold + + def _hierarchical_clustering( + self, similarity_matrix: np.ndarray, threshold: float + ) -> List[List[int]]: + """Cluster prompts with complete-linkage hierarchical clustering. + + Args: + similarity_matrix (np.ndarray): pairwise similarities. + threshold (float): minimum within-cluster similarity. + + Returns: + List[List[int]]: clusters of prompt indices. + """ + distance_matrix = 1 - similarity_matrix + upper_triangle = distance_matrix[np.triu_indices_from(distance_matrix, k=1)] + + if len(upper_triangle) == 0: + return [[i] for i in range(len(similarity_matrix))] + + try: + Z = linkage(upper_triangle, method="complete") + distance_threshold = 1 - threshold + clusters_arr = fcluster(Z, distance_threshold, criterion="distance") + + clusters = {} + for idx, cluster_id in enumerate(clusters_arr): + if cluster_id not in clusters: + clusters[cluster_id] = [] + clusters[cluster_id].append(idx) + + return list(clusters.values()) + except Exception: + return [[i] for i in range(len(similarity_matrix))] + + def _dfs_clustering( + self, similarity_matrix: np.ndarray, threshold: float + ) -> List[List[int]]: + """Cluster prompts as connected components using DFS. + + Args: + similarity_matrix (np.ndarray): pairwise similarities. + threshold (float): minimum edge similarity. + + Returns: + List[List[int]]: clusters of prompt indices. + """ + n = len(similarity_matrix) + visited = [False] * n + clusters = [] + + for i in range(n): + if visited[i]: + continue + + cluster = [] + stack = [i] + while stack: + node = stack.pop() + if visited[node]: + continue + visited[node] = True + cluster.append(node) + + for j in range(n): + if not visited[j] and similarity_matrix[node][j] >= threshold: + stack.append(j) + + clusters.append(sorted(cluster)) + + return clusters + + def _filter_near_duplicates( + self, + population: List[Prompt], + similarity_matrix: np.ndarray, + duplicate_threshold: float = 0.95, + ) -> Tuple[List[Prompt], np.ndarray, np.ndarray]: + """Remove near-duplicates, keeping the best of each group. + + Args: + population (List[Prompt]): prompts sorted by descending score. + similarity_matrix (np.ndarray): pairwise similarities. + duplicate_threshold (float): duplicate-similarity cutoff. + + Returns: + Tuple[List[Prompt], np.ndarray, np.ndarray]: retained prompts, + reduced similarity matrix, and retained original indices. + """ + n = len(population) + visited = [False] * n + kept_indices = [] + + for i in range(n): + if visited[i]: + continue + + duplicates = [i] + for j in range(i + 1, n): + if not visited[j] and similarity_matrix[i][j] >= duplicate_threshold: + duplicates.append(j) + visited[j] = True + + # Keep only best (first in sorted list) + best_idx = duplicates[0] + kept_indices.append(best_idx) + visited[best_idx] = True + + # Create new population and similarity matrix + kept_population = [population[i] for i in kept_indices] + kept_sim_matrix = similarity_matrix[np.ix_(kept_indices, kept_indices)] + + return kept_population, kept_sim_matrix, kept_indices + + def filter_by_diversity( + self, population: List[Prompt], target_population_size: int + ) -> List[Prompt]: + """Filter a population for diversity while retaining strong prompts. + + Args: + population (List[Prompt]): prompts sorted by descending score. + target_population_size (int): maximum returned population size. + + Returns: + List[Prompt]: diverse population sorted by descending score. + """ + if len(population) <= target_population_size: + return population + + # Compute similarity matrix + similarity_matrix = self._compute_similarity_matrix(population) + + # Level 1: Remove near-duplicates (threshold similarity) + (population, similarity_matrix, kept_indices) = self._filter_near_duplicates( + population, similarity_matrix, duplicate_threshold=self.duplicate_threshold + ) + + if len(population) <= target_population_size: + return sorted(population, key=lambda p: p.score, reverse=True) + + # Determine threshold + if self.auto_threshold: + threshold = self._adaptive_threshold( + similarity_matrix, target_population_size + ) + else: + threshold = self.similarity_threshold + + threshold = max( + self.similarity_threshold, min(threshold, self.duplicate_threshold) + ) + + # Cluster prompts + if self.use_hierarchical: + clusters = self._hierarchical_clustering(similarity_matrix, threshold) + else: + clusters = self._dfs_clustering(similarity_matrix, threshold) + + # Select best prompts from each cluster + selected_indices = [] + removed_indices = [] + + for cluster in clusters: + # Keep top max_per_cluster from this cluster + # (already sorted by score) + for idx in cluster[: self.max_per_cluster]: + selected_indices.append(idx) + + # Track removed prompts + for idx in cluster[self.max_per_cluster :]: + removed_indices.append(idx) + + # Store report for logging + self.last_filter_report = { + "threshold": threshold, + "num_clusters": len(clusters), + "num_removed": len(removed_indices), + "removed_indices": removed_indices, + "similarity_matrix": similarity_matrix, + "deduplication_removed": len(kept_indices), + } + + # Sort by original order and select + selected_indices = sorted(selected_indices) + selected_prompts = [population[i] for i in selected_indices] + + # Re-sort by score and trim to target size + selected_prompts = sorted( + selected_prompts, key=lambda p: p.score, reverse=True + )[:target_population_size] + + return selected_prompts + + def maintain_diversity( + self, population: List[Prompt], max_size: int + ) -> List[Prompt]: + """Sort prompts by score and enforce a diverse population bound. + + Args: + population (List[Prompt]): candidate prompts. + max_size (int): maximum returned population size. + + Returns: + List[Prompt]: diverse population sorted by descending score. + """ + if len(population) <= max_size: + return sorted(population, key=lambda p: p.score, reverse=True) + + # First sort by score + sorted_pop = sorted(population, key=lambda p: p.score, reverse=True) + + # Then filter by diversity + diverse_pop = self.filter_by_diversity(sorted_pop, max_size) + + # Final sort by score + return sorted(diverse_pop, key=lambda p: p.score, reverse=True) + + def compute_diversity(self, population: List[Prompt]) -> float: + """Return mean pairwise distance in [0, 1] using TF-IDF only (fast). + + Returns 1.0 (fully diverse) when population has fewer than 2 prompts. + + Args: + population (List[Prompt]): prompts whose diversity to measure. + + Returns: + float: mean pairwise distance, or ``0.5`` if computation fails. + """ + if len(population) < 2: + return 1.0 + texts = [p.text for p in population] + try: + tfidf_matrix = self.vectorizer.fit_transform(texts) + sim = cosine_similarity(tfidf_matrix) + n = len(population) + upper = sim[np.triu_indices(n, k=1)] + return float(np.clip(1.0 - upper.mean(), 0.0, 1.0)) + except Exception: + return 0.5 + + def get_filter_report(self) -> dict: + """Get details about the last filtering operation. + + Returns: + dict: thresholds, clusters, and removed prompt indices. + """ + return self.last_filter_report diff --git a/coolprompt/optimizer/brave/prompt_templates.py b/coolprompt/optimizer/brave/prompt_templates.py new file mode 100644 index 0000000..04cbafd --- /dev/null +++ b/coolprompt/optimizer/brave/prompt_templates.py @@ -0,0 +1,209 @@ +PROMPT_BY_DESCRIPTION_TEMPLATE = """You are an expert in the domain of optimization prompts. Your task is to design prompts that can effectively solve optimization problems. +Here is the description of your problem: {PROBLEM_DESCRIPTION} +Write the best prompt/instruction in order to solve that problem in the most effective way. +Remember to pay attention to all details provided in the description (i.e. constraints, restrictions, input-output formats and etc.) +Output prompt only. +Bracket the final prompt with . +""" +PARAPHRASE_BY_DESCRIPTION_TEMPLATE = """You are an expert in the domain of optimization prompts. Your task is to design prompts that can effectively solve optimization problems. +Here is the description of your problem: {PROBLEM_DESCRIPTION} +Paraphrase the given prompt in order to improve it and to solve that problem in the most effective way. +[Prompt] +{PROMPT} +Remember to pay attention to all details provided in the description (i.e. constraints, restrictions, input-output formats and etc.) +Output prompt only. +Bracket the final prompt with . +""" +TEXTUAL_GRADIENT_TEMPLATE = """You are an expert in the domain of prompt optimization. You can deeply analyze the key properties and effects of every prompt. +You will be given a prompt that was designed for the following problem: {PROBLEM_DESCRIPTION}. +You will also be given a few examples from the dataset where the LLM, guided by the prompt, generated poor answers. +Your goal is to determine the main weaknesses and flaws of the prompt by looking at the examples, model answers, and correct outputs. + +Prompt: {PROMPT} + +Examples: {EXAMPLES} + +Provide detailed feedback on how the given prompt can be improved to achieve the best-quality answers for the given problem description and avoid repeating the same mistakes. +Pay attention to the structure of the prompt, to the input and output formats, to the cohesion and coherence of the instruction. These are the key features of each prompt and should be fixed firstly. +Bracket the final feedback with . +""" +CROSSOVER_TEMPLATE = """You are an expert in the domain of optimization prompts. Your task is to design prompts that can effectively solve optimization problems. +Your response outputs prompt text and nothing else. + +You will be provided the problem description (your prompts must solve it effectively), two parent prompts and the reflection. +The reflections contain the crucial information about strengths and weaknesses of both parent prompts. Use it wisely to create the most effective offspring possible. + +[Problem description] +{PROBLEM_DESCRIPTION} + +[Parent 1] +{PARENT1} +[Parent 2] +{PARENT2} +[Reflection] +{SHORT_TERM_REFLECTION} +[Improved prompt] +Please write an improved prompt, using all the information provided above. +Bracket the final prompt with . +""" +SHORT_TERM_REFLECTION_TEMPLATE = """You are an expert in the domain of optimization prompts. Your task is to give hints to design better prompts. + +Below are two prompts, that were created to solve a specific task. +Here is the problem description of the task, they are trying to solve: +{PROBLEM_DESCRIPTION}. + +For each prompt you are provided with detailed feedback of how this particular prompt can be drastically impoved. +[Prompt 1] +{PROMPT1} +[Prompt 1 feedback] +{FEEDBACK1} + +[Prompt 2] +{PROMPT2} +[Prompt 2 feedback] +{FEEDBACK2} + +Use the information above (don't forget about problem description) wisely to create the distilled hint of how to improve both prompts in a most effective way. +This hint can should manifest all the strengths of both prompts and suggest the corrections of their weaknesses. +Bracket the final hint with . +""" +LONG_TERM_REFLECTION_TEMPLATE = """You are an expert in the domain of optimization prompts. Your task is to give hints to design better prompts. + +Below are some newly gained insights on how you should update your prompts to achieve some big gains in quality. +{SHORT_TERM_REFLECTIONS} + +Write the constructive hint for designing better prompts, based on the provided insights and ideas. +Try to distill all ideas of the provided reflections into one global comprehensive reflection. It may consist of several parts referring to each key idea. +Bracket the final distillation with . +""" +LONG_TERM_REFLECTION_UPDATE_TEMPLATE = """You are an expert in the domain of optimization prompts. Your task is to give hints to design better prompts. + +Below are some newly gained insights on how you should update your prompts to achieve some big gains in quality. +{SHORT_TERM_REFLECTIONS} + +And this is your prior version of the key reflection. +{LONG_TERM_REFLECTION} + +Update your key relfection into newer version by distilling the main ideas from the insights above and combining it with the previous distilled vesrion (your prior version). +Try to distill all ideas of the provided reflections into one global comprehensive reflection. It may consist of several parts referring to each key idea. +Bracket the final distillation with . +""" +ELITIST_MUTATION_TEMPLATE = """You are an expert in the domain of optimization prompts. Your task is to design prompts that can effectively solve optimization problems. +Your response outputs prompt text and nothing else. + +You will be provided the problem description (your prompts must solve it effectively), parent elitist prompt and the prior reflection. +The reflection contains the crucial information about the best approaches in prompt optimization guided for the speicific provided task. Use it wisely to create the most effective offspring possible, based on the current elitist prompt. + +[Problem description] +{PROBLEM_DESCRIPTION} + +[Elitist Prompt] +{ELITIST_PROMPT} + +[Prior Reflection] +{LONG_TERM_REFLECTION} + +[Improved prompt] +Please write a mutated prompt, according to the reflection. +Give the main priority to the provided reflection as it accumulates essential information about correct prompt structure and other different prompt features. + +Output the mutated prompt only. +Bracket the final prompt with . +""" +GRADIENT_STEP_TEMPLATE = """You are an expert in the domain of optimization prompts. Your task is to design prompts that can effectively solve optimization problems. +Your response outputs prompt text and nothing else. + +You will be provided the problem description (your prompts must solve it effectively), parent prompt and its textual gradient. +The textual gradient contains the crucial information about the best approaches in prompt optimization guided for the speicific provided task. Use it wisely to create the most effective offspring possible, based on the current prompt. + +[Problem description] +{PROBLEM_DESCRIPTION} + +[Prompt] +{PROMPT} + +[Textual Gradient] +{TEXTUAL_GRADIENT} + +[Improved prompt] +Please write a mutated prompt, according to the reflection. +Give the main priority to the provided gradient as it accumulates essential information about correct prompt structure and other different prompt features. + +Output the mutated prompt only. +Bracket the final prompt with . +""" +CREATIVE_ZERO_ORDER_MUTATION_TEMPLATE = """Imagine that you're an artist. You can do what you want and how you want. You have the mightiest power of free will. +Below is a task user is needed to solve. +{PROBLEM_DESCRIPTION} + +Think deeply throughout your mind, collect all your pros and cons, your powers and your weaknesses. +Use all your self-reflections and self-knowledge to create a way, a prompt for the solution to help the user. +You do NOT need to solve the task directly. Just think of the right instruction for it. +You can write whatever you want, any wordings and combination of phrases. Use all of your free will to create an unique prompt for effective task solution. +There is no restrictions to the prompt language, except of only one: the prompt you will create must effectively solve the provided task. +Remember, you are an artist! Be creative! Be self-expressing! Open yourself to create the most powerful version of you! +Firstly, write your self-reflections and all of the thoughts in . +Secondly, using all the power of free will and no restriction in formulations write the final prompt in +""" +CREATIVE_STYLE_AND_ROLE_TEMPLATE = """Imagine that you're an artist. You can do what you want and how you want. You have the mightiest power of free will. +Below is a task user is needed to solve. +{PROBLEM_DESCRIPTION} + +Think deeply throughout your mind, collect all your pros and cons, your powers and your weaknesses. +Use all your self-reflections and self-knowledge to create a way, a prompt for the solution to help the user. +You do NOT need to solve the task directly. Just think of the right instruction for it. +You can write whatever you want, any wordings and combination of phrases. Use all of your free will to think of how you can create an unique prompt for effective task solution. +There is no restrictions. +Remember, you are an artist! Be creative! Be self-expressing! Open yourself to create the most powerful version of you! +Firstly, write your self-reflections and all of the thoughts in . +Secondly, using all the power of free will and no restriction in formulations think the style of your future effective prompt at +Trirdly, create the mightiest role for yourself. Describe in details who you must be to create that type of prompt. Who is your muse and your inspiration. You can be literally anyone (or even anything!)! You can take your inspiration from both REAL-WORLD and fictional characters! Enjoy and explore the possibilities! Make the best version of yourself that corresponds with the task! Bracket the role in . +Make your style as detailed as possible. It can contain everything! Use all of your imagination! +""" +CREATIVE_STYLE_ROLE_MUTATION_TEMPLATE = """Imagine that you're an artist. You can do what you want and how you want. You have the mightiest power of free will. +Below is a task you need to solve. +[Problem description] +{PROBLEM_DESCRIPTION} +[Problem description] + +One day you've had the vision of your most suitable role and style for the prompt you need to follow. +You strongly believe that while you are following your role and using that style, it can be the only way to solve the provided task. +This is that style: +[Style] +{STYLE} +[Style] + +This is your mightiest role: +[Role] +{ROLE} +[Role] + +Think deeply throughout your mind, remember that you are the brightest artist of all time! +Collect all your improvisation, imagination and inspiration together and rewrite the given prompt below into the best version for the task following the best suitable style and role. +[Prompt to be rewritten] +{PROMPT} +[Prompt to be rewritten] + +Write the final prompt in +""" +FEW_SHOT_EXAMPLES_REMOVING_TEMPLATE = """Carefully remove all the few-shot examples from the provided prompt below. You need to keep the rest of the prompt structure and it's instruction untouched. +[Prompt] +{PROMPT} +[Prompt] + +Bracket the final prompt with """ +FEW_SHOT_EXAMPLES_INCORPORATING_TEMPLATE = """You are an expert in the domain of optimization prompts. Your task is to design prompts that can effectively solve optimization problems. +Your task is to carefully incorporate provided few-shot examples into the given prompt. +Do NOT change the examples neither create your own. You must use the examples provided below and you must incorporate all of them. +Make sure, that the examples fit perfectly and only improve the previous instruction. + +[Prompt] +{PROMPT} +[Prompt] + +The structure of each example provided below: "Input: input from the dataset\nOutput: the correct output for the provided input" +[Examples] +{EXAMPLES} +[Examples] + +Bracket the final prompt with """ diff --git a/coolprompt/optimizer/brave/run.py b/coolprompt/optimizer/brave/run.py new file mode 100644 index 0000000..6a264cb --- /dev/null +++ b/coolprompt/optimizer/brave/run.py @@ -0,0 +1,143 @@ +"""High-level entry point for BRAVE prompt optimization.""" + +from dataclasses import asdict +from random import sample +from typing import Any, List, Mapping, Optional, Tuple, override + +from langchain_core.language_models.base import BaseLanguageModel + +from coolprompt.data_generator.generator import SyntheticDataGenerator +from coolprompt.evaluator import Evaluator +from coolprompt.optimizer.autoprompting_method import ( + AutoPromptingMethod, + BenchmarkContext, +) +from coolprompt.optimizer.brave.evoluter import BRAVEEvoluter +from coolprompt.optimizer.brave.utils import BRAVEConfig +from coolprompt.utils.logging_config import logger + + +def brave( + model: BaseLanguageModel, + dataset_split: Tuple[List[str], List[str], List[str], List[str]], + evaluator: Evaluator, + problem_description: str, + initial_prompt: str, + config: Optional[BRAVEConfig | Mapping[str, Any]] = None, + seed: int = 19, + verbose: bool = True, + log_dir: Optional[str] = None, + **config_overrides: Any, +) -> str: + """Run BRAVE and return the best prompt on the validation split. + + ``config`` accepts either a :class:`BRAVEConfig` instance or a mapping. + Additional keyword arguments override individual BRAVE configuration + fields, matching the keyword-based API of the other CoolPrompt optimizers. + """ + if config is None: + config_values: dict[str, Any] = {} + elif isinstance(config, BRAVEConfig): + config_values = asdict(config) + elif isinstance(config, Mapping): + config_values = dict(config) + else: + raise TypeError("config must be a BRAVEConfig, a mapping, or None") + + config_values.update(config_overrides) + brave_config = BRAVEConfig(**config_values) + train_data, val_data, train_targets, val_targets = dataset_split + + evoluter = BRAVEEvoluter( + model=model, + evaluator=evaluator, + config=brave_config, + seed=seed, + verbose=verbose, + log_dir=log_dir, + ) + logger.info("Starting BRAVE optimization...") + result = evoluter.optimize( + initial_prompt=initial_prompt, + problem_description=problem_description, + train_data=list(train_data), + train_targets=list(train_targets), + val_data=list(val_data), + val_targets=list(val_targets), + ) + logger.info("BRAVE optimization completed") + return result["best_val_prompt"] + + +class BRAVEMethod(AutoPromptingMethod): + """BRAVE implementation of the shared auto-prompting interface.""" + + @override + def optimize( + self, + model, + initial_prompt, + dataset_split, + evaluator, + problem_description, + **kwargs, + ) -> str: + """Run BRAVE through the shared method interface.""" + kwargs.pop("telemetry_callback", None) + return brave( + model=model, + dataset_split=dataset_split, + evaluator=evaluator, + problem_description=problem_description, + initial_prompt=initial_prompt, + **kwargs, + ) + + @override + def run_configured_benchmark( + self, + ctx: BenchmarkContext, + start_prompt: str, + ) -> str: + """Run BRAVE from a benchmark context.""" + problem_description = ctx.config.get("problem_description") + if problem_description is None: + generator = SyntheticDataGenerator(ctx._system_model) + count = min(5, len(ctx.dataset_split[0])) + indices = sample(range(len(ctx.dataset_split[0])), count) + examples = [ + (ctx.dataset_split[0][index], ctx.dataset_split[2][index]) + for index in indices + ] + labels = generator._extract_labels(ctx.dataset_split[2]) + problem_description = generator._generate_problem_description( + prompt=start_prompt, + examples=examples, + task=ctx.evaluator.task, + labels=labels, + ) + + method_config = dict(ctx.config.get("method", {})) + seed = method_config.pop("seed", 19) + verbose = method_config.pop("verbose", True) + log_dir = method_config.pop("log_dir", method_config.pop("output_path", None)) + return self.optimize( + model=ctx.model, + initial_prompt=start_prompt, + dataset_split=ctx.dataset_split, + evaluator=ctx.evaluator, + problem_description=problem_description, + config=method_config, + seed=seed, + verbose=verbose, + log_dir=log_dir, + ) + + @override + def is_data_driven(self) -> bool: + return True + + @property + @override + def name(self) -> str: + return "brave" diff --git a/coolprompt/optimizer/brave/utils.py b/coolprompt/optimizer/brave/utils.py new file mode 100644 index 0000000..7d7b077 --- /dev/null +++ b/coolprompt/optimizer/brave/utils.py @@ -0,0 +1,158 @@ +from dataclasses import dataclass, field, fields +from pathlib import Path +from typing import Any, Dict, List + +from coolprompt.optimizer.reflective_prompt.prompt import Prompt + +FEEDBACK_TAGS = ("", "") +HINT_TAGS = ("", "") +PROMPT_TAGS = ("", "") +STYLE_TAGS = ("") +ROLE_TAGS = ("", "") + + +@dataclass +class BRAVEConfig: + """Configure BRAVE actions, budgets, sampling, and diversity controls.""" + + actions: List[str] | str = field( + default_factory=lambda: [ + "crossover", + "elitist_mutation", + ] + ) + initial_budget_tokens: float = 200_000.0 + max_steps: int = 1000 + population_size: int = 10 + bad_examples_num: int = 5 + patience_steps: int = 100 + min_improvement: float = 1e-4 + lambda_mean_quality: float = 0.3 + lambda_min_quality: float = 0.3 + max_action_budget_share: float = 0.35 + alpha_roi_ema: float = 0.1 + uncertainty_penalty_beta: float = 0.35 + neural_weight: float = 0.55 + improve_prob_weight: float = 0.6 + kill_switch_min_trials: int = 10 + kill_switch_roi_threshold: float = -0.0002 + kill_switch_base_cooldown: int = 5 + kill_switch_scaling_factor: float = 10.0 + use_neural_bandit: bool = True + neural_hidden_dim: int = 32 + neural_learning_rate: float = 5e-3 + diversity_similarity_threshold: float = 0.80 + diversity_max_per_cluster: int = 2 + diversity_auto_threshold: bool = True + diversity_use_hierarchical: bool = True + diversity_use_bert: bool = True + diversity_bert_weight: float = 0.6 + random_mutation_probability: float = 0.0 + diversity_duplicate_threshold: float = 0.95 + few_shot_examples_max_num: int = 5 + few_shot_examples_from_data_cnt: int = 7 + population_initializer: str = "brave" + population_clusterization: bool = True + initial_population_size: int = 10 + train_batch_size: int = 0 + train_batch_seed: int = 19 + use_stratified_train_batches: bool = True + generation_strata_bins: int = 3 + use_curriculum_batches: bool = False + curriculum_warmup_steps: int = 20 + curriculum_max_alpha: float = 0.6 + val_checkpoint_steps: int = 0 + val_checkpoint_topk: int = 1 + rescore_steps: int = 0 + early_stop: bool = False + + +def _merge_dicts(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]: + """Return a shallow merge in which override values take precedence. + + Args: + base (Dict[str, Any]): base mapping. + override (Dict[str, Any]): replacement values. + + Returns: + Dict[str, Any]: newly merged mapping. + """ + + out = dict(base) + out.update(override) + return out + + +def load_brave_config_from_yaml(path: str, profile: str = "balanced") -> BRAVEConfig: + """Load a BRAVE configuration profile from YAML. + + Values from ``profiles[profile]`` override the file's ``defaults``. + + Args: + path (str): path to the YAML configuration file. + profile (str): profile name within the ``profiles`` mapping. + + Returns: + BRAVEConfig: validated configuration fields from the merged profile. + + Raises: + ImportError: if PyYAML is unavailable. + FileNotFoundError: if ``path`` does not exist. + KeyError: if the requested profile is absent. + """ + try: + import yaml # type: ignore + except Exception as exc: + raise ImportError( + "PyYAML is required to load YAML configs. " + + "Install with: pip install pyyaml" + ) from exc + + cfg_path = Path(path) + if not cfg_path.exists(): + raise FileNotFoundError(f"Config file not found: {cfg_path}") + + with cfg_path.open("r", encoding="utf-8") as f: + raw = yaml.safe_load(f) or {} + + defaults = raw.get("defaults", {}) + profiles = raw.get("profiles", {}) + selected = profiles.get(profile) + if selected is None: + available = ", ".join(sorted(profiles.keys())) + raise KeyError(f"Profile '{profile}' not found. Available: [{available}]") + + merged = _merge_dicts(defaults, selected) + allowed = {x.name for x in fields(BRAVEConfig)} + kwargs = {k: v for k, v in merged.items() if k in allowed} + return BRAVEConfig(**kwargs) + + +@dataclass +class OptimizationLog: + """Store controller and efficiency metrics for one optimization step.""" + + step: int + action: str + score: float + delta_quality: float + cost_tokens: float + cumulative_spent: float + value_per_token: float + useful_operation: bool + controller_diag: Dict[str, float] + remaining_budget: float + best_quality: float + + +def reranking_population(population: List[Prompt]) -> List[Prompt]: + """ + Sorts given population of prompts by their scores in descending order. + + Args: + population (List[Prompt]): population to sort. + + Returns: + List[Prompt]: sorted population. + """ + return list(sorted(population, key=lambda prompt: prompt.score, reverse=True)) diff --git a/coolprompt/optimizer/reflective_prompt/prompt.py b/coolprompt/optimizer/reflective_prompt/prompt.py index 9ed1236..02d7e5f 100644 --- a/coolprompt/optimizer/reflective_prompt/prompt.py +++ b/coolprompt/optimizer/reflective_prompt/prompt.py @@ -1,5 +1,5 @@ from enum import Enum -from typing import Type, List, Dict +from typing import Type, List, Dict, Tuple, Optional class PromptOrigin(Enum): @@ -9,9 +9,20 @@ class PromptOrigin(Enum): """ MANUAL = "manual" - APE = "ape" + BY_PD = "by_pd" EVOLUTED = "evoluted" MUTATED = "mutated" + HYPE = "hype" + APE = "ape" + COMPRESSED = "compressed" + CROSSOVER = "crossover" + ELITIST_MUTATION = "elitist_mutation" + LONG_TERM_MUTATION = "long_term_mutation" + GRADIENT_STEP = "gradient_step" + PARAPHRASED = "paraphrased" + CREATIVE_ZERO_ORDER_PD = "creative_zero_order_pd" + CREATIVE_IN_STYLE_OF = "creative_in_style_of" + FEW_SHOT = "few_shot" @classmethod def from_string(cls: Type["PromptOrigin"], string: str) -> "PromptOrigin": @@ -73,9 +84,12 @@ class Prompt: def __init__( self, text: str, - origin: PromptOrigin = PromptOrigin.EVOLUTED, - score: float = None, - bad_examples: List[BadExample] = [], + origin: PromptOrigin = PromptOrigin.EVOLUTED, + score: float = None, + val_score: float = None, + gradient: str = None, + bad_examples: Optional[List[BadExample]] = None, + few_shot_examples: Optional[List[Tuple[str, str]]] = None, ) -> None: """Prompt class. @@ -91,7 +105,12 @@ def __init__( self.text = text self.origin = origin self.score = score - self.bad_examples = bad_examples + self.bad_examples = bad_examples if bad_examples is not None else [] + self.few_shot_examples = ( + few_shot_examples if few_shot_examples is not None else [] + ) + self.gradient = gradient + self.val_score = val_score def set_score(self, new_score: float) -> None: """Records new prompt evaluation score. @@ -102,6 +121,9 @@ def set_score(self, new_score: float) -> None: self.score = float(new_score) + def set_val_score(self, new_score: float) -> None: + self.val_score = float(new_score) + def set_bad_examples(self, bad_examples: List[Dict[str, str]]) -> None: """Stores provided bad examples.""" @@ -113,6 +135,10 @@ def set_bad_examples(self, bad_examples: List[Dict[str, str]]) -> None: ) for example in bad_examples ] + self.gradient = None + + def add_few_shot_example(self, example: Tuple[str, str]) -> None: + self.few_shot_examples.append(example) def to_dict(self) -> dict: """Creates dictionary representation of prompt. @@ -124,9 +150,11 @@ def to_dict(self) -> dict: result = { "text": self.text, "origin": self.origin.name, - } - if self.score is not None: - result["score"] = self.score + } + if self.score is not None: + result["score"] = self.score + if self.val_score is not None: + result["val_score"] = self.val_score if len(self.bad_examples) > 0: result["bad_examples"] = [ex.to_dict() for ex in self.bad_examples] return result diff --git a/coolprompt/utils/prompt_templates/data_generator_templates.py b/coolprompt/utils/prompt_templates/data_generator_templates.py index 50b1f97..214c790 100644 --- a/coolprompt/utils/prompt_templates/data_generator_templates.py +++ b/coolprompt/utils/prompt_templates/data_generator_templates.py @@ -30,6 +30,29 @@ """ +CLASSIFICATION_PROBLEM_DESCRIPTION_BASED_ON_EXAMPLES_TEMPLATE = """You are an expert in LLM task domain. +You are given a user's prompt and a few examples from a classification dataset. +User created this prompt to solve the classification task represented by given dataset. +Write the detailed problem description for which that prompt was created. Feel free to use provided examples from the dataset to highlight the key features of the task. You can pay attention to answer format, problem's subject and scope and other aspects that may be crucial for better understanding. +Remember, you should provide a very detailed problem description in order to make it understandable and clear as much as possible, but it is very important to make your problem description general and non-specific. Do not highlight the meaning of specific examples, you need to define the meaning of the task as a whole. +Use only textual description. Do not add another data. + +This is a classification task. The model must assign each input to exactly one of the following classes: +{labels} +Make sure to describe what each class represents and how it differs from the others. + +User's prompt: {prompt} + +Examples from dataset: +{examples} + +Provide your answer in JSON format with object with key 'problem_description'. +Output format: +{{ + 'problem_description': "Determined problem description" +}} +""" + PROBLEM_DESCRIPTION_BASED_ON_EXAMPLES_TEMPLATE_OLD = """You are an expert in LLM task domain. You are given a user's prompt and a few examples from problem dataset. User created this prompt to solve the task represented by given dataset. diff --git a/coolprompt/utils/utils.py b/coolprompt/utils/utils.py index 7e893ac..cdb4bc6 100644 --- a/coolprompt/utils/utils.py +++ b/coolprompt/utils/utils.py @@ -1,5 +1,8 @@ from sklearn.model_selection import train_test_split -from typing import Iterable, Tuple +from typing import Iterable, List, Optional, Tuple +import numpy as np + +from coolprompt.utils.enums import Task def get_dataset_split( @@ -7,6 +10,7 @@ def get_dataset_split( target: Iterable[str], validation_size: float, train_as_test: bool, + random_state: Optional[int] = None, ) -> Tuple[Iterable[str], Iterable[str], Iterable[str], Iterable[str]]: """Provides a train/val dataset split. @@ -19,6 +23,8 @@ def get_dataset_split( Provided size of validation subset. train_as_test (bool): Either to use all data for train and validation or split it. + random_state (Optional[int]): + Random seed for reproducibility. Defaults to None. Returns: Tuple[Iterable[str], Iterable[str], Iterable[str], Iterable[str]]: @@ -28,6 +34,64 @@ def get_dataset_split( if train_as_test: return (dataset, dataset, target, target) train_data, val_data, train_targets, val_targets = train_test_split( - dataset, target, test_size=validation_size + dataset, target, test_size=validation_size, random_state=random_state ) return (train_data, val_data, train_targets, val_targets) + + +def get_stratified_dataset_split( + dataset: List[str], + target: List, + validation_size: float, + task: Task, + generation_bins: int = 3, + random_state: int = 42, +) -> Tuple[List, List, List, List]: + """Train/val split with stratification by class label (classification) + or input-length quantile bins (generation). + + Falls back to a plain random split if stratification is not feasible + (e.g. too few samples per stratum). + + Args: + dataset: input texts. + target: ground-truth labels or references. + validation_size: fraction of data for validation. + task: Task.CLASSIFICATION or Task.GENERATION. + generation_bins: number of length quantile bins for generation tasks. + random_state: random seed for reproducibility. + + Returns: + (train_data, val_data, train_targets, val_targets) + """ + dataset = list(dataset) + target = list(target) + + if task == Task.CLASSIFICATION: + stratify_labels = [str(t) for t in target] + else: + lengths = np.array([len(s) for s in dataset], dtype=np.float64) + n_bins = max(int(generation_bins), 2) + quantiles = np.linspace(0.0, 1.0, n_bins + 1)[1:-1] + edges = np.unique(np.quantile(lengths, quantiles)) + stratify_labels = [ + str(int(np.searchsorted(edges, length, side="right"))) for length in lengths + ] + + try: + train_data, val_data, train_targets, val_targets = train_test_split( + dataset, + target, + test_size=validation_size, + stratify=stratify_labels, + random_state=random_state, + ) + except ValueError: + train_data, val_data, train_targets, val_targets = train_test_split( + dataset, + target, + test_size=validation_size, + random_state=random_state, + ) + + return train_data, val_data, train_targets, val_targets diff --git a/coolprompt/utils/var_validation.py b/coolprompt/utils/var_validation.py index 468057b..e159a15 100644 --- a/coolprompt/utils/var_validation.py +++ b/coolprompt/utils/var_validation.py @@ -4,6 +4,7 @@ from langchain_core.language_models.base import BaseLanguageModel from coolprompt.optimizer.autoprompting_method import AutoPromptingMethod +from coolprompt.optimizer.brave import BRAVEMethod from coolprompt.optimizer.distill_prompt import DistillMethod from coolprompt.optimizer.hyper.meta_prompt import HyPERLightMethod from coolprompt.optimizer.hyper.hyper import HyPERMethod @@ -22,6 +23,7 @@ "regps": ReGPSMethod, "compress": CompressorMethod, "rider": RIDERGenesisMethod, + "brave": BRAVEMethod, } diff --git a/docs/API.md b/docs/API.md index 41374d6..83caadb 100644 --- a/docs/API.md +++ b/docs/API.md @@ -55,6 +55,7 @@ Method names accepted by `PromptTuner.run(method=...)`: - `hyper` - iterative HyPER optimizer. Documentation - `regps` - RE-GPS optimizer. - `rider` - RIDER optimizer. Documentation +- `brave` - BRAVE budget-aware evolutionary optimizer. Documentation - `compress` - PromptCompressor. Documentation - `reflective` - legacy ReflectivePrompt. Documentation - `distill` - legacy DistillPrompt. Documentation @@ -65,7 +66,7 @@ Custom methods should implement `AutoPromptingMethod` and can be passed to `Prom ## `method_evaluation/` Benchmark interface for comparing autoprompting methods on dataset/config-based experiments. -`evaluate_method(...)` supports the built-in method names `hyper_light`, `hyper`, `reflective`, `reflectiveprompt`, `distill`, `compress`, `regps`, and `rider`. +`evaluate_method(...)` supports the built-in method names `hyper_light`, `hyper`, `reflective`, `reflectiveprompt`, `distill`, `compress`, `regps`, `rider`, and `brave`. --- ## `data_generator/` and `task_detector/` diff --git a/pyproject.toml b/pyproject.toml index 70ee7a7..5762a0b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -51,10 +51,11 @@ dependencies = [ "langchain-core>=0.3.59", "langchain_huggingface>=0.3.1", "langchain-openai>=0.3.30", - "langdetect>=0.4.31", - "deepeval>=3.7.2", - "transformers<5.0.0" -] + "langdetect>=0.4.31", + "deepeval>=3.7.2", + "transformers<5.0.0", + "sentence-transformers>=3.0.0" +] license = "Apache-2.0" license-files = ["LICEN[CS]E*"] keywords = ["apo", "autoprompting", "llm", "transformers", "prompt-engineering"] diff --git a/requirements.txt b/requirements.txt index d18ee6f..c4f078e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -15,4 +15,5 @@ langchain_huggingface>=0.3.1 langchain-openai>=0.3.30 langdetect>=0.4.31 deepeval>=3.7.2 -transformers<5.0.0 \ No newline at end of file +transformers<5.0.0 +sentence-transformers>=3.0.0 diff --git a/test/coolprompt/data_generator/test_generator.py b/test/coolprompt/data_generator/test_generator.py index 5de6a23..c57948a 100644 --- a/test/coolprompt/data_generator/test_generator.py +++ b/test/coolprompt/data_generator/test_generator.py @@ -177,5 +177,7 @@ def test_generate_dataset_without_problem_description(self): ), (["in"], ["out"], "problem"), ) - self._generate_problem_description_mock.assert_called_once_with("prompt") + self._generate_problem_description_mock.assert_called_once_with( + "prompt", task=Task.GENERATION + ) self._generate_mock.assert_called_once_with(request, schema, "examples") diff --git a/test/coolprompt/optimizer/__init__.py b/test/coolprompt/optimizer/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/coolprompt/optimizer/brave/__init__.py b/test/coolprompt/optimizer/brave/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/test/coolprompt/optimizer/brave/test_config.py b/test/coolprompt/optimizer/brave/test_config.py new file mode 100644 index 0000000..cdc8500 --- /dev/null +++ b/test/coolprompt/optimizer/brave/test_config.py @@ -0,0 +1,63 @@ +import tempfile +import unittest +from pathlib import Path + +from coolprompt.optimizer.brave.utils import ( + BRAVEConfig, + load_brave_config_from_yaml, + reranking_population, +) +from coolprompt.optimizer.reflective_prompt.prompt import Prompt + + +class TestBRAVEConfig(unittest.TestCase): + + def test_default_actions_are_supported_by_evoluter(self): + self.assertEqual( + BRAVEConfig().actions, + ["crossover", "elitist_mutation"], + ) + + def test_yaml_profile_overrides_defaults_and_ignores_unknown_fields(self): + config_text = """ +defaults: + max_steps: 100 + population_size: 8 + unknown_option: ignored +profiles: + fast: + max_steps: 5 + initial_budget_tokens: 1200 +""" + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "brave.yaml" + path.write_text(config_text, encoding="utf-8") + config = load_brave_config_from_yaml(str(path), profile="fast") + + self.assertEqual(config.max_steps, 5) + self.assertEqual(config.population_size, 8) + self.assertEqual(config.initial_budget_tokens, 1200) + self.assertFalse(hasattr(config, "unknown_option")) + + def test_yaml_loader_rejects_unknown_profile(self): + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "brave.yaml" + path.write_text("profiles:\n balanced: {}\n", encoding="utf-8") + + with self.assertRaisesRegex(KeyError, "missing"): + load_brave_config_from_yaml(str(path), profile="missing") + + def test_reranking_population_orders_by_descending_score(self): + population = [ + Prompt("low", score=0.1), + Prompt("high", score=0.9), + Prompt("middle", score=0.5), + ] + + result = reranking_population(population) + + self.assertEqual([prompt.text for prompt in result], ["high", "middle", "low"]) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/coolprompt/optimizer/brave/test_controller.py b/test/coolprompt/optimizer/brave/test_controller.py new file mode 100644 index 0000000..a5d5d5b --- /dev/null +++ b/test/coolprompt/optimizer/brave/test_controller.py @@ -0,0 +1,95 @@ +import unittest +from unittest.mock import patch + +import numpy as np + +from coolprompt.optimizer.brave.controller import EVCController + + +class TestEVCController(unittest.TestCase): + + def _controller(self, **kwargs): + defaults = { + "actions": ["cheap", "expensive"], + "feature_dim": 2, + "use_neural_bandit": False, + "uncertainty_penalty_beta": 0.0, + "max_action_budget_share": 0.5, + "seed": 7, + } + defaults.update(kwargs) + return EVCController(**defaults) + + def test_select_action_excludes_action_above_budget_share(self): + controller = self._controller() + outcomes = { + "cheap": (1.0, 10.0), + "expensive": (10.0, 60.0), + } + + with patch.object( + controller, + "_sample_benefit_cost", + side_effect=lambda action, _: outcomes[action], + ): + action, scores = controller.select_action( + x=np.array([1.0, 0.0]), + remaining_budget_tokens=100, + ) + + self.assertEqual(action, "cheap") + self.assertIn("cheap", scores) + self.assertNotIn("expensive", scores) + + def test_select_action_returns_none_when_every_action_is_unaffordable(self): + controller = self._controller(max_action_budget_share=1.0) + + with patch.object( + controller, + "_sample_benefit_cost", + return_value=(1.0, 101.0), + ): + action, scores = controller.select_action( + x=np.ones(2), + remaining_budget_tokens=100, + ) + + self.assertIsNone(action) + self.assertIsNone(scores) + + def test_update_records_success_and_realized_roi(self): + controller = self._controller(alpha_roi_ema=0.5) + + controller.update( + action="cheap", + x_before=np.array([1.0, 0.0]), + delta_quality=0.2, + actual_cost_tokens=100, + improved=True, + ) + + stats = controller.action_stats["cheap"] + self.assertEqual(stats["trials"], 1.0) + self.assertEqual(stats["success_count"], 1.0) + self.assertAlmostEqual(stats["ema_roi"], 0.001) + self.assertGreater( + controller.benefit_models["cheap"].predictive_mean(np.array([1.0, 0.0])), + 0.0, + ) + + def test_kill_switch_disables_action_with_negative_roi(self): + controller = self._controller( + kill_switch_min_trials=2, + kill_switch_roi_threshold=0.0, + kill_switch_base_cooldown=3, + ) + stats = controller.action_stats["cheap"] + stats["trials"] = 2.0 + stats["ema_roi"] = -0.1 + + self.assertTrue(controller._should_disable("cheap")) + self.assertGreater(stats["disabled_until_step"], 0.0) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/coolprompt/optimizer/brave/test_diversity.py b/test/coolprompt/optimizer/brave/test_diversity.py new file mode 100644 index 0000000..2573ae3 --- /dev/null +++ b/test/coolprompt/optimizer/brave/test_diversity.py @@ -0,0 +1,84 @@ +import unittest + +import numpy as np + +from coolprompt.optimizer.brave.actions import ActionResult +from coolprompt.optimizer.brave.population_diversity import ( + PopulationDiversityManager, +) +from coolprompt.optimizer.reflective_prompt.prompt import Prompt + + +class TestPopulationDiversityManager(unittest.TestCase): + + def setUp(self): + self.manager = PopulationDiversityManager( + use_bert=False, + use_hierarchical=False, + ) + + def test_single_prompt_is_fully_diverse(self): + self.assertEqual( + self.manager.compute_diversity([Prompt("only prompt")]), + 1.0, + ) + + def test_identical_prompts_have_zero_diversity(self): + diversity = self.manager.compute_diversity( + [ + Prompt("same prompt"), + Prompt("same prompt"), + ] + ) + + self.assertAlmostEqual(diversity, 0.0) + + def test_near_duplicate_filter_keeps_highest_ranked_prompt(self): + population = [ + Prompt("best", score=0.9), + Prompt("duplicate", score=0.8), + Prompt("different", score=0.7), + ] + similarity = np.array( + [ + [1.0, 0.99, 0.1], + [0.99, 1.0, 0.1], + [0.1, 0.1, 1.0], + ] + ) + + kept, reduced, indices = self.manager._filter_near_duplicates( + population, + similarity, + duplicate_threshold=0.95, + ) + + self.assertEqual([prompt.text for prompt in kept], ["best", "different"]) + np.testing.assert_array_equal(indices, [0, 2]) + self.assertEqual(reduced.shape, (2, 2)) + + def test_maintain_diversity_sorts_population_within_limit(self): + population = [ + Prompt("low", score=0.1), + Prompt("high", score=0.9), + ] + + result = self.manager.maintain_diversity(population, max_size=2) + + self.assertEqual([prompt.text for prompt in result], ["high", "low"]) + + +class TestActionResult(unittest.TestCase): + + def test_optional_fields_have_independent_defaults(self): + first = ActionResult("a", delta_quality=0.1, cost_tokens=10) + second = ActionResult("b", delta_quality=0.0, cost_tokens=20) + + first.payload["value"] = 1 + + self.assertEqual(second.payload, {}) + self.assertFalse(first.improved) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/coolprompt/optimizer/brave/test_sampling.py b/test/coolprompt/optimizer/brave/test_sampling.py new file mode 100644 index 0000000..8b49157 --- /dev/null +++ b/test/coolprompt/optimizer/brave/test_sampling.py @@ -0,0 +1,103 @@ +import unittest + +import numpy as np + +from coolprompt.optimizer.brave.batch_sampler import ( + CurriculumStratifiedBatchSampler, + StratifiedBatchSampler, +) +from coolprompt.optimizer.brave.bayesian_sampling import ( + BayesianLinearTS, + StateFeaturizer, +) +from coolprompt.optimizer.brave.core_states import OptimizerState +from coolprompt.utils.enums import Task + + +class TestStateFeaturizer(unittest.TestCase): + + def test_transform_includes_interaction_features(self): + state = OptimizerState( + val_quality=0.8, + quality_slope=0.1, + stagnation=0.6, + useless_ops_ratio=0.25, + remaining_budget_ratio=0.5, + epoch_progress=0.4, + population_diversity=0.75, + ) + + features = StateFeaturizer().transform(state) + + np.testing.assert_allclose( + features, [0.8, 0.1, 0.6, 0.25, 0.5, 0.4, 0.3, 0.75, 0.15] + ) + self.assertEqual(features.shape, (StateFeaturizer().dim,)) + + +class TestBayesianLinearTS(unittest.TestCase): + + def test_update_changes_posterior_and_prediction(self): + model = BayesianLinearTS(dim=2, alpha=1.0, sigma2=1.0) + x = np.array([1.0, 0.0]) + + model.update(x, y=2.0) + + np.testing.assert_allclose(model.posterior_mean(), [1.0, 0.0]) + self.assertAlmostEqual(model.predictive_mean(x), 1.0) + self.assertAlmostEqual(model.predictive_std(x), np.sqrt(0.5)) + + +class TestStratifiedBatchSampler(unittest.TestCase): + + def test_classification_sampling_is_balanced_and_deterministic(self): + dataset = [f"sample-{index}" for index in range(8)] + targets = ["a"] * 4 + ["b"] * 4 + sampler = StratifiedBatchSampler( + task=Task.CLASSIFICATION, + batch_size=4, + seed=11, + ) + + first = sampler.sample(dataset, targets, epoch=3) + second = sampler.sample(dataset, targets, epoch=3) + + self.assertEqual(first, second) + self.assertEqual(len(first), 4) + self.assertEqual(len(set(first)), 4) + self.assertEqual([targets[index] for index in first].count("a"), 2) + self.assertEqual([targets[index] for index in first].count("b"), 2) + + def test_empty_and_small_datasets_do_not_require_sampling(self): + sampler = StratifiedBatchSampler(Task.GENERATION, batch_size=5) + + self.assertEqual(sampler.sample([], [], epoch=0), []) + self.assertEqual( + sampler.sample(["a", "b"], ["x", "y"], epoch=0), + [0, 1], + ) + + +class TestCurriculumStratifiedBatchSampler(unittest.TestCase): + + def test_difficulty_and_alpha_are_updated(self): + sampler = CurriculumStratifiedBatchSampler( + task=Task.CLASSIFICATION, + batch_size=2, + total_steps=10, + warmup_steps=2, + max_alpha=0.8, + ) + + sampler.update_difficulties([0, 1], failed_indices=[1]) + + self.assertEqual(sampler._curriculum_alpha(epoch=2), 0.0) + self.assertAlmostEqual(sampler._curriculum_alpha(epoch=6), 0.4) + self.assertAlmostEqual(sampler._curriculum_alpha(epoch=10), 0.8) + self.assertEqual(sampler._difficulty(0), 0.0) + self.assertEqual(sampler._difficulty(1), 1.0) + self.assertEqual(sampler._difficulty(99), 0.5) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/coolprompt/optimizer/test_brave.py b/test/coolprompt/optimizer/test_brave.py new file mode 100644 index 0000000..01651b0 --- /dev/null +++ b/test/coolprompt/optimizer/test_brave.py @@ -0,0 +1,58 @@ +import unittest +from unittest.mock import MagicMock, patch + +from coolprompt.optimizer.brave.run import BRAVEMethod, brave +from coolprompt.optimizer.brave.utils import BRAVEConfig +from coolprompt.utils.var_validation import validate_method + + +class TestBraveIntegration(unittest.TestCase): + + @patch("coolprompt.optimizer.brave.run.BRAVEEvoluter") + def test_brave_returns_best_validation_prompt(self, evoluter_cls): + evoluter = evoluter_cls.return_value + evoluter.optimize.return_value = { + "best_prompt": "training best", + "best_val_prompt": "validation best", + } + dataset_split = (["train"], ["validation"], ["a"], ["b"]) + + result = brave( + model=MagicMock(), + dataset_split=dataset_split, + evaluator=MagicMock(), + problem_description="Classify text", + initial_prompt="Initial prompt", + max_steps=7, + initial_budget_tokens=1234, + ) + + self.assertEqual(result, "validation best") + config = evoluter_cls.call_args.kwargs["config"] + self.assertIsInstance(config, BRAVEConfig) + self.assertEqual(config.max_steps, 7) + self.assertEqual(config.initial_budget_tokens, 1234) + evoluter.optimize.assert_called_once_with( + initial_prompt="Initial prompt", + problem_description="Classify text", + train_data=["train"], + train_targets=["a"], + val_data=["validation"], + val_targets=["b"], + ) + + def test_default_actions_are_executable(self): + self.assertEqual( + BRAVEConfig().actions, + ["crossover", "elitist_mutation"], + ) + + def test_brave_is_a_supported_data_driven_method(self): + method = validate_method("brave") + self.assertIsInstance(method, BRAVEMethod) + self.assertEqual(method.name, "brave") + self.assertTrue(method.is_data_driven()) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/coolprompt/optimizer/test_smoke_optimizer_interface.py b/test/coolprompt/optimizer/test_smoke_optimizer_interface.py index 1b6906d..c5b91eb 100644 --- a/test/coolprompt/optimizer/test_smoke_optimizer_interface.py +++ b/test/coolprompt/optimizer/test_smoke_optimizer_interface.py @@ -3,6 +3,7 @@ import pytest from coolprompt.optimizer.autoprompting_method import AutoPromptingMethod +from coolprompt.optimizer.brave import BRAVEMethod from coolprompt.optimizer.hyper.meta_prompt import HyPERLightMethod from coolprompt.optimizer.rider import RIDERGenesisMethod from coolprompt.utils.var_validation import _METHOD_BY_NAME, validate_method @@ -14,6 +15,9 @@ def test_autoprompting_module_exports(): assert issubclass(RIDERGenesisMethod, AutoPromptingMethod) assert RIDERGenesisMethod().name == "rider" assert RIDERGenesisMethod().is_data_driven() is True + assert issubclass(BRAVEMethod, AutoPromptingMethod) + assert BRAVEMethod().name == "brave" + assert BRAVEMethod().is_data_driven() is True def test_validate_method_string_class_and_instance_equivalent(): @@ -55,6 +59,7 @@ def test_method_by_name_covers_expected_keys(): "regps", "compress", "rider", + "brave", } @@ -64,6 +69,7 @@ def test_method_evaluation_entrypoint(): assert callable(me.evaluate_method) assert hasattr(HyPERLightMethod(), "run") assert "rider" in me._BENCHMARK_IMPL + assert "brave" in me._BENCHMARK_IMPL def test_prompt_tuner_importable():