wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/integration/openai/resolver.py

@@ -0,0 +1,240 @@
+import datetime
+import io
+import logging
+from dataclasses import asdict, dataclass
+from typing import Any, Dict, List, Optional, Sequence
+
+import wandb
+from wandb.sdk.data_types import trace_tree
+from wandb.sdk.integration_utils.auto_logging import Response
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class UsageMetrics:
+    elapsed_time: float = None
+    prompt_tokens: int = None
+    completion_tokens: int = None
+    total_tokens: int = None
+
+
+@dataclass
+class Metrics:
+    usage: UsageMetrics = None
+    stats: wandb.Table = None
+    trace: trace_tree.WBTraceTree = None
+
+
+usage_metric_keys = {f"usage/{k}" for k in asdict(UsageMetrics())}
+
+
+class OpenAIRequestResponseResolver:
+    def __init__(self):
+        self.define_metrics_called = False
+
+    def __call__(
+        self,
+        args: Sequence[Any],
+        kwargs: Dict[str, Any],
+        response: Response,
+        start_time: float,  # pass to comply with the protocol, but use response["created"] instead
+        time_elapsed: float,
+    ) -> Optional[Dict[str, Any]]:
+        request = kwargs
+
+        if not self.define_metrics_called:
+            # define metrics on first call
+            for key in usage_metric_keys:
+                wandb.define_metric(key, step_metric="_timestamp")
+            self.define_metrics_called = True
+
+        try:
+            if response.get("object") == "edit":
+                return self._resolve_edit(request, response, time_elapsed)
+            elif response.get("object") == "text_completion":
+                return self._resolve_completion(request, response, time_elapsed)
+            elif response.get("object") == "chat.completion":
+                return self._resolve_chat_completion(request, response, time_elapsed)
+            else:
+                # todo: properly treat failed requests
+                logger.info(
+                    f"Unsupported OpenAI response object: {response.get('object')}"
+                )
+        except Exception as e:
+            logger.warning(f"Failed to resolve request/response: {e}")
+        return None
+
+    @staticmethod
+    def results_to_trace_tree(
+        request: Dict[str, Any],
+        response: Response,
+        results: List[trace_tree.Result],
+        time_elapsed: float,
+    ) -> trace_tree.WBTraceTree:
+        """Converts the request, response, and results into a trace tree.
+
+        params:
+            request: The request dictionary
+            response: The response object
+            results: A list of results object
+            time_elapsed: The time elapsed in seconds
+        returns:
+            A wandb trace tree object.
+        """
+        start_time_ms = int(round(response["created"] * 1000))
+        end_time_ms = start_time_ms + int(round(time_elapsed * 1000))
+        span = trace_tree.Span(
+            name=f"{response.get('model', 'openai')}_{response['object']}_{response.get('created')}",
+            attributes=dict(response),  # type: ignore
+            start_time_ms=start_time_ms,
+            end_time_ms=end_time_ms,
+            span_kind=trace_tree.SpanKind.LLM,
+            results=results,
+        )
+        model_obj = {"request": request, "response": response, "_kind": "openai"}
+        return trace_tree.WBTraceTree(root_span=span, model_dict=model_obj)
+
+    def _resolve_edit(
+        self,
+        request: Dict[str, Any],
+        response: Response,
+        time_elapsed: float,
+    ) -> Dict[str, Any]:
+        """Resolves the request and response objects for `openai.Edit`."""
+        request_str = (
+            f"\n\n**Instruction**: {request['instruction']}\n\n"
+            f"**Input**: {request['input']}\n"
+        )
+        choices = [
+            f"\n\n**Edited**: {choice['text']}\n" for choice in response["choices"]
+        ]
+
+        return self._resolve_metrics(
+            request=request,
+            response=response,
+            request_str=request_str,
+            choices=choices,
+            time_elapsed=time_elapsed,
+        )
+
+    def _resolve_completion(
+        self,
+        request: Dict[str, Any],
+        response: Response,
+        time_elapsed: float,
+    ) -> Dict[str, Any]:
+        """Resolves the request and response objects for `openai.Completion`."""
+        request_str = f"\n\n**Prompt**: {request['prompt']}\n"
+        choices = [
+            f"\n\n**Completion**: {choice['text']}\n" for choice in response["choices"]
+        ]
+
+        return self._resolve_metrics(
+            request=request,
+            response=response,
+            request_str=request_str,
+            choices=choices,
+            time_elapsed=time_elapsed,
+        )
+
+    def _resolve_chat_completion(
+        self,
+        request: Dict[str, Any],
+        response: Response,
+        time_elapsed: float,
+    ) -> Dict[str, Any]:
+        """Resolves the request and response objects for `openai.Completion`."""
+        prompt = io.StringIO()
+        for message in request["messages"]:
+            prompt.write(f"\n\n**{message['role']}**: {message['content']}\n")
+        request_str = prompt.getvalue()
+
+        choices = [
+            f"\n\n**{choice['message']['role']}**: {choice['message']['content']}\n"
+            for choice in response["choices"]
+        ]
+
+        return self._resolve_metrics(
+            request=request,
+            response=response,
+            request_str=request_str,
+            choices=choices,
+            time_elapsed=time_elapsed,
+        )
+
+    def _resolve_metrics(
+        self,
+        request: Dict[str, Any],
+        response: Response,
+        request_str: str,
+        choices: List[str],
+        time_elapsed: float,
+    ) -> Dict[str, Any]:
+        """Resolves the request and response objects for `openai.Completion`."""
+        results = [
+            trace_tree.Result(
+                inputs={"request": request_str},
+                outputs={"response": choice},
+            )
+            for choice in choices
+        ]
+        metrics = self._get_metrics_to_log(request, response, results, time_elapsed)
+        return self._convert_metrics_to_dict(metrics)
+
+    @staticmethod
+    def _get_usage_metrics(response: Response, time_elapsed: float) -> UsageMetrics:
+        """Gets the usage stats from the response object."""
+        if response.get("usage"):
+            usage_stats = UsageMetrics(**response["usage"])
+        else:
+            usage_stats = UsageMetrics()
+        usage_stats.elapsed_time = time_elapsed
+        return usage_stats
+
+    def _get_metrics_to_log(
+        self,
+        request: Dict[str, Any],
+        response: Response,
+        results: List[Any],
+        time_elapsed: float,
+    ) -> Metrics:
+        model = response.get("model") or request.get("model")
+        usage_metrics = self._get_usage_metrics(response, time_elapsed)
+
+        usage = []
+        for result in results:
+            row = {
+                "request": result.inputs["request"],
+                "response": result.outputs["response"],
+                "model": model,
+                "start_time": datetime.datetime.fromtimestamp(response["created"]),
+                "end_time": datetime.datetime.fromtimestamp(
+                    response["created"] + time_elapsed
+                ),
+                "request_id": response.get("id", None),
+                "api_type": response.get("api_type", "openai"),
+                "session_id": wandb.run.id,
+            }
+            row.update(asdict(usage_metrics))
+            usage.append(row)
+        usage_table = wandb.Table(
+            columns=list(usage[0].keys()),
+            data=[(item.values()) for item in usage],
+        )
+
+        trace = self.results_to_trace_tree(request, response, results, time_elapsed)
+
+        metrics = Metrics(stats=usage_table, trace=trace, usage=usage_metrics)
+        return metrics
+
+    @staticmethod
+    def _convert_metrics_to_dict(metrics: Metrics) -> Dict[str, Any]:
+        """Converts metrics to a dict."""
+        metrics_dict = {
+            "stats": metrics.stats,
+            "trace": metrics.trace,
+        }
+        usage_stats = {f"usage/{k}": v for k, v in asdict(metrics.usage).items()}
+        metrics_dict.update(usage_stats)
+        return metrics_dict

wandb/integration/prodigy/__init__.py

@@ -0,0 +1,3 @@
+from .prodigy import upload_dataset
+
+__all__ = ["upload_dataset"]

wandb/integration/prodigy/prodigy.py

@@ -0,0 +1,299 @@
+"""Prodigy integration for W&B.
+
+User can upload Prodigy annotated datasets directly
+from the local database to W&B in Tables format.
+
+Example usage:
+
+```python
+import wandb
+from wandb.integration.prodigy import upload_dataset
+
+run = wandb.init(project="prodigy")
+upload_dataset("name_of_dataset")
+wandb.finish()
+```
+"""
+
+import base64
+import collections.abc
+import io
+import urllib
+from copy import deepcopy
+
+import pandas as pd
+from PIL import Image
+
+import wandb
+from wandb import util
+from wandb.plot.utils import test_missing
+from wandb.sdk.lib import telemetry as wb_telemetry
+
+
+def named_entity(docs):
+    """Create a named entity visualization.
+
+    Taken from https://github.com/wandb/wandb/blob/main/wandb/plots/named_entity.py.
+    """
+    spacy = util.get_module(
+        "spacy",
+        required="part_of_speech requires the spacy library, install with `pip install spacy`",
+    )
+
+    util.get_module(
+        "en_core_web_md",
+        required="part_of_speech requires `en_core_web_md` library, install with `python -m spacy download en_core_web_md`",
+    )
+
+    # Test for required packages and missing & non-integer values in docs data
+    if test_missing(docs=docs):
+        html = spacy.displacy.render(
+            docs, style="ent", page=True, minify=True, jupyter=False
+        )
+        wandb_html = wandb.Html(html)
+        return wandb_html
+
+
+def merge(dict1, dict2):
+    """Return a new dictionary by merging two dictionaries recursively."""
+    result = deepcopy(dict1)
+
+    for key, value in dict2.items():
+        if isinstance(value, collections.abc.Mapping):
+            result[key] = merge(result.get(key, {}), value)
+        else:
+            result[key] = deepcopy(dict2[key])
+
+    return result
+
+
+def get_schema(list_data_dict, struct, array_dict_types):
+    """Get a schema of the dataset's structure and data types."""
+    # Get the structure of the JSON objects in the database
+    # This is similar to getting a JSON schema but with slightly different format
+    for _i, item in enumerate(list_data_dict):
+        # If the list contains dict objects
+        for k, v in item.items():
+            # Check if key already exists in template
+            if k not in struct.keys():
+                if isinstance(v, list):
+                    if len(v) > 0 and isinstance(v[0], list):
+                        # nested list structure
+                        struct[k] = type(v)  # type list
+                    elif len(v) > 0 and not (
+                        isinstance(v[0], list) or isinstance(v[0], dict)
+                    ):
+                        # list of singular values
+                        struct[k] = type(v)  # type list
+                    else:
+                        # list of dicts
+                        array_dict_types.append(
+                            k
+                        )  # keep track of keys that are type list[dict]
+                        struct[k] = {}
+                        struct[k] = get_schema(v, struct[k], array_dict_types)
+                elif isinstance(v, dict):
+                    struct[k] = {}
+                    struct[k] = get_schema([v], struct[k], array_dict_types)
+                else:
+                    struct[k] = type(v)
+            else:
+                # Get the value of struct[k] which is the current template
+                # Find new keys and then merge the two templates together
+                cur_struct = struct[k]
+                if isinstance(v, list):
+                    if len(v) > 0 and isinstance(v[0], list):
+                        # nested list coordinate structure
+                        # if the value in the item is currently None, then update
+                        if v is not None:
+                            struct[k] = type(v)  # type list
+                    elif len(v) > 0 and not (
+                        isinstance(v[0], list) or isinstance(v[0], dict)
+                    ):
+                        # single list with values
+                        # if the value in the item is currently None, then update
+                        if v is not None:
+                            struct[k] = type(v)  # type list
+                    else:
+                        array_dict_types.append(
+                            k
+                        )  # keep track of keys that are type list[dict]
+                        struct[k] = {}
+                        struct[k] = get_schema(v, struct[k], array_dict_types)
+                        # merge cur_struct and struct[k], remove duplicates
+                        struct[k] = merge(struct[k], cur_struct)
+                elif isinstance(v, dict):
+                    struct[k] = {}
+                    struct[k] = get_schema([v], struct[k], array_dict_types)
+                    # merge cur_struct and struct[k], remove duplicates
+                    struct[k] = merge(struct[k], cur_struct)
+                else:
+                    # if the value in the item is currently None, then update
+                    if v is not None:
+                        struct[k] = type(v)
+
+    return struct
+
+
+def standardize(item, structure, array_dict_types):
+    """Standardize all rows/entries in dataset to fit the schema.
+
+    Will look for missing values and fill it in so all rows have
+    the same items and structure.
+    """
+    for k, v in structure.items():
+        if k not in item:
+            # If the structure/field does not exist
+            if isinstance(v, dict) and (k not in array_dict_types):
+                # If key k is of type dict, and not not a type list[dict]
+                item[k] = {}
+                standardize(item[k], v, array_dict_types)
+            elif isinstance(v, dict) and (k in array_dict_types):
+                # If key k is of type dict, and is actually of type list[dict],
+                # just treat as a list and set to None by default
+                item[k] = None
+            else:
+                # Assign a default type
+                item[k] = v()
+        else:
+            # If the structure/field already exists and is a list or dict
+            if isinstance(item[k], list):
+                # ignore if item is a nested list structure or list of non-dicts
+                condition = (
+                    not (len(item[k]) > 0 and isinstance(item[k][0], list))
+                ) and (
+                    not (
+                        len(item[k]) > 0
+                        and not (
+                            isinstance(item[k][0], list) or isinstance(item[k][0], dict)
+                        )
+                    )
+                )
+                if condition:
+                    for sub_item in item[k]:
+                        standardize(sub_item, v, array_dict_types)
+            elif isinstance(item[k], dict):
+                standardize(item[k], v, array_dict_types)
+
+
+def create_table(data):
+    """Create a W&B Table.
+
+    - Create/decode images from URL/Base64
+    - Uses spacy to translate NER span data to visualizations.
+    """
+    # create table object from columns
+    table_df = pd.DataFrame(data)
+    columns = list(table_df.columns)
+    if ("spans" in table_df.columns) and ("text" in table_df.columns):
+        columns.append("spans_visual")
+    if "image" in columns:
+        columns.append("image_visual")
+    main_table = wandb.Table(columns=columns)
+
+    # Convert to dictionary format to maintain order during processing
+    matrix = table_df.to_dict(orient="records")
+
+    # Import en_core_web_md if exists
+    en_core_web_md = util.get_module(
+        "en_core_web_md",
+        required="part_of_speech requires `en_core_web_md` library, install with `python -m spacy download en_core_web_md`",
+    )
+    nlp = en_core_web_md.load(disable=["ner"])
+
+    # Go through each individual row
+    for _i, document in enumerate(matrix):
+        # Text NER span visualizations
+        if ("spans_visual" in columns) and ("text" in columns):
+            # Add visuals for spans
+            document["spans_visual"] = None
+            doc = nlp(document["text"])
+            ents = []
+            if ("spans" in document) and (document["spans"] is not None):
+                for span in document["spans"]:
+                    if ("start" in span) and ("end" in span) and ("label" in span):
+                        charspan = doc.char_span(
+                            span["start"], span["end"], span["label"]
+                        )
+                        ents.append(charspan)
+                doc.ents = ents
+                document["spans_visual"] = named_entity(docs=doc)
+
+        # Convert image link to wandb Image
+        if "image" in columns:
+            # Turn into wandb image
+            document["image_visual"] = None
+            if ("image" in document) and (document["image"] is not None):
+                isurl = urllib.parse.urlparse(document["image"]).scheme in (
+                    "http",
+                    "https",
+                )
+                isbase64 = ("data:" in document["image"]) and (
+                    ";base64" in document["image"]
+                )
+                if isurl:
+                    # is url
+                    try:
+                        im = Image.open(urllib.request.urlopen(document["image"]))
+                        document["image_visual"] = wandb.Image(im)
+                    except urllib.error.URLError:
+                        print(
+                            "Warning: Image URL "
+                            + str(document["image"])
+                            + " is invalid."
+                        )
+                        document["image_visual"] = None
+                elif isbase64:
+                    # is base64 uri
+                    imgb64 = document["image"].split("base64,")[1]
+                    try:
+                        msg = base64.b64decode(imgb64)
+                        buf = io.BytesIO(msg)
+                        im = Image.open(buf)
+                        document["image_visual"] = wandb.Image(im)
+                    except base64.binascii.Error:
+                        print(
+                            "Warning: Base64 string "
+                            + str(document["image"])
+                            + " is invalid."
+                        )
+                        document["image_visual"] = None
+                else:
+                    # is data path
+                    document["image_visual"] = wandb.Image(document["image"])
+
+        # Create row and append to table
+        values_list = list(document.values())
+        main_table.add_data(*values_list)
+    return main_table
+
+
+def upload_dataset(dataset_name):
+    """Upload dataset from local database to Weights & Biases.
+
+    Args:
+        dataset_name: The name of the dataset in the Prodigy database.
+    """
+    # Check if wandb.init has been called
+    if wandb.run is None:
+        raise ValueError("You must call wandb.init() before upload_dataset()")
+
+    with wb_telemetry.context(run=wandb.run) as tel:
+        tel.feature.prodigy = True
+
+    prodigy_db = util.get_module(
+        "prodigy.components.db",
+        required="`prodigy` library is required but not installed. Please see https://prodi.gy/docs/install",
+    )
+    # Retrieve and upload prodigy dataset
+    database = prodigy_db.connect()
+    data = database.get_dataset(dataset_name)
+
+    array_dict_types = []
+    schema = get_schema(data, {}, array_dict_types)
+
+    for i, _d in enumerate(data):
+        standardize(data[i], schema, array_dict_types)
+    table = create_table(data)
+    wandb.log({dataset_name: table})
+    print("Prodigy dataset `" + dataset_name + "` uploaded.")

wandb/integration/sacred/__init__.py

@@ -0,0 +1,117 @@
+import warnings
+
+import numpy
+from sacred.dependencies import get_digest
+from sacred.observers import RunObserver
+
+import wandb
+
+
+class WandbObserver(RunObserver):
+    """Log sacred experiment data to W&B.
+
+    Arguments:
+        Accepts all the arguments accepted by wandb.init().
+
+        name — A display name for this run, which shows up in the UI and is editable, doesn't have to be unique
+        notes — A multiline string description associated with the run
+        config — a dictionary-like object to set as initial config
+        project — the name of the project to which this run will belong
+        tags — a list of strings to associate with this run as tags
+        dir — the path to a directory where artifacts will be written (default: ./wandb)
+        entity — the team posting this run (default: your username or your default team)
+        job_type — the type of job you are logging, e.g. eval, worker, ps (default: training)
+        save_code — save the main python or notebook file to wandb to enable diffing (default: editable from your settings page)
+        group — a string by which to group other runs; see Grouping
+        reinit — whether to allow multiple calls to wandb.init in the same process (default: False)
+        id — A unique ID for this run primarily used for Resuming. It must be globally unique, and if you delete a run you can't reuse the ID. Use the name field for a descriptive, useful name for the run. The ID cannot contain special characters.
+        resume — if set to True, the run auto resumes; can also be a unique string for manual resuming; see Resuming (default: False)
+        anonymous — can be "allow", "never", or "must". This enables or explicitly disables anonymous logging. (default: never)
+        force — whether to force a user to be logged into wandb when running a script (default: False)
+        magic — (bool, dict, or str, optional): magic configuration as bool, dict, json string, yaml filename. If set to True will attempt to auto-instrument your script. (default: None)
+        sync_tensorboard — A boolean indicating whether or not copy all TensorBoard logs wandb; see Tensorboard (default: False)
+        monitor_gym — A boolean indicating whether or not to log videos generated by OpenAI Gym; see Ray Tune (default: False)
+        allow_val_change — whether to allow wandb.config values to change, by default we throw an exception if config values are overwritten. (default: False)
+
+    Examples:
+        Create sacred experiment::
+        from wandb.sacred import WandbObserver
+        ex.observers.append(WandbObserver(project='sacred_test',
+                                                name='test1'))
+        @ex.config
+        def cfg():
+            C = 1.0
+            gamma = 0.7
+        @ex.automain
+        def run(C, gamma, _run):
+            iris = datasets.load_iris()
+            per = permutation(iris.target.size)
+            iris.data = iris.data[per]
+            iris.target = iris.target[per]
+            clf = svm.SVC(C, 'rbf', gamma=gamma)
+            clf.fit(iris.data[:90],
+                    iris.target[:90])
+            return clf.score(iris.data[90:],
+                                iris.target[90:])
+    """
+
+    def __init__(self, **kwargs):
+        self.run = wandb.init(**kwargs)
+        self.resources = {}
+
+    def started_event(
+        self, ex_info, command, host_info, start_time, config, meta_info, _id
+    ):
+        # TODO: add the source code file
+        # TODO: add dependencies and metadata.
+        self.__update_config(config)
+
+    def completed_event(self, stop_time, result):
+        if result:
+            if not isinstance(result, tuple):
+                result = (
+                    result,
+                )  # transform single result to tuple so that both single & multiple results use same code
+
+            for i, r in enumerate(result):
+                if isinstance(r, float) or isinstance(r, int):
+                    wandb.log({f"result_{i}": float(r)})
+                elif isinstance(r, dict):
+                    wandb.log(r)
+                elif isinstance(r, object):
+                    artifact = wandb.Artifact(f"result_{i}.pkl", type="result")
+                    artifact.add_file(r)
+                    self.run.log_artifact(artifact)
+                elif isinstance(r, numpy.ndarray):
+                    wandb.log({f"result_{i}": wandb.Image(r)})
+                else:
+                    warnings.warn(
+                        f"logging results does not support type '{type(r)}' results. Ignoring this result",
+                        stacklevel=2,
+                    )
+
+    def artifact_event(self, name, filename, metadata=None, content_type=None):
+        if content_type is None:
+            content_type = "file"
+        artifact = wandb.Artifact(name, type=content_type)
+        artifact.add_file(filename)
+        self.run.log_artifact(artifact)
+
+    def resource_event(self, filename):
+        """TODO: Maintain resources list."""
+        if filename not in self.resources:
+            md5 = get_digest(filename)
+            self.resources[filename] = md5
+
+    def log_metrics(self, metrics_by_name, info):
+        for metric_name, metric_ptr in metrics_by_name.items():
+            for _step, value in zip(metric_ptr["steps"], metric_ptr["values"]):
+                if isinstance(value, numpy.ndarray):
+                    wandb.log({metric_name: wandb.Image(value)})
+                else:
+                    wandb.log({metric_name: value})
+
+    def __update_config(self, config):
+        for k, v in config.items():
+            self.run.config[k] = v
+        self.run.config["resources"] = []

wandb/integration/sagemaker/__init__.py

@@ -0,0 +1,12 @@
+"""wandb integration sagemaker module."""
+
+from .auth import sagemaker_auth
+from .config import parse_sm_config
+from .resources import parse_sm_resources, parse_sm_secrets
+
+__all__ = [
+    "sagemaker_auth",
+    "parse_sm_config",
+    "parse_sm_secrets",
+    "parse_sm_resources",
+]