Typo/nit fixes and add tests for BoostrapWithRandomSearch

dspy/teleprompt/random_search.py

@@ -48,16 +48,10 @@ def __init__(
         self.max_num_samples = max_bootstrapped_demos
         self.max_errors = max_errors
         self.num_candidate_sets = num_candidate_programs
-        # self.max_num_traces = 1 + int(max_bootstrapped_demos / 2.0 * self.num_candidate_sets)
-
-        # Semi-hacky way to get the parent class's _bootstrap function to stop early.
-        # self.max_bootstrapped_demos = self.max_num_traces
         self.max_labeled_demos = max_labeled_demos
 
-        print(
-            "Going to sample between", self.min_num_samples, "and", self.max_num_samples, "traces per predictor.",
-        )
-        print("Will attempt to bootstrap", self.num_candidate_sets, "candidate sets.")
+        print(f"Going to sample between {self.min_num_samples} and {self.max_num_samples} traces per predictor.")
+        print(f"Will attempt to bootstrap {self.num_candidate_sets} candidate sets.")
 
     def compile(self, student, *, teacher=None, trainset, valset=None, restrict=None, labeled_sample=True):
         self.trainset = trainset
@@ -71,36 +65,36 @@ def compile(self, student, *, teacher=None, trainset, valset=None, restrict=None
             if (restrict is not None) and (seed not in restrict):
                 continue
 
-            trainset2 = list(self.trainset)
+            trainset_copy = list(self.trainset)
 
             if seed == -3:
                 # zero-shot
-                program2 = student.reset_copy()
+                program = student.reset_copy()
 
             elif seed == -2:
                 # labels only
                 teleprompter = LabeledFewShot(k=self.max_labeled_demos)
-                program2 = teleprompter.compile(student, trainset=trainset2, sample=labeled_sample)
+                program = teleprompter.compile(student, trainset=trainset_copy, sample=labeled_sample)
 
             elif seed == -1:
                 # unshuffled few-shot
-                program = BootstrapFewShot(
+                optimizer = BootstrapFewShot(
                     metric=self.metric,
                     metric_threshold=self.metric_threshold,
                     max_bootstrapped_demos=self.max_num_samples,
                     max_labeled_demos=self.max_labeled_demos,
                     teacher_settings=self.teacher_settings,
                     max_rounds=self.max_rounds,
                 )
-                program2 = program.compile(student, teacher=teacher, trainset=trainset2)
+                program = optimizer.compile(student, teacher=teacher, trainset=trainset_copy)
 
             else:
                 assert seed >= 0, seed
 
-                random.Random(seed).shuffle(trainset2)
+                random.Random(seed).shuffle(trainset_copy)
                 size = random.Random(seed).randint(self.min_num_samples, self.max_num_samples)
 
-                teleprompter = BootstrapFewShot(
+                optimizer = BootstrapFewShot(
                     metric=self.metric,
                     metric_threshold=self.metric_threshold,
                     max_bootstrapped_demos=size,
@@ -109,7 +103,7 @@ def compile(self, student, *, teacher=None, trainset, valset=None, restrict=None
                     max_rounds=self.max_rounds,
                 )
 
-                program2 = teleprompter.compile(student, teacher=teacher, trainset=trainset2)
+                program = optimizer.compile(student, teacher=teacher, trainset=trainset_copy)
 
             evaluate = Evaluate(
                 devset=self.valset,
@@ -120,39 +114,25 @@ def compile(self, student, *, teacher=None, trainset, valset=None, restrict=None
                 display_progress=True,
             )
 
-            score, subscores = evaluate(program2, return_all_scores=True)
+            score, subscores = evaluate(program, return_all_scores=True)
 
             all_subscores.append(subscores)
 
             ############ Assertion-aware Optimization ############
-            if hasattr(program2, "_suggest_failures"):
-                score = score - program2._suggest_failures * 0.2
-            if hasattr(program2, "_assert_failures"):
-                score = 0 if program2._assert_failures > 0 else score
+            if hasattr(program, "_suggest_failures"):
+                score = score - program._suggest_failures * 0.2
+            if hasattr(program, "_assert_failures"):
+                score = 0 if program._assert_failures > 0 else score
             ######################################################
 
-            print("Score:", score, "for set:", [len(predictor.demos) for predictor in program2.predictors()])
-
             if len(scores) == 0 or score > max(scores):
-                print("New best sscore:", score, "for seed", seed)
-                best_program = program2
+                print("New best score:", score, "for seed", seed)
+                best_program = program
 
             scores.append(score)
-            print(f"Scores so far: {scores}")
-
-            print("Best score:", max(scores))
-
-            score_data.append((score, subscores, seed, program2))
-
-            if len(score_data) > 2:  # We check if there are at least 3 scores to consider
-                for k in [1, 2, 3, 5, 8, 9999]:
-                    top_3_scores = sorted(score_data, key=lambda x: x[0], reverse=True)[:k]
-
-                    # Transpose the subscores to get max per entry and then calculate their average
-                    transposed_subscores = zip(*[subscores for _, subscores, *_ in top_3_scores if subscores])
-                    avg_of_max_per_entry = sum(max(entry) for entry in transposed_subscores) / len(top_3_scores[0][1])
+            print(f"Best score so far: {max(scores)}")
 
-                    print(f"Average of max per entry across top {k} scores: {avg_of_max_per_entry}")
+            score_data.append((score, subscores, seed, program))
 
             if self.stop_at_score is not None and score >= self.stop_at_score:
                 print(f"Stopping early because score {score} is >= stop_at_score {self.stop_at_score}")

tests/teleprompt/test_random_search.py

@@ -0,0 +1,37 @@
+import dspy
+from dspy.predict import Predict
+from dspy.utils.dummies import DummyLM
+from dspy import Example
+from dspy.teleprompt import BootstrapFewShotWithRandomSearch
+
+class SimpleModule(dspy.Module):
+    def __init__(self, signature):
+        super().__init__()
+        self.predictor = Predict(signature)
+
+    def forward(self, **kwargs):
+        return self.predictor(**kwargs)
+
+def simple_metric(example, prediction, trace=None):
+    return example.output == prediction.output
+
+def test_basic_workflow():
+    """Test to ensure the basic compile flow runs without errors."""
+    student = SimpleModule("input -> output")
+    teacher = SimpleModule("input -> output")
+
+    lm = DummyLM(
+        [
+            "Initial thoughts",
+            "Finish[blue]",  # Expected output for both training and validation
+        ]
+    )
+    dspy.settings.configure(lm=lm)
+
+    optimizer = BootstrapFewShotWithRandomSearch(metric=simple_metric, max_bootstrapped_demos=1, max_labeled_demos=1)
+    trainset = [
+        Example(input="What is the color of the sky?", output="blue").with_inputs("input"),
+        Example(input="What does the fox say?", output="Ring-ding-ding-ding-dingeringeding!").with_inputs("input"),
+    ]
+    optimizer.compile(student, teacher=teacher, trainset=trainset)
+