wandb 0.18.2__py3-none-musllinux_1_2_x86_64.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).
+
+    Arguments:
+        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/fastai/__init__.py

@@ -0,0 +1,249 @@
+"""Hooks that add fast.ai v1 Learners to Weights & Biases through a callback.
+
+Requested logged data can be configured through the callback constructor.
+
+Examples:
+    WandbCallback can be used when initializing the Learner::
+
+    ```
+        from wandb.fastai import WandbCallback
+        [...]
+        learn = Learner(data, ..., callback_fns=WandbCallback)
+        learn.fit(epochs)
+    ```
+
+    Custom parameters can be given using functools.partial::
+
+    ```
+        from wandb.fastai import WandbCallback
+        from functools import partial
+        [...]
+        learn = Learner(data, ..., callback_fns=partial(WandbCallback, ...))
+        learn.fit(epochs)
+    ```
+
+    Finally, it is possible to use WandbCallback only when starting
+    training. In this case it must be instantiated::
+
+    ```
+        learn.fit(..., callbacks=WandbCallback(learn))
+    ```
+
+    or, with custom parameters::
+
+    ```
+        learn.fit(..., callbacks=WandbCallback(learn, ...))
+    ```
+"""
+
+import random
+import sys
+from pathlib import Path
+from typing import Any, Optional
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+import fastai
+from fastai.callbacks import TrackerCallback
+
+import wandb
+
+try:
+    import matplotlib
+
+    if wandb.wandb_lib.ipython._get_python_type() != "jupyter":  # type: ignore[attr-defined]
+        matplotlib.use("Agg")  # non-interactive backend (avoid tkinter issues)
+    import matplotlib.pyplot as plt
+except ImportError:
+    print("Warning: matplotlib required if logging sample image predictions")
+
+
+class WandbCallback(TrackerCallback):
+    """Callback for saving model topology, losses & metrics.
+
+    Optionally logs weights, gradients, sample predictions and best trained model.
+
+    Arguments:
+        learn (fastai.basic_train.Learner): the fast.ai learner to hook.
+        log (str): "gradients", "parameters", "all", or None. Losses & metrics are always logged.
+        save_model (bool): save model at the end of each epoch. It will also load best model at the end of training.
+        monitor (str): metric to monitor for saving best model. None uses default TrackerCallback monitor value.
+        mode (str): "auto", "min" or "max" to compare "monitor" values and define best model.
+        input_type (str): "images" or None. Used to display sample predictions.
+        validation_data (list): data used for sample predictions if input_type is set.
+        predictions (int): number of predictions to make if input_type is set and validation_data is None.
+        seed (int): initialize random generator for sample predictions if input_type is set and validation_data is None.
+    """
+
+    # Record if watch has been called previously (even in another instance)
+    _watch_called = False
+
+    def __init__(
+        self,
+        learn: "fastai.basic_train.Learner",
+        log: Optional[Literal["gradients", "parameters", "all"]] = "gradients",
+        save_model: bool = True,
+        monitor: Optional[str] = None,
+        mode: Literal["auto", "min", "max"] = "auto",
+        input_type: Optional[Literal["images"]] = None,
+        validation_data: Optional[list] = None,
+        predictions: int = 36,
+        seed: int = 12345,
+    ) -> None:
+        # Check if wandb.init has been called
+        if wandb.run is None:
+            raise ValueError("You must call wandb.init() before WandbCallback()")
+
+        # Adapted from fast.ai "SaveModelCallback"
+        if monitor is None:
+            # use default TrackerCallback monitor value
+            super().__init__(learn, mode=mode)
+        else:
+            super().__init__(learn, monitor=monitor, mode=mode)
+        self.save_model = save_model
+        self.model_path = Path(wandb.run.dir) / "bestmodel.pth"
+
+        self.log = log
+        self.input_type = input_type
+        self.best = None
+
+        # Select items for sample predictions to see evolution along training
+        self.validation_data = validation_data
+        if input_type and not self.validation_data:
+            wandb_random = random.Random(seed)  # For repeatability
+            predictions = min(predictions, len(learn.data.valid_ds))
+            indices = wandb_random.sample(range(len(learn.data.valid_ds)), predictions)
+            self.validation_data = [learn.data.valid_ds[i] for i in indices]
+
+    def on_train_begin(self, **kwargs: Any) -> None:
+        """Call watch method to log model topology, gradients & weights."""
+        # Set self.best, method inherited from "TrackerCallback" by "SaveModelCallback"
+        super().on_train_begin()
+
+        # Ensure we don't call "watch" multiple times
+        if not WandbCallback._watch_called:
+            WandbCallback._watch_called = True
+
+            # Logs model topology and optionally gradients and weights
+            wandb.watch(self.learn.model, log=self.log)
+
+    def on_epoch_end(
+        self, epoch: int, smooth_loss: float, last_metrics: list, **kwargs: Any
+    ) -> None:
+        """Log training loss, validation loss and custom metrics & log prediction samples & save model."""
+        if self.save_model:
+            # Adapted from fast.ai "SaveModelCallback"
+            current = self.get_monitor_value()
+            if current is not None and self.operator(current, self.best):
+                print(
+                    "Better model found at epoch {} with {} value: {}.".format(
+                        epoch, self.monitor, current
+                    )
+                )
+                self.best = current
+
+                # Save within wandb folder
+                with self.model_path.open("wb") as model_file:
+                    self.learn.save(model_file)
+
+        # Log sample predictions if learn.predict is available
+        if self.validation_data:
+            try:
+                self._wandb_log_predictions()
+            except FastaiError as e:
+                wandb.termwarn(e.message)
+                self.validation_data = None  # prevent from trying again on next loop
+            except Exception as e:
+                wandb.termwarn(f"Unable to log prediction samples.\n{e}")
+                self.validation_data = None  # prevent from trying again on next loop
+
+        # Log losses & metrics
+        # Adapted from fast.ai "CSVLogger"
+        logs = {
+            name: stat
+            for name, stat in list(
+                zip(self.learn.recorder.names, [epoch, smooth_loss] + last_metrics)
+            )
+        }
+        wandb.log(logs)
+
+    def on_train_end(self, **kwargs: Any) -> None:
+        """Load the best model."""
+        if self.save_model:
+            # Adapted from fast.ai "SaveModelCallback"
+            if self.model_path.is_file():
+                with self.model_path.open("rb") as model_file:
+                    self.learn.load(model_file, purge=False)
+                    print(f"Loaded best saved model from {self.model_path}")
+
+    def _wandb_log_predictions(self) -> None:
+        """Log prediction samples."""
+        pred_log = []
+
+        if self.validation_data is None:
+            return
+
+        for x, y in self.validation_data:
+            try:
+                pred = self.learn.predict(x)
+            except Exception:
+                raise FastaiError(
+                    'Unable to run "predict" method from Learner to log prediction samples.'
+                )
+
+            # scalar -> likely to be a category
+            # tensor of dim 1 -> likely to be multicategory
+            if not pred[1].shape or pred[1].dim() == 1:
+                pred_log.append(
+                    wandb.Image(
+                        x.data,
+                        caption=f"Ground Truth: {y}\nPrediction: {pred[0]}",
+                    )
+                )
+
+            # most vision datasets have a "show" function we can use
+            elif hasattr(x, "show"):
+                # log input data
+                pred_log.append(wandb.Image(x.data, caption="Input data", grouping=3))
+
+                # log label and prediction
+                for im, capt in ((pred[0], "Prediction"), (y, "Ground Truth")):
+                    # Resize plot to image resolution
+                    # from https://stackoverflow.com/a/13714915
+                    my_dpi = 100
+                    fig = plt.figure(frameon=False, dpi=my_dpi)
+                    h, w = x.size
+                    fig.set_size_inches(w / my_dpi, h / my_dpi)
+                    ax = plt.Axes(fig, [0.0, 0.0, 1.0, 1.0])
+                    ax.set_axis_off()
+                    fig.add_axes(ax)
+
+                    # Superpose label or prediction to input image
+                    x.show(ax=ax, y=im)
+                    pred_log.append(wandb.Image(fig, caption=capt))
+                    plt.close(fig)
+
+            # likely to be an image
+            elif hasattr(y, "shape") and (
+                (len(y.shape) == 2) or (len(y.shape) == 3 and y.shape[0] in [1, 3, 4])
+            ):
+                pred_log.extend(
+                    [
+                        wandb.Image(x.data, caption="Input data", grouping=3),
+                        wandb.Image(pred[0].data, caption="Prediction"),
+                        wandb.Image(y.data, caption="Ground Truth"),
+                    ]
+                )
+
+            # we just log input data
+            else:
+                pred_log.append(wandb.Image(x.data, caption="Input data"))
+
+            wandb.log({"Prediction Samples": pred_log}, commit=False)
+
+
+class FastaiError(wandb.Error):
+    pass

wandb/integration/gym/__init__.py

@@ -0,0 +1,105 @@
+import re
+import sys
+from typing import Optional
+
+import wandb
+import wandb.util
+
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
+
+_gym_version_lt_0_26: Optional[bool] = None
+_gymnasium_version_lt_1_0_0: Optional[bool] = None
+
+_required_error_msg = (
+    "Couldn't import the gymnasium python package, "
+    "install with `pip install gymnasium`"
+)
+GymLib = Literal["gym", "gymnasium"]
+
+
+def monitor():
+    """Monitor a gym environment.
+
+    Supports both gym and gymnasium.
+    """
+    gym_lib: Optional[GymLib] = None
+
+    # gym is not maintained anymore, gymnasium is the drop-in replacement - prefer it
+    if wandb.util.get_module("gymnasium") is not None:
+        gym_lib = "gymnasium"
+    elif wandb.util.get_module("gym") is not None:
+        gym_lib = "gym"
+
+    if gym_lib is None:
+        raise wandb.Error(_required_error_msg)
+
+    global _gym_version_lt_0_26
+    global _gymnasium_version_lt_1_0_0
+
+    if _gym_version_lt_0_26 is None or _gymnasium_version_lt_1_0_0 is None:
+        if gym_lib == "gym":
+            import gym
+        else:
+            import gymnasium as gym  # type: ignore
+
+        from wandb.util import parse_version
+
+        gym_lib_version = parse_version(gym.__version__)
+        _gym_version_lt_0_26 = gym_lib_version < parse_version("0.26.0")
+        _gymnasium_version_lt_1_0_0 = gym_lib_version < parse_version("1.0.0a1")
+
+    path = "path"  # Default path
+    if gym_lib == "gymnasium" and not _gymnasium_version_lt_1_0_0:
+        vcr_recorder_attribute = "RecordVideo"
+        wrappers = wandb.util.get_module(
+            f"{gym_lib}.wrappers",
+            required=_required_error_msg,
+        )
+        recorder = getattr(wrappers, vcr_recorder_attribute)
+    else:
+        vcr = wandb.util.get_module(
+            f"{gym_lib}.wrappers.monitoring.video_recorder",
+            required=_required_error_msg,
+        )
+        # Breaking change in gym 0.26.0
+        if _gym_version_lt_0_26:
+            vcr_recorder_attribute = "ImageEncoder"
+            recorder = getattr(vcr, vcr_recorder_attribute)
+            path = "output_path"  # Override path for older gym versions
+        else:
+            vcr_recorder_attribute = "VideoRecorder"
+            recorder = getattr(vcr, vcr_recorder_attribute)
+
+    recorder.orig_close = recorder.close
+
+    def close(self):
+        recorder.orig_close(self)
+        if not self.enabled:
+            return
+        if wandb.run:
+            m = re.match(r".+(video\.\d+).+", getattr(self, path))
+            key = m.group(1) if m else "videos"
+            wandb.log({key: wandb.Video(getattr(self, path))})
+
+    def del_(self):
+        self.orig_close()
+
+    if not _gym_version_lt_0_26:
+        recorder.__del__ = del_
+    recorder.close = close
+
+    if gym_lib == "gymnasium" and not _gymnasium_version_lt_1_0_0:
+        wrapper_name = vcr_recorder_attribute
+    else:
+        wrapper_name = f"monitoring.video_recorder.{vcr_recorder_attribute}"
+
+    wandb.patched["gym"].append(
+        [
+            f"{gym_lib}.wrappers.{wrapper_name}",
+            "close",
+        ]
+    )

wandb/integration/huggingface/__init__.py

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

wandb/integration/huggingface/huggingface.py

@@ -0,0 +1,18 @@
+import logging
+
+from wandb.sdk.integration_utils.auto_logging import AutologAPI
+
+from .resolver import HuggingFacePipelineRequestResponseResolver
+
+logger = logging.getLogger(__name__)
+
+resolver = HuggingFacePipelineRequestResponseResolver()
+
+autolog = AutologAPI(
+    name="transformers",
+    symbols=("Pipeline.__call__",),
+    resolver=resolver,
+    telemetry_feature="hf_pipeline_autolog",
+)
+
+autolog.get_latest_id = resolver.get_latest_id

wandb/integration/huggingface/resolver.py

@@ -0,0 +1,213 @@
+import logging
+import os
+from datetime import datetime
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+import pytz
+
+import wandb
+from wandb.sdk.integration_utils.auto_logging import Response
+from wandb.sdk.lib.runid import generate_id
+
+logger = logging.getLogger(__name__)
+
+SUPPORTED_PIPELINE_TASKS = [
+    "text-classification",
+    "sentiment-analysis",
+    "question-answering",
+    "summarization",
+    "translation",
+    "text2text-generation",
+    "text-generation",
+    # "conversational",
+]
+
+PIPELINES_WITH_TOP_K = [
+    "text-classification",
+    "sentiment-analysis",
+    "question-answering",
+]
+
+
+class HuggingFacePipelineRequestResponseResolver:
+    """Resolver for HuggingFace's pipeline request and responses, providing necessary data transformations and formatting.
+
+    This is based off (from wandb.sdk.integration_utils.auto_logging import RequestResponseResolver)
+    """
+
+    autolog_id = None
+
+    def __call__(
+        self,
+        args: Sequence[Any],
+        kwargs: Dict[str, Any],
+        response: Response,
+        start_time: float,
+        time_elapsed: float,
+    ) -> Optional[Dict[str, Any]]:
+        """Main call method for this class.
+
+        :param args: list of arguments
+        :param kwargs: dictionary of keyword arguments
+        :param response: the response from the request
+        :param start_time: time when request started
+        :param time_elapsed: time elapsed for the request
+        :returns: packed data as a dictionary for logging to wandb, None if an exception occurred
+        """
+        try:
+            pipe, input_data = args[:2]
+            task = pipe.task
+
+            # Translation tasks are in the form of `translation_x_to_y`
+            if task in SUPPORTED_PIPELINE_TASKS or task.startswith("translation"):
+                model = self._get_model(pipe)
+                if model is None:
+                    return None
+                model_alias = model.name_or_path
+                timestamp = datetime.now(pytz.utc)
+
+                input_data, response = self._transform_task_specific_data(
+                    task, input_data, response
+                )
+                formatted_data = self._format_data(task, input_data, response, kwargs)
+                packed_data = self._create_table(
+                    formatted_data, model_alias, timestamp, time_elapsed
+                )
+                table_name = os.environ.get("WANDB_AUTOLOG_TABLE_NAME", f"{task}")
+                # TODO: Let users decide the name in a way that does not use an environment variable
+
+                return {
+                    table_name: wandb.Table(
+                        columns=packed_data[0], data=packed_data[1:]
+                    )
+                }
+
+            logger.warning(
+                f"The task: `{task}` is not yet supported.\nPlease contact `wandb` to notify us if you would like support for this task"
+            )
+        except Exception as e:
+            logger.warning(e)
+        return None
+
+    # TODO: This should have a dependency on PreTrainedModel. i.e. isinstance(PreTrainedModel)
+    # from transformers.modeling_utils import PreTrainedModel
+    # We do not want this dependency explicitly in our codebase so we make a very general
+    # assumption about the structure of the pipeline which may have unintended consequences
+    def _get_model(self, pipe) -> Optional[Any]:
+        """Extracts model from the pipeline.
+
+        :param pipe: the HuggingFace pipeline
+        :returns: Model if available, None otherwise
+        """
+        model = pipe.model
+        try:
+            return model.model
+        except AttributeError:
+            logger.info(
+                "Model does not have a `.model` attribute. Assuming `pipe.model` is the correct model."
+            )
+            return model
+
+    @staticmethod
+    def _transform_task_specific_data(
+        task: str, input_data: Union[List[Any], Any], response: Union[List[Any], Any]
+    ) -> Tuple[Union[List[Any], Any], Union[List[Any], Any]]:
+        """Transform input and response data based on specific tasks.
+
+        :param task: the task name
+        :param input_data: the input data
+        :param response: the response data
+        :returns: tuple of transformed input_data and response
+        """
+        if task == "question-answering":
+            input_data = input_data if isinstance(input_data, list) else [input_data]
+            input_data = [data.__dict__ for data in input_data]
+        elif task == "conversational":
+            # We only grab the latest input/output pair from the conversation
+            # Logging the whole conversation renders strangely.
+            input_data = input_data if isinstance(input_data, list) else [input_data]
+            input_data = [data.__dict__["past_user_inputs"][-1] for data in input_data]
+
+            response = response if isinstance(response, list) else [response]
+            response = [data.__dict__["generated_responses"][-1] for data in response]
+        return input_data, response
+
+    def _format_data(
+        self,
+        task: str,
+        input_data: Union[List[Any], Any],
+        response: Union[List[Any], Any],
+        kwargs: Dict[str, Any],
+    ) -> List[Dict[str, Any]]:
+        """Formats input data, response, and kwargs into a list of dictionaries.
+
+        :param task: the task name
+        :param input_data: the input data
+        :param response: the response data
+        :param kwargs: dictionary of keyword arguments
+        :returns: list of dictionaries containing formatted data
+        """
+        input_data = input_data if isinstance(input_data, list) else [input_data]
+        response = response if isinstance(response, list) else [response]
+
+        formatted_data = []
+        for i_text, r_text in zip(input_data, response):
+            # Unpack single element responses for better rendering in wandb UI when it is a task without top_k
+            # top_k = 1 would unpack the response into a single element while top_k > 1 would be a list
+            # this would cause the UI to not properly concatenate the tables of the same task by omitting the elements past the first
+            if (
+                (isinstance(r_text, list))
+                and (len(r_text) == 1)
+                and task not in PIPELINES_WITH_TOP_K
+            ):
+                r_text = r_text[0]
+            formatted_data.append(
+                {"input": i_text, "response": r_text, "kwargs": kwargs}
+            )
+        return formatted_data
+
+    def _create_table(
+        self,
+        formatted_data: List[Dict[str, Any]],
+        model_alias: str,
+        timestamp: float,
+        time_elapsed: float,
+    ) -> List[List[Any]]:
+        """Creates a table from formatted data, model alias, timestamp, and elapsed time.
+
+        :param formatted_data: list of dictionaries containing formatted data
+        :param model_alias: alias of the model
+        :param timestamp: timestamp of the data
+        :param time_elapsed: time elapsed from the beginning
+        :returns: list of lists, representing a table of data. [0]th element = columns. [1]st element = data
+        """
+        header = [
+            "ID",
+            "Model Alias",
+            "Timestamp",
+            "Elapsed Time",
+            "Input",
+            "Response",
+            "Kwargs",
+        ]
+        table = [header]
+        autolog_id = generate_id(length=16)
+
+        for data in formatted_data:
+            row = [
+                autolog_id,
+                model_alias,
+                timestamp,
+                time_elapsed,
+                data["input"],
+                data["response"],
+                data["kwargs"],
+            ]
+            table.append(row)
+
+        self.autolog_id = autolog_id
+
+        return table
+
+    def get_latest_id(self):
+        return self.autolog_id