Skip to content

huggingface_tracer

ChatCompletionOutput(impacts=None) dataclass

Bases: ChatCompletionOutput

Wrapper of huggingface_hub.ChatCompletionOutput with ImpactsOutput

ChatCompletionStreamOutput(impacts=None) dataclass

Bases: ChatCompletionStreamOutput

Wrapper of huggingface_hub.ChatCompletionStreamOutput with ImpactsOutput

HuggingfaceInstrumentor()

Instrumentor initialized by EcoLogits to automatically wrap all HuggingFace Hub calls

Source code in ecologits/tracers/huggingface_tracer.py
def __init__(self) -> None:
    self.wrapped_methods = [
        {
            "module": "huggingface_hub.inference._client",
            "name": "InferenceClient.chat_completion",
            "wrapper": huggingface_chat_wrapper,
        },
        {
            "module": "huggingface_hub.inference._generated._async_client",
            "name": "AsyncInferenceClient.chat_completion",
            "wrapper": huggingface_async_chat_wrapper,
        },
    ]

huggingface_chat_wrapper(wrapped, instance, args, kwargs)

Function that wraps a HuggingFace Hub answer with computed impacts

Parameters:

Name Type Description Default
wrapped Callable

Callable that returns the LLM response

required
instance InferenceClient

Never used - for compatibility with wrapt

required
args Any

Arguments of the callable

required
kwargs Any

Keyword arguments of the callable

required

Returns:

Type Description
Union[ChatCompletionOutput, Iterable[ChatCompletionStreamOutput]]

A wrapped ChatCompletionOutput or Iterable[ChatCompletionStreamOutput] with impacts

Source code in ecologits/tracers/huggingface_tracer.py
def huggingface_chat_wrapper(
    wrapped: Callable,
    instance: InferenceClient,
    args: Any,
    kwargs: Any
) -> Union[ChatCompletionOutput, Iterable[ChatCompletionStreamOutput]]:
    """
    Function that wraps a HuggingFace Hub answer with computed impacts

    Args:
        wrapped: Callable that returns the LLM response
        instance: Never used - for compatibility with `wrapt`
        args: Arguments of the callable
        kwargs: Keyword arguments of the callable

    Returns:
        A wrapped `ChatCompletionOutput` or `Iterable[ChatCompletionStreamOutput]` with impacts
    """

    if kwargs.get("stream", False):
        return huggingface_chat_wrapper_stream(wrapped, instance, args, kwargs)
    else:
        return huggingface_chat_wrapper_non_stream(wrapped, instance, args, kwargs)

huggingface_async_chat_wrapper(wrapped, instance, args, kwargs) async

Function that wraps a HuggingFace Hub answer with computed impacts in async mode

Parameters:

Name Type Description Default
wrapped Callable

Async callable that returns the LLM response

required
instance AsyncInferenceClient

Never used - for compatibility with wrapt

required
args Any

Arguments of the callable

required
kwargs Any

Keyword arguments of the callable

required

Returns:

Type Description
Union[ChatCompletionOutput, AsyncIterable[ChatCompletionStreamOutput]]

A wrapped ChatCompletionOutput or AsyncIterable[ChatCompletionStreamOutput]] with impacts

Source code in ecologits/tracers/huggingface_tracer.py
async def huggingface_async_chat_wrapper(
    wrapped: Callable,
    instance: AsyncInferenceClient,
    args: Any,
    kwargs: Any
) -> Union[ChatCompletionOutput, AsyncIterable[ChatCompletionStreamOutput]]:
    """
    Function that wraps a HuggingFace Hub answer with computed impacts in async mode

    Args:
        wrapped: Async callable that returns the LLM response
        instance: Never used - for compatibility with `wrapt`
        args: Arguments of the callable
        kwargs: Keyword arguments of the callable

    Returns:
        A wrapped `ChatCompletionOutput` or `AsyncIterable[ChatCompletionStreamOutput]]` with impacts
    """

    if kwargs.get("stream", False):
        return huggingface_async_chat_wrapper_stream(wrapped, instance, args, kwargs)
    else:
        return await huggingface_async_chat_wrapper_non_stream(wrapped, instance, args, kwargs)