diff --git a/dspy/teleprompt/mipro_optimizer_v2.py b/dspy/teleprompt/mipro_optimizer_v2.py index c96d71d3f5..63b824ef03 100644 --- a/dspy/teleprompt/mipro_optimizer_v2.py +++ b/dspy/teleprompt/mipro_optimizer_v2.py @@ -2,9 +2,7 @@ import random from collections import defaultdict from typing import TYPE_CHECKING, Any, Callable, Literal - import numpy as np - import dspy from dspy.evaluate.evaluate import Evaluate from dspy.propose import GroundedProposer @@ -23,6 +21,8 @@ if TYPE_CHECKING: import optuna +from optuna.samplers import BaseSampler + logger = logging.getLogger(__name__) # Constants @@ -114,6 +114,8 @@ def compile( fewshot_aware_proposer: bool = True, requires_permission_to_run: bool | None = None, # deprecated provide_traceback: bool | None = None, + sampler: BaseSampler | None = None, + n_jobs: int = 1 ) -> Any: if requires_permission_to_run == False: logger.warning( @@ -257,6 +259,8 @@ def compile( minibatch_size, minibatch_full_eval_steps, seed, + sampler, + n_jobs ) return best_program @@ -492,6 +496,8 @@ def _propose_instructions( return instruction_candidates + + def _optimize_prompt_parameters( self, program: Any, @@ -504,6 +510,8 @@ def _optimize_prompt_parameters( minibatch_size: int, minibatch_full_eval_steps: int, seed: int, + sampler: BaseSampler, + n_jobs: int ) -> Any | None: import optuna @@ -643,8 +651,16 @@ def objective(trial): return score - sampler = optuna.samplers.TPESampler(seed=seed, multivariate=True) - study = optuna.create_study(direction="maximize", sampler=sampler) + if not sampler: + _sampler = optuna.samplers.TPESampler(seed=seed, multivariate=True) + elif sampler.__class__.__name__ == 'qEISampler': + # batch Expected Improvement Sampler; allows parallelization in optuna + _sampler = sampler + _sampler.create_search_space(instruction_candidates, demo_candidates) + else: + _sampler = sampler + + study = optuna.create_study(direction="maximize", sampler=_sampler) default_params = {f"{i}_predictor_instruction": 0 for i in range(len(program.predictors()))} if demo_candidates: @@ -657,7 +673,8 @@ def objective(trial): value=default_score, ) study.add_trial(trial) - study.optimize(objective, n_trials=num_trials) + study.optimize(objective, n_trials=num_trials, n_jobs=n_jobs) + self.study = study # Attach logs to best program if best_program is not None and self.track_stats: @@ -766,7 +783,8 @@ def _select_and_insert_instructions_and_demos( predictor.demos = demo_candidates[i][demos_idx] trial_logs[trial_num][f"{i}_predictor_demos"] = demos_idx chosen_params.append(f"Predictor {i}: Few-Shot Set {demos_idx}") - raw_chosen_params[f"{i}_predictor_demos"] = instruction_idx + raw_chosen_params[f"{i}_predictor_demos"] = demos_idx +# raw_chosen_params[f"{i}_predictor_demos"] = instruction_idx return chosen_params, raw_chosen_params