avoid importing optuna in dspy import (#8258)

Author: TomeHirata

dspy/teleprompt/mipro_optimizer_v2.py

@@ -4,12 +4,11 @@
 import sys
 import textwrap
 import time
+from typing import TYPE_CHECKING
 from collections import defaultdict
 from typing import Any, Callable, Dict, List, Literal, Optional, Tuple
 
 import numpy as np
-import optuna
-from optuna.distributions import CategoricalDistribution
 
 import dspy
 from dspy.evaluate.evaluate import Evaluate
@@ -26,6 +25,9 @@
     set_signature,
 )
 
+if TYPE_CHECKING:
+    import optuna
+
 logger = logging.getLogger(__name__)
 
 # Constants
@@ -496,6 +498,7 @@ def _optimize_prompt_parameters(
         minibatch_full_eval_steps: int,
         seed: int,
     ) -> Optional[Any]:
+        import optuna
         # Run optimization
         optuna.logging.set_verbosity(optuna.logging.WARNING)
         logger.info("==> STEP 3: FINDING OPTIMAL PROMPT PARAMETERS <==")
@@ -727,7 +730,7 @@ def _select_and_insert_instructions_and_demos(
         candidate_program: Any,
         instruction_candidates: Dict[int, List[str]],
         demo_candidates: Optional[List],
-        trial: optuna.trial.Trial,
+        trial: "optuna.trial.Trial",
         trial_logs: Dict,
         trial_num: int,
     ) -> List[str]:
@@ -756,6 +759,7 @@ def _select_and_insert_instructions_and_demos(
         return chosen_params, raw_chosen_params
 
     def _get_param_distributions(self, program, instruction_candidates, demo_candidates):
+        from optuna.distributions import CategoricalDistribution
         param_distributions = {}
 
         for i in range(len(instruction_candidates)):
@@ -780,10 +784,11 @@ def _perform_full_evaluation(
         score_data,
         best_score: float,
         best_program: Any,
-        study: optuna.Study,
+        study: "optuna.Study",
         instruction_candidates: List,
         demo_candidates: List,
     ):
+        import optuna
         logger.info(f"===== Trial {trial_num + 1} / {adjusted_num_trials} - Full Evaluation =====")
 
         # Identify best program to evaluate fully

dspy/teleprompt/teleprompt_optuna.py

@@ -1,5 +1,3 @@
-import optuna
-
 from dspy.evaluate.evaluate import Evaluate
 from dspy.teleprompt.teleprompt import Teleprompter
 
@@ -55,6 +53,7 @@ def objective(self, trial):
         return score
 
     def compile(self, student, *, teacher=None, max_demos, trainset, valset=None):
+        import optuna
         self.trainset = trainset
         self.valset = valset or trainset
         self.student = student.reset_copy()