Add reference doc for databricks.py, avatar.py and simba_utils.py

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.
+            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

dspy/predict/avatar/avatar.py

@@ -20,12 +20,23 @@ 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,
-        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

dspy/teleprompt/simba_utils.py

@@ -32,7 +32,22 @@ def prepare_models_for_resampling(program: dspy.Module, n: int, teacher_settings
     return models
 
 def wrap_program(program: dspy.Module, metric: Callable):
-    def wrapped_program(example):
+    """
+    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 take a dspy.Example as argument and returns a dict with the following keys when called: {
+            "prediction",
+            "trace",
+            "score",
+            "example",
+            "output_metadata",
+        }
+    """
+    def wrapped_program(example: dspy.Example):
         with dspy.context(trace=[]):
             prediction, trace, score = None, None, 0.0
             try:
@@ -70,7 +85,22 @@ 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:
+        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"]
@@ -103,7 +133,18 @@ 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:
+        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
@@ -207,7 +248,13 @@ 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
+    Returns:
+        str: A listing of input_fields and output_fields from program.named_predictors
+    """
     separator = "-" * 80
     output = [separator]