import dsp import dspy from dspy.teleprompt.teleprompt import Teleprompter from dspy.signatures import Signature from dspy.evaluate.evaluate import Evaluate from collections import defaultdict import random from dspy.teleprompt import BootstrapFewShot import numpy as np import optuna import math """ USAGE SUGGESTIONS: The following code can be used to compile a optimized signature teleprompter using the BayesianSignatureOptimizer, and evaluate it on an end task: from dspy.teleprompt import BayesianSignatureOptimizer teleprompter = BayesianSignatureOptimizer(prompt_model=prompt_model, task_model=task_model, metric=metric, n=10, init_temperature=1.0) kwargs = dict(num_threads=NUM_THREADS, display_progress=True, display_table=0) compiled_prompt_opt = teleprompter.compile(program, devset=devset[:DEV_NUM], optuna_trials_num=100, max_bootstrapped_demos=3, max_labeled_demos=5, eval_kwargs=kwargs) eval_score = evaluate(compiled_prompt_opt, devset=evalset[:EVAL_NUM], **kwargs) Note that this teleprompter takes in the following parameters: * prompt_model: The model used for prompt generation. When unspecified, defaults to the model set in settings (ie. dspy.settings.configure(lm=task_model)). * task_model: The model used for prompt generation. When unspecified, defaults to the model set in settings (ie. dspy.settings.configure(lm=task_model)). * metric: The task metric used for optimization. * n: The number of new prompts and sets of fewshot examples to generate and evaluate. Default=10. * init_temperature: The temperature used to generate new prompts. Higher roughly equals more creative. Default=1.0. * verbose: Tells the method whether or not to print intermediate steps. * track_stats: Tells the method whether or not to track statistics about the optimization process. If True, the method will track a dictionary with a key corresponding to the trial number, and a value containing a dict with the following keys: * program: the program being evaluated at a given trial * score: the last average evaluated score for the program * pruned: whether or not this program was pruned This information will be returned as attributes of the best program. """ class BasicGenerateInstruction(Signature): """You are an instruction optimizer for large language models. I will give you a ``signature`` of fields (inputs and outputs) in English. Your task is to propose an instruction that will lead a good language model to perform the task well. Don't be afraid to be creative.""" basic_instruction = dspy.InputField(desc="The initial instructions before optimization") proposed_instruction = dspy.OutputField(desc="The improved instructions for the language model") proposed_prefix_for_output_field = dspy.OutputField(desc="The string at the end of the prompt, which will help the model start solving the task") class BasicGenerateInstructionWithDataObservations(Signature): """You are an instruction optimizer for large language models. I will give you a ``signature`` of fields (inputs and outputs) in English. I will also give you some ``observations`` I have made about the dataset and task. Your task is to propose an instruction that will lead a good language model to perform the task well. Don't be afraid to be creative.""" basic_instruction = dspy.InputField(desc="The initial instructions before optimization") observations = dspy.InputField(desc="Observations about the dataset and task") proposed_instruction = dspy.OutputField(desc="The improved instructions for the language model") proposed_prefix_for_output_field = dspy.OutputField(desc="The string at the end of the prompt, which will help the model start solving the task") class BasicGenerateInstructionWithExamples(dspy.Signature): ("""You are an instruction optimizer for large language models. I will give you a ``signature`` of fields (inputs and outputs) in English. Specifically, I will also provide you with the current ``basic instruction`` that is being used for this task. I will also provide you with some ``examples`` of the expected inputs and outputs. Your task is to propose an instruction that will lead a good language model to perform the task well. Don't be afraid to be creative.""") # attempted_instructions = dspy.InputField(format=str, desc="Previously attempted task instructions, along with their resulting validation score, and an example of the instruction in use on a sample from our dataset.") basic_instruction = dspy.InputField(desc="The initial instructions before optimization") # examples = dspy.InputField(format=dsp.passages2text, desc="Example(s) of the task") examples = dspy.InputField(format=dsp.passages2text, desc="Example(s) of the task") proposed_instruction = dspy.OutputField(desc="The improved instructions for the language model") proposed_prefix_for_output_field = dspy.OutputField(desc="The string at the end of the prompt, which will help the model start solving the task") class BasicGenerateInstructionWithExamplesAndDataObservations(dspy.Signature): ("""You are an instruction optimizer for large language models. I will give you a ``signature`` of fields (inputs and outputs) in English. Specifically, I will also provide you with the current ``basic instruction`` that is being used for this task. I will also provide you with some ``observations`` I have made about the dataset and task, along with some ``examples`` of the expected inputs and outputs. Your task is to propose a new improved instruction and prefix for the output field that will lead a good language model to perform the task well. Don't be afraid to be creative.""") basic_instruction = dspy.InputField(desc="The initial instructions before optimization") observations = dspy.InputField(desc="Observations about the dataset and task") examples = dspy.InputField(format=dsp.passages2text, desc="Example(s) of the task") proposed_instruction = dspy.OutputField(desc="The improved instructions for the language model") proposed_prefix_for_output_field = dspy.OutputField(desc="The string at the end of the prompt, which will help the model start solving the task") class ObservationSummarizer(dspy.Signature): ("""Given a series of observations I have made about my dataset, please summarize them into a brief 2-3 sentence summary which highlights only the most important details.""") observations = dspy.InputField(desc="Observations I have made about my dataset") summary = dspy.OutputField(desc="Two to Three sentence summary of only the most significant highlights of my observations") class DatasetDescriptor(dspy.Signature): ("""Given several examples from a dataset please write observations about trends that hold for most or all of the samples. """ """Some areas you may consider in your observations: topics, content, syntax, conciceness, etc. """ """It will be useful to make an educated guess as to the nature of the task this dataset will enable. Don't be afraid to be creative""") examples = dspy.InputField(desc="Sample data points from the dataset") observations = dspy.OutputField(desc="Somethings that holds true for most or all of the data you observed") class DatasetDescriptorWithPriorObservations(dspy.Signature): ("""Given several examples from a dataset please write observations about trends that hold for most or all of the samples. """ """I will also provide you with a few observations I have already made. Please add your own observations or if you feel the observations are comprehensive say 'COMPLETE' """ """Some areas you may consider in your observations: topics, content, syntax, conciceness, etc. """ """It will be useful to make an educated guess as to the nature of the task this dataset will enable. Don't be afraid to be creative""") examples = dspy.InputField(desc="Sample data points from the dataset") prior_observations = dspy.InputField(desc="Some prior observations I made about the data") observations = dspy.OutputField(desc="Somethings that holds true for most or all of the data you observed or COMPLETE if you have nothing to add") class BayesianSignatureOptimizer(Teleprompter): def __init__(self, prompt_model=None, task_model=None, teacher_settings={}, n=10, metric=None, init_temperature=1.0, verbose=False, track_stats=False, view_data_batch_size=10): self.n = n self.metric = metric self.init_temperature = init_temperature self.prompt_model = prompt_model if prompt_model is not None else dspy.settings.lm self.task_model = task_model if task_model is not None else dspy.settings.lm self.verbose = verbose self.track_stats = track_stats self.teacher_settings = teacher_settings self.view_data_batch_size = view_data_batch_size def _print_full_program(self, program): for i,predictor in enumerate(program.predictors()): if self.verbose: print(f"Predictor {i}") if (hasattr(predictor, 'extended_signature')): if self.verbose: print(f"i: {predictor.extended_signature.instructions}") if self.verbose: print(f"p: {predictor.extended_signature.fields[-1].name}") else: if self.verbose: print(f"i: {predictor.extended_signature1.instructions}") if self.verbose: print(f"p: {predictor.extended_signature1.fields[-1].name}") if self.verbose: print("\n") def _print_model_history(self, model, n=1): if self.verbose: print(f"Model ({model}) History:") model.inspect_history(n=n) def _observe_data(self, trainset): upper_lim = min(len(trainset), self.view_data_batch_size) observation = dspy.Predict(DatasetDescriptor, n=1, temperature=1.0)(examples=(trainset[0:upper_lim].__repr__())) observations = observation["observations"] skips = 0 for b in range(self.view_data_batch_size, len(trainset), self.view_data_batch_size): upper_lim = min(len(trainset), b+self.view_data_batch_size) output = dspy.Predict(DatasetDescriptorWithPriorObservations, n=1, temperature=1.0)(prior_observations=observations, examples=(trainset[b:upper_lim].__repr__())) if len(output["observations"]) >= 8 and output["observations"][:8].upper() == "COMPLETE": skips += 1 if skips >= 5: break continue observations += output["observations"] summary = dspy.Predict(ObservationSummarizer, n=1, temperature=1.0)(observations=observations) return summary.summary def _create_example_string(self, fields, example): # Building the output string output = [] for field in fields: name = field.name separator = field.separator input_variable = field.input_variable # Determine the value from input_data or prediction_data value = example.get(input_variable) # Construct the string for the current field field_str = f"{name}{separator}{value}" output.append(field_str) # Joining all the field strings return '\n'.join(output) def _generate_first_N_candidates(self, module, N, view_data, view_examples, demo_candidates, devset): candidates = {} evaluated_candidates = defaultdict(dict) if view_data: # Create data observations self.observations = None with dspy.settings.context(lm=self.prompt_model): self.observations = self._observe_data(devset).replace("Observations:","").replace("Summary:","") if view_examples: example_sets = {} for predictor in module.predictors(): # Get all augmented examples example_set = {} all_sets_of_examples = demo_candidates[id(predictor)] # Get all generated sets of examples for example_set_i, set_of_examples in enumerate(all_sets_of_examples): if example_set_i != 0: # Skip the no examples case for example in set_of_examples: # Get each individual example in the set if "augmented" in example.keys(): if example["augmented"]: if example_set_i not in example_set: example_set[example_set_i] = [] fields_to_use = predictor.signature.fields input_variable_names = [field.input_variable for field in fields_to_use] example_with_only_signature_fields = {key: value for key, value in example.items() if key in input_variable_names} example_string = self._create_example_string(fields_to_use, example_with_only_signature_fields) example_set[example_set_i].append(example_string) example_sets[id(predictor)] = example_set else: example_set[example_set_i] = [] example_sets[id(predictor)] = example_set # Seed the prompt optimizer zero shot with just the instruction, generate BREADTH new prompts for predictor in module.predictors(): basic_instruction = None basic_prefix = None if (hasattr(predictor, 'extended_signature')): basic_instruction = predictor.extended_signature.instructions basic_prefix = predictor.extended_signature.fields[-1].name else: basic_instruction = predictor.extended_signature1.instructions basic_prefix = predictor.extended_signature1.fields[-1].name with dspy.settings.context(lm=self.prompt_model): # Data & Examples if view_data and view_examples: instruct = None for i in range(1,self.n): new_instruct = dspy.Predict(BasicGenerateInstructionWithExamplesAndDataObservations, n=1, temperature=self.init_temperature)(basic_instruction=basic_instruction, observations=self.observations, examples=example_sets[id(predictor)][i]) if not instruct: instruct = new_instruct else: instruct.completions.proposed_instruction.extend(new_instruct.completions.proposed_instruction) instruct.completions.proposed_prefix_for_output_field.extend(new_instruct.completions.proposed_prefix_for_output_field) # Just data elif view_data: instruct = dspy.Predict(BasicGenerateInstructionWithDataObservations, n=N-1, temperature=self.init_temperature)(basic_instruction=basic_instruction, observations=self.observations) # Just examples elif view_examples: instruct = None for i in range(1,self.n): # Note: skip over the first example set which is empty new_instruct = dspy.Predict(BasicGenerateInstructionWithExamples, n=1, temperature=self.init_temperature)(basic_instruction=basic_instruction, examples=example_sets[id(predictor)][i]) if not instruct: instruct = new_instruct else: instruct.completions.proposed_instruction.extend(new_instruct.completions.proposed_instruction) instruct.completions.proposed_prefix_for_output_field.extend(new_instruct.completions.proposed_prefix_for_output_field) # Neither else: instruct = dspy.Predict(BasicGenerateInstruction, n=N-1, temperature=self.init_temperature)(basic_instruction=basic_instruction) # Add in our initial prompt as a candidate as well instruct.completions.proposed_instruction.insert(0, basic_instruction) instruct.completions.proposed_prefix_for_output_field.insert(0, basic_prefix) candidates[id(predictor)] = instruct.completions evaluated_candidates[id(predictor)] = {} if self.verbose: self._print_model_history(self.prompt_model) return candidates, evaluated_candidates def compile(self, student, *, devset, optuna_trials_num, max_bootstrapped_demos, max_labeled_demos, eval_kwargs, seed=42, view_data=True, view_examples=True): random.seed(seed) # Set up program and evaluation function module = student.deepcopy() evaluate = Evaluate(devset=devset, metric=self.metric, **eval_kwargs) # Generate N few shot example sets demo_candidates = {} for i in range(self.n): if i == 0: # Story empty set of demos as default for index 0 for module_p in module.predictors(): if id(module_p) not in demo_candidates.keys(): demo_candidates[id(module_p)] = [] demo_candidates[id(module_p)].append([]) else: if self.verbose: print(f"Creating basic bootstrap: {i}/{self.n-1}") # Create a new basic bootstrap few - shot program . rng = random.Random(i) shuffled_devset = devset[:] # Create a copy of devset rng.shuffle(shuffled_devset) # Shuffle the copy tp = BootstrapFewShot(metric = self.metric, max_bootstrapped_demos=max_bootstrapped_demos, max_labeled_demos=max_labeled_demos, teacher_settings=self.teacher_settings) candidate_program = tp.compile(student=module.deepcopy(), trainset=shuffled_devset) # Store the candidate demos for module_p, candidate_p in zip(module.predictors(), candidate_program.predictors()): if id(module_p) not in demo_candidates.keys(): demo_candidates[id(module_p)] = [] demo_candidates[id(module_p)].append(candidate_p.demos) # Generate N candidate prompts instruction_candidates, _ = self._generate_first_N_candidates(module, self.n, view_data, view_examples, demo_candidates, devset) # Initialize variables to store the best program and its score best_score = float('-inf') best_program = None trial_num = 0 trial_logs = {} # Define our trial objective def create_objective(baseline_program, instruction_candidates, demo_candidates, evaluate, devset): def objective(trial): nonlocal best_program, best_score, trial_num, trial_logs # Allow access to the outer variables candidate_program = baseline_program.deepcopy() # Suggest the instruction to use for our predictor if self.verbose: print(f"Starting trial num: {trial_num}") trial_logs[trial_num] = {} for p_old, p_new in zip(baseline_program.predictors(), candidate_program.predictors()): # Get instruction candidates for our given predictor p_instruction_candidates = instruction_candidates[id(p_old)] p_demo_candidates = demo_candidates[id(p_old)] # Suggest the index of the instruction candidate to use in our trial instruction_idx = trial.suggest_categorical(f"{id(p_old)}_predictor_instruction",range(len(p_instruction_candidates))) demos_idx = trial.suggest_categorical(f"{id(p_old)}_predictor_demos",range(len(p_demo_candidates))) trial_logs[trial_num][f"{id(p_old)}_predictor_instruction"] = instruction_idx trial_logs[trial_num][f"{id(p_old)}_predictor_demos"] = demos_idx # Get the selected instruction candidate selected_candidate = p_instruction_candidates[instruction_idx] selected_instruction = selected_candidate.proposed_instruction.strip('"').strip() selected_prefix = selected_candidate.proposed_prefix_for_output_field.strip('"').strip() # Use this candidates in our program p_new.extended_signature.instructions = selected_instruction p_new.extended_signature.fields[-1] = p_new.extended_signature.fields[-1]._replace(name=selected_prefix) # Get the selected demos selected_demos = p_demo_candidates[demos_idx] # Use these demos in our program p_new.demos = selected_demos if self.verbose: print("Evaling the following program:") self._print_full_program(candidate_program) trial_logs[trial_num]["program"] = candidate_program # Evaluate with the new prompts total_score = 0 batch_size = 100 num_batches = math.ceil(len(devset) / batch_size) for i in range(num_batches): start_index = i * batch_size end_index = min((i + 1) * batch_size, len(devset)) split_dev = devset[start_index:end_index] split_score = evaluate(candidate_program, devset=split_dev, display_table=0) if self.verbose: print(f"{i}st split score: {split_score}") total_score += split_score * len(split_dev) curr_weighted_avg_score = total_score / min((i+1)*100,len(devset)) if self.verbose: print(f"curr average score: {curr_weighted_avg_score}") trial.report(curr_weighted_avg_score, i) # Handle pruning based on the intermediate value. if trial.should_prune(): if self.verbose: print(f"Optuna decided to prune!") trial_logs[trial_num]["score"] = curr_weighted_avg_score trial_logs[trial_num]["pruned"] = True trial_num += 1 raise optuna.TrialPruned() if self.verbose: print(f"Fully evaled score: {curr_weighted_avg_score}") self._print_model_history(self.task_model, n=1) score = curr_weighted_avg_score trial_logs[trial_num]["score"] = curr_weighted_avg_score trial_logs[trial_num]["pruned"] = False # Update the best program if the current score is better if score > best_score: best_score = score best_program = candidate_program.deepcopy() trial_num += 1 return score return objective # Run the trial objective_function = create_objective(module, instruction_candidates, demo_candidates, evaluate, devset) sampler = optuna.samplers.TPESampler(seed=seed) study = optuna.create_study(direction="maximize", sampler=sampler) score = study.optimize(objective_function, n_trials=optuna_trials_num) if best_program is not None and self.track_stats: best_program.trial_logs = trial_logs return best_program