wandb 0.22.1__py3-none-win_arm64.whl

wandb/integration/diffusers/resolvers/utils.py

@@ -0,0 +1,102 @@
+import inspect
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence
+
+import wandb
+from wandb.util import get_module
+
+if TYPE_CHECKING:
+    np_array = get_module("numpy.array")
+    torch_float_tensor = get_module("torch.FloatTensor")
+
+
+def chunkify(input_list, chunk_size) -> List:
+    chunk_size = max(1, chunk_size)
+    return [
+        input_list[i : i + chunk_size] for i in range(0, len(input_list), chunk_size)
+    ]
+
+
+def get_updated_kwargs(
+    pipeline: Any, args: Sequence[Any], kwargs: Dict[str, Any]
+) -> Dict[str, Any]:
+    pipeline_call_parameters = list(
+        inspect.signature(pipeline.__call__).parameters.items()
+    )
+    for idx, arg in enumerate(args):
+        kwargs[pipeline_call_parameters[idx][0]] = arg
+    for pipeline_parameter in pipeline_call_parameters:
+        if pipeline_parameter[0] not in kwargs:
+            kwargs[pipeline_parameter[0]] = pipeline_parameter[1].default
+    if "generator" in kwargs:
+        generator = kwargs["generator"]
+        kwargs["generator"] = (
+            {
+                "seed": generator.initial_seed(),
+                "device": generator.device,
+                "random_state": generator.get_state().cpu().numpy().tolist(),
+            }
+            if generator is not None
+            else None
+        )
+    if "ip_adapter_image" in kwargs:
+        if kwargs["ip_adapter_image"] is not None:
+            wandb.log({"IP-Adapter-Image": wandb.Image(kwargs["ip_adapter_image"])})
+    return kwargs
+
+
+def postprocess_pils_to_np(image: List) -> "np_array":
+    np = get_module(
+        "numpy",
+        required="Please ensure NumPy is installed. You can run `pip install numpy` to install it.",
+    )
+    return np.stack(
+        [np.transpose(np.array(img).astype("uint8"), axes=(2, 0, 1)) for img in image],
+        axis=0,
+    )
+
+
+def postprocess_np_arrays_for_video(
+    images: List["np_array"], normalize: Optional[bool] = False
+) -> "np_array":
+    np = get_module(
+        "numpy",
+        required="Please ensure NumPy is installed. You can run `pip install numpy` to install it.",
+    )
+    images = [(img * 255).astype("uint8") for img in images] if normalize else images
+    return np.transpose(np.stack((images), axis=0), axes=(0, 3, 1, 2))
+
+
+def decode_sdxl_t2i_latents(pipeline: Any, latents: "torch_float_tensor") -> List:
+    """Decode latents generated by [`diffusers.StableDiffusionXLPipeline`](https://huggingface.co/docs/diffusers/main/en/api/pipelines/stable_diffusion/stable_diffusion_xl#stable-diffusion-xl).
+
+    Args:
+        pipeline: (diffusers.DiffusionPipeline) The Diffusion Pipeline from
+            [`diffusers`](https://huggingface.co/docs/diffusers).
+        latents (torch.FloatTensor): The generated latents.
+
+    Returns:
+        List of `PIL` images corresponding to the generated latents.
+    """
+    torch = get_module(
+        "torch",
+        required="Please ensure PyTorch is installed. You can check out https://pytorch.org/get-started/locally/#start-locally for installation instructions.",
+    )
+    with torch.no_grad():
+        needs_upcasting = (
+            pipeline.vae.dtype == torch.float16 and pipeline.vae.config.force_upcast
+        )
+        if needs_upcasting:
+            pipeline.upcast_vae()
+            latents = latents.to(
+                next(iter(pipeline.vae.post_quant_conv.parameters())).dtype
+            )
+        images = pipeline.vae.decode(
+            latents / pipeline.vae.config.scaling_factor, return_dict=False
+        )[0]
+        if needs_upcasting:
+            pipeline.vae.to(dtype=torch.float16)
+        if pipeline.watermark is not None:
+            images = pipeline.watermark.apply_watermark(images)
+        images = pipeline.image_processor.postprocess(images, output_type="pil")
+        pipeline.maybe_free_model_hooks()
+        return images

wandb/integration/dspy/__init__.py

@@ -0,0 +1,5 @@
+"""W&B DSPy integration package."""
+
+from .dspy import WandbDSPyCallback
+
+__all__ = ["WandbDSPyCallback"]

wandb/integration/dspy/dspy.py

@@ -0,0 +1,422 @@
+"""DSPy ↔ Weights & Biases integration."""
+
+from __future__ import annotations
+
+import logging
+import os
+from collections.abc import Mapping, Sequence
+from typing import Any, Literal
+
+import wandb
+import wandb.util
+from wandb.sdk.wandb_run import Run
+
+dspy = wandb.util.get_module(
+    name="dspy",
+    required=(
+        "To use the W&B DSPy integration you need to have the `dspy` "
+        "python package installed.  Install it with `uv pip install dspy`."
+    ),
+    lazy=False,
+)
+if dspy is not None:
+    assert dspy.__version__ >= "3.0.0", (
+        "DSPy 3.0.0 or higher is required. You have " + dspy.__version__
+    )
+
+
+logger = logging.getLogger(__name__)
+
+
+def _flatten_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+    """Flatten a list of nested row dicts into flat key/value dicts.
+
+    Args:
+        rows (list[dict[str, Any]]): List of nested dictionaries to flatten.
+
+    Returns:
+        list[dict[str, Any]]: List of flattened dictionaries.
+
+    """
+
+    def _flatten(
+        d: dict[str, Any], parent_key: str = "", sep: str = "."
+    ) -> dict[str, Any]:
+        items = []
+        for k, v in d.items():
+            new_key = f"{parent_key}{sep}{k}" if parent_key else k
+            if isinstance(v, dict):
+                items.extend(_flatten(v, new_key, sep=sep).items())
+            else:
+                items.append((new_key, v))
+        return dict(items)
+
+    return [_flatten(row) for row in rows]
+
+
+class WandbDSPyCallback(dspy.utils.BaseCallback):
+    """W&B callback for tracking DSPy evaluation and optimization.
+
+    This callback logs evaluation scores, per-step predictions (optional), and
+    a table capturing the DSPy program signature over time. It can also save
+    the best program as a W&B Artifact for reproducibility.
+
+    Examples:
+        Basic usage within DSPy settings:
+
+        ```python
+        import dspy
+        import wandb
+        from wandb.integration.dspy import WandbDSPyCallback
+
+        with wandb.init(project="dspy-optimization") as run:
+            dspy.settings.callbacks.append(WandbDSPyCallback(run=run))
+            # Run your DSPy optimization/evaluation
+        ```
+    """
+
+    def __init__(self, log_results: bool = True, run: Run | None = None) -> None:
+        """Initialize the callback.
+
+        Args:
+            log_results (bool): Whether to log per-evaluation prediction tables.
+            run (Run | None): Optional W&B run to use. Defaults to the
+                current global run if available.
+
+        Raises:
+            wandb.Error: If no active run is provided or found.
+        """
+        # If no run is provided, use the current global run if available.
+        if run is None:
+            if wandb.run is None:
+                raise wandb.Error(
+                    "You must call `wandb.init()` before instantiating WandbDSPyCallback()."
+                )
+            run = wandb.run
+
+        self.log_results = log_results
+
+        with wandb.wandb_lib.telemetry.context(run=run) as tel:
+            tel.feature.dspy_callback = True
+
+        self._run = run
+        self._did_log_config: bool = False
+        self._program_info: dict[str, Any] = {}
+        self._program_table: wandb.Table | None = None
+        self._row_idx: int = 0
+
+    def _flatten_dict(
+        self, nested: Any, parent_key: str = "", sep: str = "."
+    ) -> dict[str, Any]:
+        """Recursively flatten arbitrarily nested mappings and sequences.
+
+        Args:
+            nested (Any): Nested structure of mappings/lists to flatten.
+            parent_key (str): Prefix to prepend to keys in the flattened output.
+            sep (str): Key separator for nested fields.
+
+        Returns:
+            dict[str, Any]: Flattened dictionary representation.
+        """
+        flat: dict[str, Any] = {}
+
+        def _walk(obj: Any, base: str) -> None:
+            if isinstance(obj, Mapping):
+                for k, v in obj.items():
+                    new_key = f"{base}{sep}{k}" if base else str(k)
+                    _walk(v, new_key)
+            elif isinstance(obj, Sequence) and not isinstance(
+                obj, (str, bytes, bytearray)
+            ):
+                for idx, v in enumerate(obj):
+                    new_key = f"{base}{sep}{idx}" if base else str(idx)
+                    _walk(v, new_key)
+            else:
+                # Base can be empty only if the top-level is a scalar; guard against that.
+                key = base if base else ""
+                if key:
+                    flat[key] = obj
+
+        _walk(nested, parent_key)
+        return flat
+
+    def _extract_fields(self, fields: list[dict[str, Any]]) -> dict[str, str]:
+        """Convert signature fields to a flat mapping of strings.
+
+        Note:
+            The input is expected to be a dict-like mapping from field names to
+            field metadata. Values are stringified for logging.
+
+        Args:
+            fields (list[dict[str, Any]]): Mapping of field name to metadata.
+
+        Returns:
+            dict[str, str]: Mapping of field name to string value.
+        """
+        return {k: str(v) for k, v in fields.items()}
+
+    def _extract_program_info(self, program_obj: Any) -> dict[str, Any]:
+        """Extract signature-related info from a DSPy program.
+
+        Attempts to read the program signature, instructions, input and output
+        fields from a DSPy `Predict` parameter if available.
+
+        Args:
+            program_obj (Any): DSPy program/module instance.
+
+        Returns:
+            dict[str, Any]: Flattened dictionary of signature metadata.
+        """
+        info_dict = {}
+
+        if program_obj is None:
+            return info_dict
+
+        try:
+            sig = next(
+                param.signature
+                for _, param in program_obj.named_parameters()
+                if isinstance(param, dspy.Predict)
+            )
+
+            if getattr(sig, "signature", None):
+                info_dict["signature"] = sig.signature
+            if getattr(sig, "instructions", None):
+                info_dict["instructions"] = sig.instructions
+            if getattr(sig, "input_fields", None):
+                input_fields = sig.input_fields
+                info_dict["input_fields"] = self._extract_fields(input_fields)
+            if getattr(sig, "output_fields", None):
+                output_fields = sig.output_fields
+                info_dict["output_fields"] = self._extract_fields(output_fields)
+
+            return self._flatten_dict(info_dict)
+        except Exception as e:
+            logger.warning(
+                "Failed to extract program info from Evaluate instance: %s", e
+            )
+        return info_dict
+
+    def on_evaluate_start(
+        self,
+        call_id: str,
+        instance: Any,
+        inputs: dict[str, Any],
+    ) -> None:
+        """Handle start of a DSPy evaluation call.
+
+        Logs non-private fields from the evaluator instance to W&B config and
+        captures program signature info for later logging.
+
+        Args:
+            call_id (str): Unique identifier for the evaluation call.
+            instance (Any): The evaluation instance (e.g., `dspy.Evaluate`).
+            inputs (dict[str, Any]): Inputs passed to the evaluation (may
+                include a `program` key with the DSPy program).
+        """
+        if not self._did_log_config:
+            instance_vars = vars(instance) if hasattr(instance, "__dict__") else {}
+            serializable = {
+                k: v for k, v in instance_vars.items() if not k.startswith("_")
+            }
+            if "devset" in serializable:
+                # we don't want to log the devset in the config
+                del serializable["devset"]
+
+            self._run.config.update(serializable)
+            self._did_log_config = True
+
+        # 2) Build/append program signature tables from the 'program' inputs
+        if program_obj := inputs.get("program"):
+            self._program_info = self._extract_program_info(program_obj)
+
+    def on_evaluate_end(
+        self,
+        call_id: str,
+        outputs: Any | None,
+        exception: Exception | None = None,
+    ) -> None:
+        """Handle end of a DSPy evaluation call.
+
+        If available, logs a numeric `score` metric and (optionally) per-step
+        prediction tables. Always appends a row to the program-signature table.
+
+        Args:
+            call_id (str): Unique identifier for the evaluation call.
+            outputs (Any | None): Evaluation outputs; supports
+                `dspy.evaluate.evaluate.EvaluationResult`.
+            exception (Exception | None): Exception raised during evaluation, if any.
+        """
+        # The `BaseCallback` does not define the interface for the `outputs` parameter,
+        # Currently, we know of `EvaluationResult` which is a subclass of `dspy.Prediction`.
+        # We currently support this type and will warn the user if a different type is passed.
+        score: float | None = None
+        if exception is None:
+            if isinstance(outputs, dspy.evaluate.evaluate.EvaluationResult):
+                # log the float score as a wandb metric
+                score = outputs.score
+                wandb.log({"score": float(score)}, step=self._row_idx)
+
+                # Log the predictions as a separate table for each eval end.
+                # We know that results if of type `list[tuple["dspy.Example", "dspy.Example", Any]]`
+                results = outputs.results
+                if self.log_results:
+                    rows = self._parse_results(results)
+                    if rows:
+                        self._log_predictions_table(rows)
+            else:
+                wandb.termwarn(
+                    f"on_evaluate_end received unexpected outputs type: {type(outputs)}. "
+                    "Expected dspy.evaluate.evaluate.EvaluationResult; skipping logging score and `log_results`."
+                )
+        else:
+            wandb.termwarn(
+                f"on_evaluate_end received exception: {exception}. "
+                "Skipping logging score and `log_results`."
+            )
+
+        # Log the program signature iteratively
+        if self._program_table is None:
+            columns = ["step", *self._program_info.keys()]
+            if isinstance(score, float):
+                columns.append("score")
+            self._program_table = wandb.Table(columns=columns, log_mode="INCREMENTAL")
+
+        if self._program_table is not None:
+            values = list(self._program_info.values())
+            if isinstance(score, float):
+                values.append(score)
+
+            self._program_table.add_data(
+                self._row_idx,
+                *values,
+            )
+            self._run.log(
+                {"program_signature": self._program_table}, step=self._row_idx
+            )
+
+        self._row_idx += 1
+
+    def _parse_results(
+        self,
+        results: list[tuple[dspy.Example, dspy.Prediction | dspy.Completions, bool]],
+    ) -> list[dict[str, Any]]:
+        """Normalize evaluation results into serializable row dicts.
+
+        Args:
+            results (list[tuple]): Sequence of `(example, prediction, is_correct)`
+                tuples from DSPy evaluation.
+
+        Returns:
+            list[dict[str, Any]]: Rows with `example`, `prediction`, `is_correct`.
+        """
+        _rows: list[dict[str, Any]] = []
+        for example, prediction, is_correct in results:
+            if isinstance(prediction, dspy.Prediction):
+                prediction_dict = prediction.toDict()
+            if isinstance(prediction, dspy.Completions):
+                prediction_dict = prediction.items()
+
+            row: dict[str, Any] = {
+                "example": example.toDict(),
+                "prediction": prediction_dict,
+                "is_correct": is_correct,
+            }
+            _rows.append(row)
+
+        return _rows
+
+    def _log_predictions_table(self, rows: list[dict[str, Any]]) -> None:
+        """Log a W&B Table of predictions for the current evaluation step.
+
+        Args:
+            rows (list[dict[str, Any]]): Prediction rows to log.
+        """
+        rows = _flatten_rows(rows)
+        columns = list(rows[0].keys())
+
+        data: list[list[Any]] = [list(row.values()) for row in rows]
+
+        preds_table = wandb.Table(columns=columns, data=data, log_mode="IMMUTABLE")
+        self._run.log({f"predictions_{self._row_idx}": preds_table}, step=self._row_idx)
+
+    def log_best_model(
+        self,
+        model: dspy.Module,
+        *,
+        save_program: bool = True,
+        save_dir: str | None = None,
+        filetype: Literal["json", "pkl"] = "json",
+        aliases: Sequence[str] = ("best", "latest"),
+        artifact_name: str = "dspy-program",
+    ) -> None:
+        """Save and log the best DSPy program as a W&B Artifact.
+
+        You can choose to save the full program (architecture + state) or only
+        the state to a single file (JSON or pickle).
+
+        Args:
+            model (dspy.Module): DSPy module to save.
+            save_program (bool): Save full program directory if True; otherwise
+                save only the state file. Defaults to `True`.
+            save_dir (str): Directory to store program files before logging. Defaults to a
+                subdirectory `dspy_program` within the active run's files directory
+                (i.e., `wandb.run.dir`).
+            filetype (Literal["json", "pkl"]): State file format when
+                `save_program` is False. Defaults to `json`.
+            aliases (Sequence[str]): Aliases for the logged Artifact version. Defaults to `("best", "latest")`.
+            artifact_name (str): Base name for the Artifact. Defaults to `dspy-program`.
+
+        Examples:
+            Save the complete program and add aliases:
+
+            ```python
+            callback.log_best_model(
+                optimized_program, save_program=True, aliases=("best", "production")
+            )
+            ```
+
+            Save only the state as JSON:
+
+            ```python
+            callback.log_best_model(
+                optimized_program, save_program=False, filetype="json"
+            )
+            ```
+        """
+        # Derive metadata to help discoverability in the UI
+        info_dict = self._extract_program_info(model)
+        metadata = {
+            "dspy_version": getattr(dspy, "__version__", "unknown"),
+            "module_class": model.__class__.__name__,
+            **info_dict,
+        }
+        artifact = wandb.Artifact(
+            name=f"{artifact_name}-{self._run.id}",
+            type="model",
+            metadata=metadata,
+        )
+
+        # Resolve and normalize the save directory in a cross-platform way
+        if save_dir is None:
+            save_dir = os.path.join(self._run.dir, "dspy_program")
+        save_dir = os.path.normpath(save_dir)
+
+        try:
+            os.makedirs(save_dir, exist_ok=True)
+        except Exception as exc:
+            wandb.termwarn(
+                f"Could not create or access directory '{save_dir}': {exc}. Skipping artifact logging."
+            )
+            return
+        # Save per requested mode
+        if save_program:
+            model.save(save_dir, save_program=True)
+            artifact.add_dir(save_dir)
+        else:
+            filename = f"program.{filetype}"
+            file_path = os.path.join(save_dir, filename)
+            model.save(file_path, save_program=False)
+            artifact.add_file(file_path)
+
+        self._run.log_artifact(artifact, aliases=list(aliases))