afnio.autodiff.lm_ops#

Classes

ChatCompletion(*args, **kwargs)

Implements a chat completion operation using the specified language model within the afnio framework, supporting automatic differentiation.
class afnio.autodiff.lm_ops.ChatCompletion(*args, **kwargs)[source]#

Bases: Function

Implements a chat completion operation using the specified language model within the afnio framework, supporting automatic differentiation.

This class inherits from Function and requires both the forward and backward methods to be defined.

Features:#

  • Mini-Batching: Processes multiple input dictionaries simultaneously to improve

    throughput.

  • Asynchronous Execution: Both the forward and backward passes are optimized to

    run asynchronous calls for each mini-batch, reducing latency.

  • Gradient Computation: Supports automatic differentiation for all Variables

    in messages and inputs arguments, maintaining the order of gradients.

The ChatCompletion function generates a Variable responses by passing a composite prompt, built from a list of messages and optional inputs, to the forward_model_client. Each message is a dictionary with a ‘role’ (e.g., ‘system’, ‘user’) and a list of Variable objects as ‘content’. inputs is a dictionary containing strings, list of strings or Variables providing dynamic values to fill placeholders within message templates. If inputs contain lists of strings or Variables which .data field is a list, the response’s .data field will be a list, corresponding to the batched results. Otherwise, the .data field will be a scalar string. Additional behavior, such as temperature or token limits, can be customized through completion_args.

Example with scalar inputs:

>>> system = Variable(
...     "You are a helpful assistant.",
...     role="system instruction",
...     requires_grad=True
... )
>>> user = Variable("Translate 'Hello' to {language}.", role="user query")
>>> messages = [
...     {"role": "system", "content": [system]},
...     {"role": "user", "content": [user]},
... ]
>>> inputs = {"language": Variable("Italian", role="language")}
>>> response = ChatCompletion.apply(
...     model_client,
...     messages,
...     inputs=inputs,
...     temperature=0.7
... )
>>> print(response.data)
'Ciao'
'Hola'
>>> feedback = Variable("Use only capital letters.", role="feedback")
>>> response.backward(feedback)
>>> system.grad[0].data
'The system instruction should enforce the use of capital letters only.'

Example with batched inputs:

>>> system = Variable(
...     "You are a helpful assistant.",
...     role="system instruction",
...     requires_grad=True
... )
>>> user = Variable("Translate 'Hello' to {language}.", role="user query")
>>> messages = [
...     {"role": "system", "content": [system]},
...     {"role": "user", "content": [user]},
... ]
>>> inputs = {
...     "language": [
...         Variable("Italian", role="language"),
...         Variable("Spanish", role="language")
...     ]
... }
>>> response = ChatCompletion.apply(
...     model_client,
...     messages,
...     inputs=inputs,
...     temperature=0.7
... )
>>> print(response.data)
['Ciao', 'Hola']
classmethod apply(*args, **kwargs)#

Applies the forward function of the custom Function class.

This method handles cases where setup_context is defined to set up the ctx (context) object separately or within the forward method itself.

static backward(ctx, grad_output)[source]#

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non variable outputs of the forward function), and it should return as many variables, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Variable or is a Variable not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve variables saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(ctx, forward_model_client, messages, inputs=None, **completion_args)[source]#

Define the forward of the custom autodiff Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

  • It must accept a context ctx as the first argument, followed by any number of arguments (variables or other types).

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

  • The forward no longer accepts a ctx argument.

  • Instead, you must also override the afnio.autodiff.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Variables should not be stored directly on ctx. Instead, variables should be saved either with ctx.save_for_backward() if they are intended to be used in backward.

static setup_context(ctx, inputs, output)#

There are two ways to define the forward pass of an autodiff.Function.

Either:

  1. Override forward with the signature forward(ctx, *args, **kwargs). setup_context is not overridden. Setting up the ctx for backward happens inside the forward.

  2. Override forward with the signature forward(*args, **kwargs) and override setup_context. Setting up the ctx for backward happens inside setup_context (as opposed to inside the forward)