From be61b2ad1b0d973ce37d29bf137cc38e8bc6eaf8 Mon Sep 17 00:00:00 2001 From: banana-bison Date: Sat, 6 Dec 2025 17:38:27 +0530 Subject: [PATCH 1/2] Add reference doc for databricks.py, avatar.py and simba_utils.py Signed-off-by: banana-bison --- dspy/clients/databricks.py | 38 ++++++++++++++++++++++++++ dspy/predict/avatar/avatar.py | 11 ++++++++ dspy/teleprompt/simba_utils.py | 49 ++++++++++++++++++++++++++++++++++ 3 files changed, 98 insertions(+) diff --git a/dspy/clients/databricks.py b/dspy/clients/databricks.py index 4c5c4db8f3..2e6c79ad3b 100644 --- a/dspy/clients/databricks.py +++ b/dspy/clients/databricks.py @@ -55,6 +55,14 @@ def deploy_finetuned_model( databricks_token: str | None = None, deploy_timeout: int = 900, ): + """ + Args: + model: Name of the model as stored in the databricks store + data_format: One of TrainDataFormat.CHAT or TrainDataFormat.COMPLETION based on what the model's supposed to do + databricks_host: URL of databricks host. Useful in case of hosted runtime + databricks_token: Databrick token + deploy_timeout: Time in seconds. Deployment is cancelled if this time period exceeds. (Default: 900s) + """ workspace_client = _get_workspace_client() model_version = next(workspace_client.model_versions.list(model)).version @@ -172,6 +180,22 @@ def finetune( train_data_format: TrainDataFormat | str | None = "chat", train_kwargs: dict[str, Any] | None = None, ) -> str: + """ + Args: + job (TrainingJobDatabricks): Training job to be finetuned + model: Model name + train_data (list[dict[str, Any]]): Training data + train_data_format: One of "chat" or "completion" (Default: "chat") + train_kwargs: Can take following options: + 1. train_data_path: path to training data in Databricks provider + 2. register_to: Required for finetune on Databricks + 3. databricks_host (Optional): Databricks host URL + 4. databricks_token (Optional): Databricks Token + 5. skip_deploy (Optional): Skip deploying the model. (Default: False) + 6. deploy_timeout (Optional): Time in seconds. (Default: 900s) + Returns: + str: Path to the model in the format "databricks/{job.endpoint_name}" + """ if isinstance(train_data_format, str): if train_data_format == "chat": train_data_format = TrainDataFormat.CHAT @@ -303,6 +327,10 @@ def _create_directory_in_databricks_unity_catalog(w: "WorkspaceClient", databric def _save_data_to_local_file(train_data: list[dict[str, Any]], data_format: TrainDataFormat): + """ + Save train_data to a local file in JSON lines format with follwing filename: "finetuning_{UUID}.jsonl". UUID + here is based on UUID Version 4 of RFC 9562 + """ import uuid file_name = f"finetuning_{uuid.uuid4()}.jsonl" @@ -322,6 +350,11 @@ def _save_data_to_local_file(train_data: list[dict[str, Any]], data_format: Trai def _validate_chat_data(data: dict[str, Any]): + """ + This function raises a ValueError exception if data is malformed in these ways: + 1. Data must be a dict with a 'messages' key + 2. Value of 'messages' key must be a dict with 'role' and 'content' keys + """ if "messages" not in data: raise ValueError( "Each finetuning data must be a dict with a 'messages' key when `task=CHAT_COMPLETION`, but " @@ -344,6 +377,11 @@ def _validate_chat_data(data: dict[str, Any]): def _validate_completion_data(data: dict[str, Any]): + """ + This function raises a ValueError exception if data is malformed in these ways: + 1. Data must be a dict with a 'prompt' key. + 2. Value of 'prompt' key must be a dict with 'response' and 'completion' keys + """ if "prompt" not in data: raise ValueError( "Each finetuning data must be a dict with a 'prompt' key when `task=INSTRUCTION_FINETUNE`, but " diff --git a/dspy/predict/avatar/avatar.py b/dspy/predict/avatar/avatar.py index 53142d21c3..455e101350 100644 --- a/dspy/predict/avatar/avatar.py +++ b/dspy/predict/avatar/avatar.py @@ -20,6 +20,17 @@ def get_number_with_suffix(number: int) -> str: class Avatar(dspy.Module): + """ + Module based on the Avatar Optimizer + + This modules implements the Avatar optimizer (https://arxiv.org/pdf/2406.11200) that allows effective use of tools using contrastive reasoning with batch-wise sampling. See paper for details. + + Args: + signature (Type[dspy.Signature]): The signature of the module. + tools (list[dspy.Tool]): a list of dspy tools to be run + max_iters (Int): The maximum number of iterations to retry code generation and execution. + verbose (bool): Toggle verbosity + """ def __init__( self, signature, diff --git a/dspy/teleprompt/simba_utils.py b/dspy/teleprompt/simba_utils.py index a47c00c821..2a8a7754d3 100644 --- a/dspy/teleprompt/simba_utils.py +++ b/dspy/teleprompt/simba_utils.py @@ -1,3 +1,5 @@ +## Miscellaneous utilities for the simba optimizer + import inspect import logging import textwrap @@ -32,6 +34,21 @@ def prepare_models_for_resampling(program: dspy.Module, n: int, teacher_settings return models def wrap_program(program: dspy.Module, metric: Callable): + """ + Wraps a program into a function that returns a dictionary of various metrics + + Args: + program (dspy.Module): dspy.Module that contains instructions to run the LM. + metric (Callable[str,str]): A function that takes examples from your data and output of the LM and compares them + Returns: + A functions that returns a dict with the following keys when called: { + "prediction", + "trace", + "score", + "example", + "output_metadata", + } + """ def wrapped_program(example): with dspy.context(trace=[]): prediction, trace, score = None, None, 0.0 @@ -71,6 +88,21 @@ def wrapped_program(example): return wrapped_program def append_a_demo(demo_input_field_maxlen): + """ + One of the strategies (chosen at random) of the SIMBA optimizer. The other one is append_a_rule + Args: + demo_input_field_maxlen: Max length of characters in the input field + Returns: + A function of the following form: + Args: + bucket: A list of dictionaries with atleast a "score" key + system: A program selected from a list of programs on the basis of softmax sampling of top k baseline candidates + predictor2name: A dict that maps predictors to their names + name2predictor: A dict that maps names to their predictors + batch_10p_score: 10th percentile score + Returns: + bool: True if demo successful, False if bucket[0] score is less than 10th percentile + """ def append_a_demo_(bucket, system, **kwargs): predictor2name, name2predictor = kwargs["predictor2name"], kwargs["name2predictor"] batch_10p_score = kwargs["batch_10p_score"] @@ -104,6 +136,17 @@ def append_a_demo_(bucket, system, **kwargs): def append_a_rule(bucket, system, **kwargs): + """ + One of the strategies (chosen at random) of the SIMBA optimizer. The other one is append_a_demo + Args: + bucket: A list of dictionaries with atleast a "score" key + system: A program selected from a list of programs on the basis of softmax sampling of top k baseline candidates + predictor2name: A dict that maps predictors to their names + name2predictor: A dict that maps names to their predictors + batch_10p_score: 10th percentile score + Returns: + bool: True if demo successful, False if bucket[0] score is less than 10th percentile or greater than 90th percentile + """ predictor2name = kwargs["predictor2name"] batch_10p_score, batch_90p_score = kwargs["batch_10p_score"], kwargs["batch_90p_score"] prompt_model = kwargs["prompt_model"] or dspy.settings.lm @@ -208,6 +251,12 @@ class OfferFeedback(dspy.Signature): ) def inspect_modules(program): + """ + Args: + program (dspy.Module): A dspy.Module + Returns: + str: A listing of input_fields and output_fields from program.named_predictors + """ separator = "-" * 80 output = [separator] From 7855d3895a3e03a3883bde4b1ddce388c5c273ca Mon Sep 17 00:00:00 2001 From: banana-bison Date: Sun, 12 Apr 2026 12:53:54 +0530 Subject: [PATCH 2/2] address review comments Signed-off-by: banana-bison --- dspy/clients/databricks.py | 16 +--------------- dspy/predict/avatar/avatar.py | 8 ++++---- dspy/teleprompt/simba_utils.py | 12 +++++------- 3 files changed, 10 insertions(+), 26 deletions(-) diff --git a/dspy/clients/databricks.py b/dspy/clients/databricks.py index 2e6c79ad3b..a2f9f06fa8 100644 --- a/dspy/clients/databricks.py +++ b/dspy/clients/databricks.py @@ -59,7 +59,7 @@ def deploy_finetuned_model( Args: model: Name of the model as stored in the databricks store data_format: One of TrainDataFormat.CHAT or TrainDataFormat.COMPLETION based on what the model's supposed to do - databricks_host: URL of databricks host. Useful in case of hosted runtime + databricks_host: URL of databricks host. databricks_token: Databrick token deploy_timeout: Time in seconds. Deployment is cancelled if this time period exceeds. (Default: 900s) """ @@ -327,10 +327,6 @@ def _create_directory_in_databricks_unity_catalog(w: "WorkspaceClient", databric def _save_data_to_local_file(train_data: list[dict[str, Any]], data_format: TrainDataFormat): - """ - Save train_data to a local file in JSON lines format with follwing filename: "finetuning_{UUID}.jsonl". UUID - here is based on UUID Version 4 of RFC 9562 - """ import uuid file_name = f"finetuning_{uuid.uuid4()}.jsonl" @@ -350,11 +346,6 @@ def _save_data_to_local_file(train_data: list[dict[str, Any]], data_format: Trai def _validate_chat_data(data: dict[str, Any]): - """ - This function raises a ValueError exception if data is malformed in these ways: - 1. Data must be a dict with a 'messages' key - 2. Value of 'messages' key must be a dict with 'role' and 'content' keys - """ if "messages" not in data: raise ValueError( "Each finetuning data must be a dict with a 'messages' key when `task=CHAT_COMPLETION`, but " @@ -377,11 +368,6 @@ def _validate_chat_data(data: dict[str, Any]): def _validate_completion_data(data: dict[str, Any]): - """ - This function raises a ValueError exception if data is malformed in these ways: - 1. Data must be a dict with a 'prompt' key. - 2. Value of 'prompt' key must be a dict with 'response' and 'completion' keys - """ if "prompt" not in data: raise ValueError( "Each finetuning data must be a dict with a 'prompt' key when `task=INSTRUCTION_FINETUNE`, but " diff --git a/dspy/predict/avatar/avatar.py b/dspy/predict/avatar/avatar.py index 455e101350..2736023961 100644 --- a/dspy/predict/avatar/avatar.py +++ b/dspy/predict/avatar/avatar.py @@ -33,10 +33,10 @@ class Avatar(dspy.Module): """ def __init__( self, - signature, - tools, - max_iters=3, - verbose=False, + signature: Type[dspy.Signature], + tools: list[dspy.Tool], + max_iters: int = 3, + verbose: bool = False, ): self.signature = ensure_signature(signature) self.input_fields = self.signature.input_fields diff --git a/dspy/teleprompt/simba_utils.py b/dspy/teleprompt/simba_utils.py index 2a8a7754d3..13676aeebc 100644 --- a/dspy/teleprompt/simba_utils.py +++ b/dspy/teleprompt/simba_utils.py @@ -1,5 +1,3 @@ -## Miscellaneous utilities for the simba optimizer - import inspect import logging import textwrap @@ -41,7 +39,7 @@ def wrap_program(program: dspy.Module, metric: Callable): program (dspy.Module): dspy.Module that contains instructions to run the LM. metric (Callable[str,str]): A function that takes examples from your data and output of the LM and compares them Returns: - A functions that returns a dict with the following keys when called: { + A functions that take a dspy.Example as argument and returns a dict with the following keys when called: { "prediction", "trace", "score", @@ -49,7 +47,7 @@ def wrap_program(program: dspy.Module, metric: Callable): "output_metadata", } """ - def wrapped_program(example): + def wrapped_program(example: dspy.Example): with dspy.context(trace=[]): prediction, trace, score = None, None, 0.0 try: @@ -87,7 +85,7 @@ def wrapped_program(example): return wrapped_program -def append_a_demo(demo_input_field_maxlen): +def append_a_demo(demo_input_field_maxlen: int): """ One of the strategies (chosen at random) of the SIMBA optimizer. The other one is append_a_rule Args: @@ -135,7 +133,7 @@ def append_a_demo_(bucket, system, **kwargs): return append_a_demo_ -def append_a_rule(bucket, system, **kwargs): +def append_a_rule(bucket: list[dict], system: dspy.Module, **kwargs): """ One of the strategies (chosen at random) of the SIMBA optimizer. The other one is append_a_demo Args: @@ -250,7 +248,7 @@ class OfferFeedback(dspy.Signature): "like the successful trajectory rather than the lower-scoring trajectory." ) -def inspect_modules(program): +def inspect_modules(program: dspy.Module): """ Args: program (dspy.Module): A dspy.Module