wandb 0.21.2__py3-none-macosx_12_0_arm64.whl

wandb/integration/catboost/catboost.py

@@ -0,0 +1,182 @@
+"""catboost init."""
+
+from pathlib import Path
+from types import SimpleNamespace
+from typing import List, Union
+
+from catboost import CatBoostClassifier, CatBoostRegressor  # type: ignore
+
+import wandb
+from wandb.sdk.lib import telemetry as wb_telemetry
+
+
+class WandbCallback:
+    """`WandbCallback` automatically integrates CatBoost with wandb.
+
+    Args:
+        - metric_period: (int) if you are passing `metric_period` to your CatBoost model please pass the same value here (default=1).
+
+    Passing `WandbCallback` to CatBoost will:
+    - log training and validation metrics at every `metric_period`
+    - log iteration at every `metric_period`
+
+    Example:
+        ```
+        train_pool = Pool(
+            train[features], label=train["label"], cat_features=cat_features
+        )
+        test_pool = Pool(test[features], label=test["label"], cat_features=cat_features)
+
+        model = CatBoostRegressor(
+            iterations=100,
+            loss_function="Cox",
+            eval_metric="Cox",
+        )
+
+        model.fit(
+            train_pool,
+            eval_set=test_pool,
+            callbacks=[WandbCallback()],
+        )
+        ```
+    """
+
+    def __init__(self, metric_period: int = 1):
+        if wandb.run is None:
+            raise wandb.Error("You must call `wandb.init()` before `WandbCallback()`")
+
+        with wb_telemetry.context() as tel:
+            tel.feature.catboost_wandb_callback = True
+
+        self.metric_period: int = metric_period
+
+    def after_iteration(self, info: SimpleNamespace) -> bool:
+        if info.iteration % self.metric_period == 0:
+            for data, metric in info.metrics.items():
+                for metric_name, log in metric.items():
+                    # todo: replace with wandb.run._log once available
+                    wandb.log({f"{data}-{metric_name}": log[-1]}, commit=False)
+            # todo: replace with wandb.run._log once available
+            wandb.log({f"iteration@metric-period-{self.metric_period}": info.iteration})
+
+        return True
+
+
+def _checkpoint_artifact(
+    model: Union[CatBoostClassifier, CatBoostRegressor], aliases: List[str]
+) -> None:
+    """Upload model checkpoint as W&B artifact."""
+    if wandb.run is None:
+        raise wandb.Error(
+            "You must call `wandb.init()` before `_checkpoint_artifact()`"
+        )
+
+    model_name = f"model_{wandb.run.id}"
+    # save the model in the default `cbm` format
+    model_path = Path(wandb.run.dir) / "model"
+
+    model.save_model(model_path)
+
+    model_artifact = wandb.Artifact(name=model_name, type="model")
+    model_artifact.add_file(str(model_path))
+    wandb.log_artifact(model_artifact, aliases=aliases)
+
+
+def _log_feature_importance(
+    model: Union[CatBoostClassifier, CatBoostRegressor],
+) -> None:
+    """Log feature importance with default settings."""
+    if wandb.run is None:
+        raise wandb.Error(
+            "You must call `wandb.init()` before `_checkpoint_artifact()`"
+        )
+
+    feat_df = model.get_feature_importance(prettified=True)
+
+    fi_data = [
+        [feat, feat_imp]
+        for feat, feat_imp in zip(feat_df["Feature Id"], feat_df["Importances"])
+    ]
+    table = wandb.Table(data=fi_data, columns=["Feature", "Importance"])
+    # todo: replace with wandb.run._log once available
+    wandb.log(
+        {
+            "Feature Importance": wandb.plot.bar(
+                table, "Feature", "Importance", title="Feature Importance"
+            )
+        },
+        commit=False,
+    )
+
+
+def log_summary(
+    model: Union[CatBoostClassifier, CatBoostRegressor],
+    log_all_params: bool = True,
+    save_model_checkpoint: bool = False,
+    log_feature_importance: bool = True,
+) -> None:
+    """`log_summary` logs useful metrics about catboost model after training is done.
+
+    Args:
+        model: it can be CatBoostClassifier or CatBoostRegressor.
+        log_all_params: (boolean) if True (default) log the model hyperparameters as W&B config.
+        save_model_checkpoint: (boolean) if True saves the model upload as W&B artifacts.
+        log_feature_importance: (boolean) if True (default) logs feature importance as W&B bar chart using the default setting of `get_feature_importance`.
+
+    Using this along with `wandb_callback` will:
+
+    - save the hyperparameters as W&B config,
+    - log `best_iteration` and `best_score` as `wandb.summary`,
+    - save and upload your trained model to Weights & Biases Artifacts (when `save_model_checkpoint = True`)
+    - log feature importance plot.
+
+    Example:
+        ```python
+        train_pool = Pool(
+            train[features], label=train["label"], cat_features=cat_features
+        )
+        test_pool = Pool(test[features], label=test["label"], cat_features=cat_features)
+
+        model = CatBoostRegressor(
+            iterations=100,
+            loss_function="Cox",
+            eval_metric="Cox",
+        )
+
+        model.fit(
+            train_pool,
+            eval_set=test_pool,
+            callbacks=[WandbCallback()],
+        )
+
+        log_summary(model)
+        ```
+    """
+    if wandb.run is None:
+        raise wandb.Error("You must call `wandb.init()` before `log_summary()`")
+
+    if not (isinstance(model, (CatBoostClassifier, CatBoostRegressor))):
+        raise wandb.Error(
+            "Model should be an instance of CatBoostClassifier or CatBoostRegressor"
+        )
+
+    with wb_telemetry.context() as tel:
+        tel.feature.catboost_log_summary = True
+
+    # log configs
+    params = model.get_all_params()
+    if log_all_params:
+        wandb.config.update(params)
+
+    # log best score and iteration
+    wandb.run.summary["best_iteration"] = model.get_best_iteration()
+    wandb.run.summary["best_score"] = model.get_best_score()
+
+    # log model
+    if save_model_checkpoint:
+        aliases = ["best"] if params["use_best_model"] else ["last"]
+        _checkpoint_artifact(model, aliases=aliases)
+
+    # Feature importance
+    if log_feature_importance:
+        _log_feature_importance(model)

wandb/integration/cohere/__init__.py

@@ -0,0 +1,3 @@
+__all__ = ("autolog",)
+
+from .cohere import autolog

wandb/integration/cohere/cohere.py

@@ -0,0 +1,21 @@
+import logging
+
+from wandb.sdk.integration_utils.auto_logging import AutologAPI
+
+from .resolver import CohereRequestResponseResolver
+
+logger = logging.getLogger(__name__)
+
+
+autolog = AutologAPI(
+    name="Cohere",
+    symbols=(
+        "Client.generate",
+        "Client.chat",
+        "Client.classify",
+        "Client.summarize",
+        "Client.rerank",
+    ),
+    resolver=CohereRequestResponseResolver(),
+    telemetry_feature="cohere_autolog",
+)

wandb/integration/cohere/resolver.py

@@ -0,0 +1,347 @@
+import logging
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import wandb
+from wandb.sdk.integration_utils.auto_logging import Response
+from wandb.sdk.lib.runid import generate_id
+
+logger = logging.getLogger(__name__)
+
+
+def subset_dict(
+    original_dict: Dict[str, Any], keys_subset: Sequence[str]
+) -> Dict[str, Any]:
+    """Create a subset of a dictionary using a subset of keys.
+
+    :param original_dict: The original dictionary.
+    :param keys_subset: The subset of keys to extract.
+    :return: A dictionary containing only the specified keys.
+    """
+    return {key: original_dict[key] for key in keys_subset if key in original_dict}
+
+
+def reorder_and_convert_dict_list_to_table(
+    data: List[Dict[str, Any]], order: List[str]
+) -> Tuple[List[str], List[List[Any]]]:
+    """Convert a list of dictionaries to a pair of column names and corresponding values, with the option to order specific dictionaries.
+
+    :param data: A list of dictionaries.
+    :param order: A list of keys specifying the desired order for specific dictionaries. The remaining dictionaries will be ordered based on their original order.
+    :return: A pair of column names and corresponding values.
+    """
+    final_columns = []
+    keys_present = set()
+
+    # First, add all ordered keys to the final columns
+    for key in order:
+        if key not in keys_present:
+            final_columns.append(key)
+            keys_present.add(key)
+
+    # Then, add any keys present in the dictionaries but not in the order
+    for d in data:
+        for key in d:
+            if key not in keys_present:
+                final_columns.append(key)
+                keys_present.add(key)
+
+    # Then, construct the table of values
+    values = []
+    for d in data:
+        row = []
+        for key in final_columns:
+            row.append(d.get(key, None))
+        values.append(row)
+
+    return final_columns, values
+
+
+def flatten_dict(
+    dictionary: Dict[str, Any], parent_key: str = "", sep: str = "-"
+) -> Dict[str, Any]:
+    """Flatten a nested dictionary, joining keys using a specified separator.
+
+    :param dictionary: The dictionary to flatten.
+    :param parent_key: The base key to prepend to each key.
+    :param sep: The separator to use when joining keys.
+    :return: A flattened dictionary.
+    """
+    flattened_dict = {}
+    for key, value in dictionary.items():
+        new_key = f"{parent_key}{sep}{key}" if parent_key else key
+        if isinstance(value, dict):
+            flattened_dict.update(flatten_dict(value, new_key, sep=sep))
+        else:
+            flattened_dict[new_key] = value
+    return flattened_dict
+
+
+def collect_common_keys(list_of_dicts: List[Dict[str, Any]]) -> Dict[str, List[Any]]:
+    """Collect the common keys of a list of dictionaries. For each common key, put its values into a list in the order they appear in the original dictionaries.
+
+    :param list_of_dicts: The list of dictionaries to inspect.
+    :return: A dictionary with each common key and its corresponding list of values.
+    """
+    common_keys = set.intersection(*map(set, list_of_dicts))
+    common_dict = {key: [] for key in common_keys}
+    for d in list_of_dicts:
+        for key in common_keys:
+            common_dict[key].append(d[key])
+    return common_dict
+
+
+class CohereRequestResponseResolver:
+    """Class to resolve the request/response from the Cohere API and convert it to a dictionary that can be logged."""
+
+    def __call__(
+        self,
+        args: Sequence[Any],
+        kwargs: Dict[str, Any],
+        response: Response,
+        start_time: float,
+        time_elapsed: float,
+    ) -> Optional[Dict[str, Any]]:
+        """Process the response from the Cohere API and convert it to a dictionary that can be logged.
+
+        :param args: The arguments of the original function.
+        :param kwargs: The keyword arguments of the original function.
+        :param response: The response from the Cohere API.
+        :param start_time: The start time of the request.
+        :param time_elapsed: The time elapsed for the request.
+        :return: A dictionary containing the parsed response and timing information.
+        """
+        try:
+            # Each of the different endpoints map to one specific response type
+            # We want to 'type check' the response without directly importing the packages type
+            # It may make more sense to pass the invoked symbol from the AutologAPI instead
+            response_type = str(type(response)).split("'")[1].split(".")[-1]
+
+            # Initialize parsed_response to None to handle the case where the response type is unsupported
+            parsed_response = None
+            if response_type == "Generations":
+                parsed_response = self._resolve_generate_response(response)
+                # TODO: Remove hard-coded default model name
+                table_column_order = [
+                    "start_time",
+                    "query_id",
+                    "model",
+                    "prompt",
+                    "text",
+                    "token_likelihoods",
+                    "likelihood",
+                    "time_elapsed_(seconds)",
+                    "end_time",
+                ]
+                default_model = "command"
+            elif response_type == "Chat":
+                parsed_response = self._resolve_chat_response(response)
+                table_column_order = [
+                    "start_time",
+                    "query_id",
+                    "model",
+                    "conversation_id",
+                    "response_id",
+                    "query",
+                    "text",
+                    "prompt",
+                    "preamble",
+                    "chat_history",
+                    "chatlog",
+                    "time_elapsed_(seconds)",
+                    "end_time",
+                ]
+                default_model = "command"
+            elif response_type == "Classifications":
+                parsed_response = self._resolve_classify_response(response)
+                kwargs = self._resolve_classify_kwargs(kwargs)
+                table_column_order = [
+                    "start_time",
+                    "query_id",
+                    "model",
+                    "id",
+                    "input",
+                    "prediction",
+                    "confidence",
+                    "time_elapsed_(seconds)",
+                    "end_time",
+                ]
+                default_model = "embed-english-v2.0"
+            elif response_type == "SummarizeResponse":
+                parsed_response = self._resolve_summarize_response(response)
+                table_column_order = [
+                    "start_time",
+                    "query_id",
+                    "model",
+                    "response_id",
+                    "text",
+                    "additional_command",
+                    "summary",
+                    "time_elapsed_(seconds)",
+                    "end_time",
+                    "length",
+                    "format",
+                ]
+                default_model = "summarize-xlarge"
+            elif response_type == "Reranking":
+                parsed_response = self._resolve_rerank_response(response)
+                table_column_order = [
+                    "start_time",
+                    "query_id",
+                    "model",
+                    "id",
+                    "query",
+                    "top_n",
+                    # This is a nested dict key that got flattened
+                    "document-text",
+                    "relevance_score",
+                    "index",
+                    "time_elapsed_(seconds)",
+                    "end_time",
+                ]
+                default_model = "rerank-english-v2.0"
+            else:
+                logger.info(f"Unsupported Cohere response object: {response}")
+
+            return self._resolve(
+                args,
+                kwargs,
+                parsed_response,
+                start_time,
+                time_elapsed,
+                response_type,
+                table_column_order,
+                default_model,
+            )
+        except Exception as e:
+            logger.warning(f"Failed to resolve request/response: {e}")
+        return None
+
+    # These helper functions process the response from different endpoints of the Cohere API.
+    # Since the response objects for different endpoints have different structures,
+    # we need different logic to process them.
+
+    def _resolve_generate_response(self, response: Response) -> List[Dict[str, Any]]:
+        return_list = []
+        for _response in response:
+            # Built in Cohere.*.Generations function to color token_likelihoods and return a dict of response data
+            _response_dict = _response._visualize_helper()
+            try:
+                _response_dict["token_likelihoods"] = wandb.Html(
+                    _response_dict["token_likelihoods"]
+                )
+            except (KeyError, ValueError):
+                pass
+            return_list.append(_response_dict)
+
+        return return_list
+
+    def _resolve_chat_response(self, response: Response) -> List[Dict[str, Any]]:
+        return [
+            subset_dict(
+                response.__dict__,
+                [
+                    "response_id",
+                    "generation_id",
+                    "query",
+                    "text",
+                    "conversation_id",
+                    "prompt",
+                    "chatlog",
+                    "preamble",
+                ],
+            )
+        ]
+
+    def _resolve_classify_response(self, response: Response) -> List[Dict[str, Any]]:
+        # The labels key is a dict returning the scores for the classification probability for each label provided
+        # We flatten this nested dict for ease of consumption in the wandb UI
+        return [flatten_dict(_response.__dict__) for _response in response]
+
+    def _resolve_classify_kwargs(self, kwargs: Dict[str, Any]) -> Dict[str, Any]:
+        # Example texts look strange when rendered in Wandb UI as it is a list of text and label
+        # We extract each value into its own column
+        example_texts = []
+        example_labels = []
+        for example in kwargs["examples"]:
+            example_texts.append(example.text)
+            example_labels.append(example.label)
+        kwargs.pop("examples")
+        kwargs["example_texts"] = example_texts
+        kwargs["example_labels"] = example_labels
+        return kwargs
+
+    def _resolve_summarize_response(self, response: Response) -> List[Dict[str, Any]]:
+        return [{"response_id": response.id, "summary": response.summary}]
+
+    def _resolve_rerank_response(self, response: Response) -> List[Dict[str, Any]]:
+        # The documents key contains a dict containing the content of the document which is at least "text"
+        # We flatten this nested dict for ease of consumption in the wandb UI
+        flattened_response_dicts = [
+            flatten_dict(_response.__dict__) for _response in response
+        ]
+        # ReRank returns each document provided a top_n value so we aggregate into one view so users can paginate a row
+        # As opposed to each row being one of the top_n responses
+        return_dict = collect_common_keys(flattened_response_dicts)
+        return_dict["id"] = response.id
+        return [return_dict]
+
+    def _resolve(
+        self,
+        args: Sequence[Any],
+        kwargs: Dict[str, Any],
+        parsed_response: List[Dict[str, Any]],
+        start_time: float,
+        time_elapsed: float,
+        response_type: str,
+        table_column_order: List[str],
+        default_model: str,
+    ) -> Dict[str, Any]:
+        """Convert a list of dictionaries to a pair of column names and corresponding values, with the option to order specific dictionaries.
+
+        :param args: The arguments passed to the API client.
+        :param kwargs: The keyword arguments passed to the API client.
+        :param parsed_response: The parsed response from the API.
+        :param start_time: The start time of the API request.
+        :param time_elapsed: The time elapsed during the API request.
+        :param response_type: The type of the API response.
+        :param table_column_order: The desired order of columns in the resulting table.
+        :param default_model: The default model to use if not specified in the response.
+        :return: A dictionary containing the formatted response.
+        """
+        # Args[0] is the client object where we can grab specific metadata about the underlying API status
+        query_id = generate_id(length=16)
+        parsed_args = subset_dict(
+            args[0].__dict__,
+            ["api_version", "batch_size", "max_retries", "num_workers", "timeout"],
+        )
+
+        start_time_dt = datetime.fromtimestamp(start_time)
+        end_time_dt = datetime.fromtimestamp(start_time + time_elapsed)
+
+        timings = {
+            "start_time": start_time_dt,
+            "end_time": end_time_dt,
+            "time_elapsed_(seconds)": time_elapsed,
+        }
+
+        packed_data = []
+        for _parsed_response in parsed_response:
+            _packed_dict = {
+                "query_id": query_id,
+                **kwargs,
+                **_parsed_response,
+                **timings,
+                **parsed_args,
+            }
+            if "model" not in _packed_dict:
+                _packed_dict["model"] = default_model
+            packed_data.append(_packed_dict)
+
+        columns, data = reorder_and_convert_dict_list_to_table(
+            packed_data, table_column_order
+        )
+
+        request_response_table = wandb.Table(data=data, columns=columns)
+
+        return {f"{response_type}": request_response_table}

wandb/integration/diffusers/__init__.py

@@ -0,0 +1,3 @@
+from .autologger import autolog
+
+__all__ = ["autolog"]

wandb/integration/diffusers/autologger.py

@@ -0,0 +1,76 @@
+import logging
+
+from wandb.sdk.integration_utils.auto_logging import AutologAPI
+
+from .pipeline_resolver import DiffusersPipelineResolver
+
+logger = logging.getLogger(__name__)
+
+autolog = AutologAPI(
+    name="diffusers",
+    symbols=(
+        "DiffusionPipeline.__call__",
+        "AutoPipelineForText2Image.__call__",
+        "AutoPipelineForImage2Image.__call__",
+        "AutoPipelineForInpainting.__call__",
+        "StableDiffusionPipeline.__call__",
+        "KandinskyCombinedPipeline.__call__",
+        "KandinskyV22CombinedPipeline.__call__",
+        "LatentConsistencyModelPipeline.__call__",
+        "LDMTextToImagePipeline.__call__",
+        "StableDiffusionPanoramaPipeline.__call__",
+        "StableDiffusionParadigmsPipeline.__call__",
+        "PixArtAlphaPipeline.__call__",
+        "StableDiffusionSAGPipeline.__call__",
+        "SemanticStableDiffusionPipeline.__call__",
+        "WuerstchenCombinedPipeline.__call__",
+        "AltDiffusionPipeline.__call__",
+        "StableDiffusionAttendAndExcitePipeline.__call__",
+        "StableDiffusionXLPipeline.__call__",
+        "StableDiffusionXLImg2ImgPipeline.__call__",
+        "IFPipeline.__call__",
+        "BlipDiffusionPipeline.__call__",
+        "BlipDiffusionControlNetPipeline.__call__",
+        "StableDiffusionControlNetPipeline.__call__",
+        "StableDiffusionControlNetImg2ImgPipeline.__call__",
+        "StableDiffusionControlNetInpaintPipeline.__call__",
+        "CycleDiffusionPipeline.__call__",
+        "StableDiffusionInstructPix2PixPipeline.__call__",
+        "PaintByExamplePipeline.__call__",
+        "RePaintPipeline.__call__",
+        "KandinskyImg2ImgCombinedPipeline.__call__",
+        "KandinskyInpaintCombinedPipeline.__call__",
+        "KandinskyV22Img2ImgCombinedPipeline.__call__",
+        "KandinskyV22InpaintCombinedPipeline.__call__",
+        "Kandinsky3Pipeline.__call__",
+        "Kandinsky3Img2ImgPipeline.__call__",
+        "AnimateDiffPipeline.__call__",
+        "AudioLDMPipeline.__call__",
+        "AudioLDM2Pipeline.__call__",
+        "MusicLDMPipeline.__call__",
+        "StableDiffusionPix2PixZeroPipeline.__call__",
+        "PNDMPipeline.__call__",
+        "ShapEPipeline.__call__",
+        "StableDiffusionImg2ImgPipeline.__call__",
+        "StableDiffusionInpaintPipeline.__call__",
+        "StableDiffusionDepth2ImgPipeline.__call__",
+        "StableDiffusionImageVariationPipeline.__call__",
+        "StableDiffusionPipelineSafe.__call__",
+        "StableDiffusionUpscalePipeline.__call__",
+        "StableDiffusionAdapterPipeline.__call__",
+        "StableDiffusionGLIGENPipeline.__call__",
+        "StableDiffusionModelEditingPipeline.__call__",
+        "VersatileDiffusionTextToImagePipeline.__call__",
+        "VersatileDiffusionImageVariationPipeline.__call__",
+        "VersatileDiffusionDualGuidedPipeline.__call__",
+        "LDMPipeline.__call__",
+        "TextToVideoSDPipeline.__call__",
+        "TextToVideoZeroPipeline.__call__",
+        "StableVideoDiffusionPipeline.__call__",
+        "AmusedPipeline.__call__",
+        "StableDiffusionXLControlNetPipeline.__call__",
+        "StableDiffusionXLControlNetImg2ImgPipeline.__call__",
+    ),
+    resolver=DiffusersPipelineResolver(),
+    telemetry_feature="diffusers_autolog",
+)

wandb/integration/diffusers/pipeline_resolver.py

@@ -0,0 +1,50 @@
+from typing import Any, Dict, Sequence
+
+from wandb.sdk.integration_utils.auto_logging import Response
+
+from .resolvers import (
+    SUPPORTED_MULTIMODAL_PIPELINES,
+    DiffusersMultiModalPipelineResolver,
+)
+
+
+class DiffusersPipelineResolver:
+    """Resolver for `DiffusionPipeline` request and responses from [HuggingFace Diffusers](https://huggingface.co/docs/diffusers/index), providing necessary data transformations, formatting, and logging.
+
+    This is based off `wandb.sdk.integration_utils.auto_logging.RequestResponseResolver`.
+    """
+
+    def __init__(self) -> None:
+        self.wandb_table = None
+        self.pipeline_call_count = 1
+
+    def __call__(
+        self,
+        args: Sequence[Any],
+        kwargs: Dict[str, Any],
+        response: Response,
+        start_time: float,
+        time_elapsed: float,
+    ) -> Any:
+        """Main call method for the `DiffusersPipelineResolver` class.
+
+        Args:
+            args: (Sequence[Any]) List of arguments.
+            kwargs: (Dict[str, Any]) Dictionary of keyword arguments.
+            response: (wandb.sdk.integration_utils.auto_logging.Response) The response from
+                the request.
+            start_time: (float) Time when request started.
+            time_elapsed: (float) Time elapsed for the request.
+
+        Returns:
+            Packed data as a dictionary for logging to wandb, None if an exception occurred.
+        """
+        pipeline_name = args[0].__class__.__name__
+        resolver = None
+        if pipeline_name in SUPPORTED_MULTIMODAL_PIPELINES:
+            resolver = DiffusersMultiModalPipelineResolver(
+                pipeline_name, self.pipeline_call_count
+            )
+            self.pipeline_call_count += 1
+        loggable_dict = resolver(args, kwargs, response, start_time, time_elapsed)
+        return loggable_dict