wandb 0.15.3__py3-none-any.whl → 0.15.5__py3-none-any.whl

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 explicity 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

wandb/integration/langchain/wandb_tracer.py

@@ -14,22 +14,10 @@ integration will not break user code. The one exception to the rule is at import
 LangChain is not installed, or the symbols are not in the same place, the appropriate error
 will be raised when importing this module.
 """
-import sys
-
-if sys.version_info >= (3, 8):
-    from typing import TypedDict
-else:
-    from typing_extensions import TypedDict
-
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence, Union
-
 from packaging import version
 
-import wandb
 import wandb.util
-from wandb.sdk.data_types import trace_tree
-from wandb.sdk.lib import telemetry as wb_telemetry
-from wandb.sdk.lib.paths import StrPath
+from wandb.sdk.lib import deprecate
 
 langchain = wandb.util.get_module(
     name="langchain",
@@ -37,174 +25,23 @@ langchain = wandb.util.get_module(
     "package installed. Please install it with `pip install langchain`.",
 )
 
-if version.parse(langchain.__version__) < version.parse("0.0.170"):
+if version.parse(langchain.__version__) < version.parse("0.0.188"):
     raise ValueError(
-        "The Weights & Biases Langchain integration does not support versions 0.0.169 and lower. "
-        "To ensure proper functionality, please use version 0.0.170 or higher."
+        "The Weights & Biases Langchain integration does not support versions 0.0.187 and lower. "
+        "To ensure proper functionality, please use version 0.0.188 or higher."
     )
 
-# We want these imports after the import_langchain() call, so that we can
-# catch the ImportError if langchain is not installed.
-
 # isort: off
-from langchain.callbacks.tracers.base import BaseTracer  # noqa: E402, I001
-
-from .util import (  # noqa: E402
-    print_wandb_init_message,
-    safely_convert_lc_run_to_wb_span,
-    # safely_convert_model_to_dict,
-    # safely_get_span_producing_model,
-)
-
-if TYPE_CHECKING:
-    from langchain.callbacks.base import BaseCallbackHandler
-    from langchain.callbacks.tracers.schemas import Run
-
-    from wandb import Settings as WBSettings
-    from wandb.wandb_run import Run as WBRun
-
-
-class WandbRunArgs(TypedDict):
-    job_type: Optional[str]
-    dir: Optional[StrPath]
-    config: Union[Dict, str, None]
-    project: Optional[str]
-    entity: Optional[str]
-    reinit: Optional[bool]
-    tags: Optional[Sequence]
-    group: Optional[str]
-    name: Optional[str]
-    notes: Optional[str]
-    magic: Optional[Union[dict, str, bool]]
-    config_exclude_keys: Optional[List[str]]
-    config_include_keys: Optional[List[str]]
-    anonymous: Optional[str]
-    mode: Optional[str]
-    allow_val_change: Optional[bool]
-    resume: Optional[Union[bool, str]]
-    force: Optional[bool]
-    tensorboard: Optional[bool]
-    sync_tensorboard: Optional[bool]
-    monitor_gym: Optional[bool]
-    save_code: Optional[bool]
-    id: Optional[str]
-    settings: Union["WBSettings", Dict[str, Any], None]
-
-
-class WandbTracer(BaseTracer):
-    """Callback Handler that logs to Weights and Biases.
-
-    This handler will log the model architecture and run traces to Weights and Biases. This will
-    ensure that all LangChain activity is logged to W&B.
-    """
-
-    _run: Optional["WBRun"] = None
-    _run_args: Optional[WandbRunArgs] = None
-
-    @classmethod
-    def init(
-        cls,
-        run_args: Optional[WandbRunArgs] = None,
-        include_stdout: bool = True,
-        additional_handlers: Optional[List["BaseCallbackHandler"]] = None,
-    ) -> None:
-        """Method provided for backwards compatibility. Please directly construct `WandbTracer` instead."""
-        message = """Global autologging is not currently supported for the LangChain integration.
-Please directly construct a `WandbTracer` and add it to the list of callbacks. For example:
-
-LLMChain(llm, callbacks=[WandbTracer()])
-# end of notebook / script:
-WandbTracer.finish()"""
-        wandb.termlog(message)
-
-    def __init__(self, run_args: Optional[WandbRunArgs] = None, **kwargs: Any) -> None:
-        """Initializes the WandbTracer.
-
-        Parameters:
-            run_args: (dict, optional) Arguments to pass to `wandb.init()`. If not provided, `wandb.init()` will be
-                called with no arguments. Please refer to the `wandb.init` for more details.
-
-        To use W&B to monitor all LangChain activity, add this tracer like any other langchain callback
-        ```
-        from wandb.integration.langchain import WandbTracer
-        LLMChain(llm, callbacks=[WandbTracer()])
-        # end of notebook / script:
-        WandbTracer.finish()
-        ```.
-        """
-        super().__init__(**kwargs)
-        self._run_args = run_args
-        self._ensure_run(should_print_url=(wandb.run is None))
-
-    @staticmethod
-    def finish() -> None:
-        """Waits for all asynchronous processes to finish and data to upload.
-
-        Proxy for `wandb.finish()`.
-        """
-        wandb.finish()
-
-    def _log_trace_from_run(self, run: "Run") -> None:
-        """Logs a LangChain Run to W*B as a W&B Trace."""
-        self._ensure_run()
-
-        root_span = safely_convert_lc_run_to_wb_span(run)
-        if root_span is None:
-            return
-
-        model_dict = None
-
-        # TODO: Uncomment this once we have a way to get the model from a run
-        # model = safely_get_span_producing_model(run)
-        # if model is not None:
-        #     model_dict = safely_convert_model_to_dict(model)
-
-        model_trace = trace_tree.WBTraceTree(
-            root_span=root_span,
-            model_dict=model_dict,
+from langchain.callbacks.tracers import WandbTracer  # noqa: E402, I001
+
+
+class WandbTracer(WandbTracer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        deprecate.deprecate(
+            field_name=deprecate.Deprecated.langchain_tracer,
+            warning_message="This feature is deprecated and has been moved to `langchain`. Enable tracing by setting "
+            "LANGCHAIN_WANDB_TRACING=true in your environment. See the documentation at "
+            "https://python.langchain.com/docs/ecosystem/integrations/agent_with_wandb_tracing for guidance. "
+            "Replace your current import with `from langchain.callbacks.tracers import WandbTracer`.",
         )
-        wandb.run.log({"langchain_trace": model_trace})
-
-    def _ensure_run(self, should_print_url=False) -> None:
-        """Ensures an active W&B run exists.
-
-        If not, will start a new run with the provided run_args.
-        """
-        if wandb.run is None:
-            # Make a shallow copy of the run args, so we don't modify the original
-            run_args = self._run_args or {}  # type: ignore
-            run_args: dict = {**run_args}  # type: ignore
-
-            # Prefer to run in silent mode since W&B has a lot of output
-            # which can be undesirable when dealing with text-based models.
-            if "settings" not in run_args:  # type: ignore
-                run_args["settings"] = {"silent": True}  # type: ignore
-
-            # Start the run and add the stream table
-            wandb.init(**run_args)
-
-            if should_print_url:
-                print_wandb_init_message(wandb.run.settings.run_url)
-
-        with wb_telemetry.context(wandb.run) as tel:
-            tel.feature.langchain_tracer = True
-
-    # Start of required methods (these methods are required by the BaseCallbackHandler interface)
-    @property
-    def always_verbose(self) -> bool:
-        """Whether to call verbose callbacks even if verbose is False."""
-        return True
-
-    def _generate_id(self) -> Optional[Union[int, str]]:
-        """Generate an id for a run."""
-        return None
-
-    def _persist_run(self, run: "Run") -> None:
-        """Persist a run."""
-        try:
-            self._log_trace_from_run(run)
-        except Exception:
-            # Silently ignore errors to not break user code
-            pass
-
-    # End of required methods

wandb/integration/openai/__init__.py

@@ -1,5 +1,3 @@
 __all__ = ("autolog",)
 
-from .openai import AutologOpenAI
-
-autolog = AutologOpenAI()
+from .openai import autolog

wandb/integration/openai/openai.py

@@ -1,151 +1,19 @@
-import functools
 import logging
-import sys
-from typing import Any, Dict, List, Optional
 
-import wandb.sdk
-import wandb.util
-from wandb.sdk.lib import telemetry as wb_telemetry
-from wandb.sdk.lib.timer import Timer
+from wandb.sdk.integration_utils.auto_logging import AutologAPI
 
 from .resolver import OpenAIRequestResponseResolver
 
-if sys.version_info >= (3, 8):
-    from typing import Literal
-else:
-    from typing_extensions import Literal
-
-
 logger = logging.getLogger(__name__)
 
 
-AutologOpenAIInitArgs = Optional[Dict[str, Any]]
-
-
-class PatchOpenAIAPI:
-    symbols: List[Literal["Edit", "Completion", "ChatCompletion"]] = [
-        "Edit",
-        "Completion",
-        "ChatCompletion",
-    ]
-
-    def __init__(self) -> None:
-        """Patches the OpenAI API to log traces to W&B."""
-        self.original_methods: Dict[str, Any] = {}
-        self.resolver = OpenAIRequestResponseResolver()
-        self._openai = None
-
-    @property
-    def openai(self) -> Any:
-        """Returns the openai module."""
-        if self._openai is None:
-            self._openai = wandb.util.get_module(
-                name="openai",
-                required="To use the W&B OpenAI Autolog, you need to have the `openai` python "
-                "package installed. Please install it with `pip install openai`.",
-                lazy=False,
-            )
-        return self._openai
-
-    def patch(self, run: "wandb.sdk.wandb_run.Run") -> None:
-        """Patches the OpenAI API to log traces to W&B."""
-        for symbol in self.symbols:
-            original = getattr(self.openai, symbol).create
-
-            def method_factory(original_method: Any):
-                @functools.wraps(original_method)
-                def create(*args, **kwargs):
-                    with Timer() as timer:
-                        result = original_method(*args, **kwargs)
-                    try:
-                        trace = self.resolver(kwargs, result, timer.elapsed)
-                        if trace is not None:
-                            run.log({"trace": trace})
-                    except Exception:
-                        # logger.warning(e)
-                        pass
-                    return result
-
-                return create
-
-            # save original method
-            self.original_methods[symbol] = original
-            # monkeypatch
-            getattr(self.openai, symbol).create = method_factory(original)
-
-    def unpatch(self) -> None:
-        """Unpatches the OpenAI API."""
-        for symbol, original in self.original_methods.items():
-            getattr(self.openai, symbol).create = original
-
-
-class AutologOpenAI:
-    def __init__(self) -> None:
-        """Autolog OpenAI API calls to W&B."""
-        self._patch_openai_api = PatchOpenAIAPI()
-        self._run: Optional["wandb.sdk.wandb_run.Run"] = None
-        self.__run_created_by_autolog: bool = False
-
-    @property
-    def _is_enabled(self) -> bool:
-        """Returns whether autologging is enabled."""
-        return self._run is not None
-
-    def __call__(self, init: AutologOpenAIInitArgs = None) -> None:
-        """Enable OpenAI autologging."""
-        self.enable(init=init)
-
-    def _run_init(self, init: AutologOpenAIInitArgs = None) -> None:
-        """Handle wandb run initialization."""
-        # - autolog(init: dict = {...}) calls wandb.init(**{...})
-        #   regardless of whether there is a wandb.run or not,
-        #   we only track if the run was created by autolog
-        #    - todo: autolog(init: dict | run = run) would use the user-provided run
-        # - autolog() uses the wandb.run if there is one, otherwise it calls wandb.init()
-        if init:
-            _wandb_run = wandb.run
-            # we delegate dealing with the init dict to wandb.init()
-            self._run = wandb.init(**init)
-            if _wandb_run != self._run:
-                self.__run_created_by_autolog = True
-        elif wandb.run is None:
-            self._run = wandb.init()
-            self.__run_created_by_autolog = True
-        else:
-            self._run = wandb.run
-
-    def enable(self, init: AutologOpenAIInitArgs = None) -> None:
-        """Enable OpenAI autologging.
-
-        Args:
-            init: Optional dictionary of arguments to pass to wandb.init().
-
-        """
-        if self._is_enabled:
-            logger.info(
-                "OpenAI autologging is already enabled, disabling and re-enabling."
-            )
-            self.disable()
-
-        logger.info("Enabling OpenAI autologging.")
-        self._run_init(init=init)
-
-        self._patch_openai_api.patch(self._run)
-
-        with wb_telemetry.context(self._run) as tel:
-            tel.feature.openai_autolog = True
-
-    def disable(self) -> None:
-        """Disable OpenAI autologging."""
-        if self._run is None:
-            return
-
-        logger.info("Disabling OpenAI autologging.")
-
-        if self.__run_created_by_autolog:
-            self._run.finish()
-            self.__run_created_by_autolog = False
-
-        self._run = None
-
-        self._patch_openai_api.unpatch()
+autolog = AutologAPI(
+    name="OpenAI",
+    symbols=(
+        "Edit.create",
+        "Completion.create",
+        "ChatCompletion.create",
+    ),
+    resolver=OpenAIRequestResponseResolver(),
+    telemetry_feature="openai_autolog",
+)