wandb 0.18.0__py3-none-any.whl → 0.18.1__py3-none-any.whl

wandb/sdk/data_types/audio.py

@@ -0,0 +1,165 @@
+import hashlib
+import os
+from typing import Optional
+
+from wandb import util
+from wandb.sdk.lib import filesystem, runid
+
+from . import _dtypes
+from ._private import MEDIA_TMP
+from .base_types.media import BatchableMedia
+
+
+class Audio(BatchableMedia):
+    """Wandb class for audio clips.
+
+    Arguments:
+        data_or_path: (string or numpy array) A path to an audio file
+            or a numpy array of audio data.
+        sample_rate: (int) Sample rate, required when passing in raw
+            numpy array of audio data.
+        caption: (string) Caption to display with audio.
+    """
+
+    _log_type = "audio-file"
+
+    def __init__(self, data_or_path, sample_rate=None, caption=None):
+        """Accept a path to an audio file or a numpy array of audio data."""
+        super().__init__()
+        self._duration = None
+        self._sample_rate = sample_rate
+        self._caption = caption
+
+        if isinstance(data_or_path, str):
+            if self.path_is_reference(data_or_path):
+                self._path = data_or_path
+                self._sha256 = hashlib.sha256(data_or_path.encode("utf-8")).hexdigest()
+                self._is_tmp = False
+            else:
+                self._set_file(data_or_path, is_tmp=False)
+        else:
+            if sample_rate is None:
+                raise ValueError(
+                    'Argument "sample_rate" is required when instantiating wandb.Audio with raw data.'
+                )
+
+            soundfile = util.get_module(
+                "soundfile",
+                required='Raw audio requires the soundfile package. To get it, run "pip install soundfile"',
+            )
+
+            tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".wav")
+            soundfile.write(tmp_path, data_or_path, sample_rate)
+            self._duration = len(data_or_path) / float(sample_rate)
+
+            self._set_file(tmp_path, is_tmp=True)
+
+    @classmethod
+    def get_media_subdir(cls):
+        return os.path.join("media", "audio")
+
+    @classmethod
+    def from_json(cls, json_obj, source_artifact):
+        return cls(
+            source_artifact.get_entry(json_obj["path"]).download(),
+            caption=json_obj["caption"],
+        )
+
+    def bind_to_run(
+        self, run, key, step, id_=None, ignore_copy_err: Optional[bool] = None
+    ):
+        if self.path_is_reference(self._path):
+            raise ValueError(
+                "Audio media created by a reference to external storage cannot currently be added to a run"
+            )
+
+        return super().bind_to_run(run, key, step, id_, ignore_copy_err)
+
+    def to_json(self, run):
+        json_dict = super().to_json(run)
+        json_dict.update(
+            {
+                "_type": self._log_type,
+                "caption": self._caption,
+            }
+        )
+        return json_dict
+
+    @classmethod
+    def seq_to_json(cls, seq, run, key, step):
+        audio_list = list(seq)
+
+        util.get_module(
+            "soundfile",
+            required="wandb.Audio requires the soundfile package. To get it, run: pip install soundfile",
+        )
+        base_path = os.path.join(run.dir, "media", "audio")
+        filesystem.mkdir_exists_ok(base_path)
+        meta = {
+            "_type": "audio",
+            "count": len(audio_list),
+            "audio": [a.to_json(run) for a in audio_list],
+        }
+        sample_rates = cls.sample_rates(audio_list)
+        if sample_rates:
+            meta["sampleRates"] = sample_rates
+        durations = cls.durations(audio_list)
+        if durations:
+            meta["durations"] = durations
+        captions = cls.captions(audio_list)
+        if captions:
+            meta["captions"] = captions
+
+        return meta
+
+    @classmethod
+    def durations(cls, audio_list):
+        return [a._duration for a in audio_list]
+
+    @classmethod
+    def sample_rates(cls, audio_list):
+        return [a._sample_rate for a in audio_list]
+
+    @classmethod
+    def captions(cls, audio_list):
+        captions = [a._caption for a in audio_list]
+        if all(c is None for c in captions):
+            return False
+        else:
+            return ["" if c is None else c for c in captions]
+
+    def resolve_ref(self):
+        if self.path_is_reference(self._path):
+            # this object was already created using a ref:
+            return self._path
+        source_artifact = self._artifact_source.artifact
+
+        resolved_name = source_artifact._local_path_to_name(self._path)
+        if resolved_name is not None:
+            target_entry = source_artifact.manifest.get_entry_by_path(resolved_name)
+            if target_entry is not None:
+                return target_entry.ref
+
+        return None
+
+    def __eq__(self, other):
+        if self.path_is_reference(self._path) or self.path_is_reference(other._path):
+            # one or more of these objects is an unresolved reference -- we'll compare
+            # their reference paths instead of their SHAs:
+            return (
+                self.resolve_ref() == other.resolve_ref()
+                and self._caption == other._caption
+            )
+
+        return super().__eq__(other) and self._caption == other._caption
+
+    def __ne__(self, other):
+        return not self.__eq__(other)
+
+
+class _AudioFileType(_dtypes.Type):
+    name = "audio-file"
+    types = [Audio]
+
+
+_dtypes.TypeRegistry.add(_AudioFileType)

wandb/sdk/data_types/bokeh.py

@@ -0,0 +1,70 @@
+import codecs
+import json
+import os
+
+from wandb import util
+from wandb.sdk.lib import runid
+
+from . import _dtypes
+from ._private import MEDIA_TMP
+from .base_types.media import Media
+
+
+class Bokeh(Media):
+    """Wandb class for Bokeh plots.
+
+    Arguments:
+        val: Bokeh plot
+    """
+
+    _log_type = "bokeh-file"
+
+    def __init__(self, data_or_path):
+        super().__init__()
+        bokeh = util.get_module("bokeh", required=True)
+        if isinstance(data_or_path, str) and os.path.exists(data_or_path):
+            with open(data_or_path) as file:
+                b_json = json.load(file)
+            self.b_obj = bokeh.document.Document.from_json(b_json)
+            self._set_file(data_or_path, is_tmp=False, extension=".bokeh.json")
+        elif isinstance(data_or_path, bokeh.model.Model):
+            _data = bokeh.document.Document()
+            _data.add_root(data_or_path)
+            # serialize/deserialize pairing followed by sorting attributes ensures
+            # that the file's sha's are equivalent in subsequent calls
+            self.b_obj = bokeh.document.Document.from_json(_data.to_json())
+            b_json = self.b_obj.to_json()
+            if "references" in b_json["roots"]:
+                b_json["roots"]["references"].sort(key=lambda x: x["id"])
+
+            tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".bokeh.json")
+            with codecs.open(tmp_path, "w", encoding="utf-8") as fp:
+                util.json_dump_safer(b_json, fp)
+            self._set_file(tmp_path, is_tmp=True, extension=".bokeh.json")
+        elif not isinstance(data_or_path, bokeh.document.Document):
+            raise TypeError(
+                "Bokeh constructor accepts Bokeh document/model or path to Bokeh json file"
+            )
+
+    def get_media_subdir(self):
+        return os.path.join("media", "bokeh")
+
+    def to_json(self, run):
+        # TODO: (tss) this is getting redundant for all the media objects. We can probably
+        # pull this into Media#to_json and remove this type override for all the media types.
+        # There are only a few cases where the type is different between artifacts and runs.
+        json_dict = super().to_json(run)
+        json_dict["_type"] = self._log_type
+        return json_dict
+
+    @classmethod
+    def from_json(cls, json_obj, source_artifact):
+        return cls(source_artifact.get_entry(json_obj["path"]).download())
+
+
+class _BokehFileType(_dtypes.Type):
+    name = "bokeh-file"
+    types = [Bokeh]
+
+
+_dtypes.TypeRegistry.add(_BokehFileType)

wandb/sdk/data_types/graph.py

@@ -0,0 +1,405 @@
+import codecs
+import os
+import pprint
+
+from wandb import util
+from wandb.sdk.data_types._private import MEDIA_TMP
+from wandb.sdk.data_types.base_types.media import Media, _numpy_arrays_to_lists
+from wandb.sdk.data_types.base_types.wb_value import WBValue
+from wandb.sdk.lib import runid
+
+
+def _nest(thing):
+    # Use tensorflows nest function if available, otherwise just wrap object in an array"""
+
+    tfutil = util.get_module("tensorflow.python.util")
+    if tfutil:
+        return tfutil.nest.flatten(thing)
+    else:
+        return [thing]
+
+
+class Edge(WBValue):
+    """Edge used in `Graph`."""
+
+    def __init__(self, from_node, to_node):
+        self._attributes = {}
+        self.from_node = from_node
+        self.to_node = to_node
+
+    def __repr__(self):
+        temp_attr = dict(self._attributes)
+        del temp_attr["from_node"]
+        del temp_attr["to_node"]
+        temp_attr["from_id"] = self.from_node.id
+        temp_attr["to_id"] = self.to_node.id
+        return str(temp_attr)
+
+    def to_json(self, run=None):
+        return [self.from_node.id, self.to_node.id]
+
+    @property
+    def name(self):
+        """Optional, not necessarily unique."""
+        return self._attributes.get("name")
+
+    @name.setter
+    def name(self, val):
+        self._attributes["name"] = val
+        return val
+
+    @property
+    def from_node(self):
+        return self._attributes.get("from_node")
+
+    @from_node.setter
+    def from_node(self, val):
+        self._attributes["from_node"] = val
+        return val
+
+    @property
+    def to_node(self):
+        return self._attributes.get("to_node")
+
+    @to_node.setter
+    def to_node(self, val):
+        self._attributes["to_node"] = val
+        return val
+
+
+class Node(WBValue):
+    """Node used in `Graph`."""
+
+    def __init__(
+        self,
+        id=None,
+        name=None,
+        class_name=None,
+        size=None,
+        parameters=None,
+        output_shape=None,
+        is_output=None,
+        num_parameters=None,
+        node=None,
+    ):
+        self._attributes = {"name": None}
+        self.in_edges = {}  # indexed by source node id
+        self.out_edges = {}  # indexed by dest node id
+        # optional object (e.g. PyTorch Parameter or Module) that this Node represents
+        self.obj = None
+
+        if node is not None:
+            self._attributes.update(node._attributes)
+            del self._attributes["id"]
+            self.obj = node.obj
+
+        if id is not None:
+            self.id = id
+        if name is not None:
+            self.name = name
+        if class_name is not None:
+            self.class_name = class_name
+        if size is not None:
+            self.size = size
+        if parameters is not None:
+            self.parameters = parameters
+        if output_shape is not None:
+            self.output_shape = output_shape
+        if is_output is not None:
+            self.is_output = is_output
+        if num_parameters is not None:
+            self.num_parameters = num_parameters
+
+    def to_json(self, run=None):
+        return self._attributes
+
+    def __repr__(self):
+        return repr(self._attributes)
+
+    @property
+    def id(self):
+        """Must be unique in the graph."""
+        return self._attributes.get("id")
+
+    @id.setter
+    def id(self, val):
+        self._attributes["id"] = val
+        return val
+
+    @property
+    def name(self):
+        """Usually the type of layer or sublayer."""
+        return self._attributes.get("name")
+
+    @name.setter
+    def name(self, val):
+        self._attributes["name"] = val
+        return val
+
+    @property
+    def class_name(self):
+        """Usually the type of layer or sublayer."""
+        return self._attributes.get("class_name")
+
+    @class_name.setter
+    def class_name(self, val):
+        self._attributes["class_name"] = val
+        return val
+
+    @property
+    def functions(self):
+        return self._attributes.get("functions", [])
+
+    @functions.setter
+    def functions(self, val):
+        self._attributes["functions"] = val
+        return val
+
+    @property
+    def parameters(self):
+        return self._attributes.get("parameters", [])
+
+    @parameters.setter
+    def parameters(self, val):
+        self._attributes["parameters"] = val
+        return val
+
+    @property
+    def size(self):
+        return self._attributes.get("size")
+
+    @size.setter
+    def size(self, val):
+        """Tensor size."""
+        self._attributes["size"] = tuple(val)
+        return val
+
+    @property
+    def output_shape(self):
+        return self._attributes.get("output_shape")
+
+    @output_shape.setter
+    def output_shape(self, val):
+        """Tensor output_shape."""
+        self._attributes["output_shape"] = val
+        return val
+
+    @property
+    def is_output(self):
+        return self._attributes.get("is_output")
+
+    @is_output.setter
+    def is_output(self, val):
+        """Tensor is_output."""
+        self._attributes["is_output"] = val
+        return val
+
+    @property
+    def num_parameters(self):
+        return self._attributes.get("num_parameters")
+
+    @num_parameters.setter
+    def num_parameters(self, val):
+        """Tensor num_parameters."""
+        self._attributes["num_parameters"] = val
+        return val
+
+    @property
+    def child_parameters(self):
+        return self._attributes.get("child_parameters")
+
+    @child_parameters.setter
+    def child_parameters(self, val):
+        """Tensor child_parameters."""
+        self._attributes["child_parameters"] = val
+        return val
+
+    @property
+    def is_constant(self):
+        return self._attributes.get("is_constant")
+
+    @is_constant.setter
+    def is_constant(self, val):
+        """Tensor is_constant."""
+        self._attributes["is_constant"] = val
+        return val
+
+    @classmethod
+    def from_keras(cls, layer):
+        node = cls()
+
+        try:
+            output_shape = layer.output_shape
+        except AttributeError:
+            output_shape = ["multiple"]
+
+        node.id = layer.name
+        node.name = layer.name
+        node.class_name = layer.__class__.__name__
+        node.output_shape = output_shape
+        node.num_parameters = layer.count_params()
+
+        return node
+
+
+class Graph(Media):
+    """Wandb class for graphs.
+
+    This class is typically used for saving and displaying neural net models.  It
+    represents the graph as an array of nodes and edges.  The nodes can have
+    labels that can be visualized by wandb.
+
+    Examples:
+        Import a keras model:
+        ```
+        Graph.from_keras(keras_model)
+        ```
+
+    Attributes:
+        format (string): Format to help wandb display the graph nicely.
+        nodes ([wandb.Node]): List of wandb.Nodes
+        nodes_by_id (dict): dict of ids -> nodes
+        edges ([(wandb.Node, wandb.Node)]): List of pairs of nodes interpreted as edges
+        loaded (boolean): Flag to tell whether the graph is completely loaded
+        root (wandb.Node): root node of the graph
+    """
+
+    _log_type = "graph-file"
+
+    def __init__(self, format="keras"):
+        super().__init__()
+        # LB: TODO: I think we should factor criterion and criterion_passed out
+        self.format = format
+        self.nodes = []
+        self.nodes_by_id = {}
+        self.edges = []
+        self.loaded = False
+        self.criterion = None
+        self.criterion_passed = False
+        self.root = None  # optional root Node if applicable
+
+    def _to_graph_json(self, run=None):
+        # Needs to be its own function for tests
+        return {
+            "format": self.format,
+            "nodes": [node.to_json() for node in self.nodes],
+            "edges": [edge.to_json() for edge in self.edges],
+        }
+
+    def bind_to_run(self, *args, **kwargs):
+        data = self._to_graph_json()
+        tmp_path = os.path.join(MEDIA_TMP.name, runid.generate_id() + ".graph.json")
+        data = _numpy_arrays_to_lists(data)
+        with codecs.open(tmp_path, "w", encoding="utf-8") as fp:
+            util.json_dump_safer(data, fp)
+        self._set_file(tmp_path, is_tmp=True, extension=".graph.json")
+        if self.is_bound():
+            return
+        super().bind_to_run(*args, **kwargs)
+
+    @classmethod
+    def get_media_subdir(cls):
+        return os.path.join("media", "graph")
+
+    def to_json(self, run):
+        json_dict = super().to_json(run)
+        json_dict["_type"] = self._log_type
+        return json_dict
+
+    def __getitem__(self, nid):
+        return self.nodes_by_id[nid]
+
+    def pprint(self):
+        for edge in self.edges:
+            pprint.pprint(edge.attributes)
+        for node in self.nodes:
+            pprint.pprint(node.attributes)
+
+    def add_node(self, node=None, **node_kwargs):
+        if node is None:
+            node = Node(**node_kwargs)
+        elif node_kwargs:
+            raise ValueError(
+                f"Only pass one of either node ({node}) or other keyword arguments ({node_kwargs})"
+            )
+        self.nodes.append(node)
+        self.nodes_by_id[node.id] = node
+
+        return node
+
+    def add_edge(self, from_node, to_node):
+        edge = Edge(from_node, to_node)
+        self.edges.append(edge)
+
+        return edge
+
+    @classmethod
+    def from_keras(cls, model):
+        # TODO: his method requires a refactor to work with the keras 3.
+        graph = cls()
+        # Shamelessly copied (then modified) from keras/keras/utils/layer_utils.py
+        sequential_like = cls._is_sequential(model)
+
+        relevant_nodes = None
+        if not sequential_like:
+            relevant_nodes = []
+            for v in model._nodes_by_depth.values():
+                relevant_nodes += v
+
+        layers = model.layers
+        for i in range(len(layers)):
+            node = Node.from_keras(layers[i])
+            if hasattr(layers[i], "_inbound_nodes"):
+                for in_node in layers[i]._inbound_nodes:
+                    if relevant_nodes and in_node not in relevant_nodes:
+                        # node is not part of the current network
+                        continue
+                    for in_layer in _nest(in_node.inbound_layers):
+                        inbound_keras_node = Node.from_keras(in_layer)
+
+                        if inbound_keras_node.id not in graph.nodes_by_id:
+                            graph.add_node(inbound_keras_node)
+                        inbound_node = graph.nodes_by_id[inbound_keras_node.id]
+
+                        graph.add_edge(inbound_node, node)
+            graph.add_node(node)
+        return graph
+
+    @classmethod
+    def _is_sequential(cls, model):
+        sequential_like = True
+
+        if (
+            model.__class__.__name__ != "Sequential"
+            and hasattr(model, "_is_graph_network")
+            and model._is_graph_network
+        ):
+            nodes_by_depth = model._nodes_by_depth.values()
+            nodes = []
+            for v in nodes_by_depth:
+                # TensorFlow2 doesn't insure inbound is always a list
+                inbound = v[0].inbound_layers
+                if not hasattr(inbound, "__len__"):
+                    inbound = [inbound]
+                if (len(v) > 1) or (len(v) == 1 and len(inbound) > 1):
+                    # if the model has multiple nodes
+                    # or if the nodes have multiple inbound_layers
+                    # the model is no longer sequential
+                    sequential_like = False
+                    break
+                nodes += v
+            if sequential_like:
+                # search for shared layers
+                for layer in model.layers:
+                    flag = False
+                    if hasattr(layer, "_inbound_nodes"):
+                        for node in layer._inbound_nodes:
+                            if node in nodes:
+                                if flag:
+                                    sequential_like = False
+                                    break
+                                else:
+                                    flag = True
+                    if not sequential_like:
+                        break
+        return sequential_like