Skip to content

processors

Processors to control generation in steerable models.

base_logits_processor

Base class for logits processors.

OutlinesLogitsProcessor

Base class for logits processors. This class implements a shared __call__ method is called by the models and returns the processed logits. It relies on the process_logits method that must be implemented by the subclasses to do the actual processing. The tensor_adapter attribute, created at initialization based on the tensor library name specified in the constructor, is used to manipulate the tensors using the appropriate library for the model (numpy, torch...).

Source code in outlines/processors/base_logits_processor.py
class OutlinesLogitsProcessor:
    """Base class for logits processors.
    This class implements a shared `__call__` method is called by the models
    and returns the processed logits. It relies on the `process_logits` method
    that must be implemented by the subclasses to do the actual processing. The
    `tensor_adapter` attribute, created at initialization based on the
    tensor library name specified in the constructor, is used to manipulate the
    tensors using the appropriate library for the model (numpy, torch...).
    """
    tensor_adapter: TensorAdapterImplementation

    def __init__(self, tensor_library_name: str):
        """
        Parameters
        ----------
        tensor_library_name
            The name of the library to use to manipulate tensors. Possible
            values are "jax", "mlx", "numpy", "tensorflow" and "torch". You
            must choose the library that your model is using.
        """
        # Temporary fix as torch raises a warning that can cause can an error
        # with python 3.12.
        if tensor_library_name == "torch":
            import torch._dynamo

            torch._dynamo.config.suppress_errors = True

        tensor_adapter_class = tensor_adapters.get(tensor_library_name)
        if tensor_adapter_class is None:
            raise NotImplementedError(
                f"Library {tensor_library_name} is not available"
            )
        self.tensor_adapter = tensor_adapter_class()  # type: ignore

    def reset(self):
        """Reset the logits processor for a new generation

        Only implement this method in subclasses if the logits processor
        needs to be reset for a new generation.

        """
        pass

    @abstractmethod
    def process_logits(
        self, input_ids: TensorType, logits: TensorType
    ) -> TensorType:
        """Main method to implement for logits processors subclasses.
        This method applies a mask on the logits to bias the generation.
        It is called by the `__call__` method that standardizes the shape of
        `input_ids` and `logits` to ensure they are 2D tensors.
        Elements to keep in mind when designing universal logits processors:
        - logits processors are only used once and never re-applied for a new
        sequence generator
        - Some models only pass output_ids, some models such as llamacpp and
        transformers prefix with input_ids
        - Some sampling methods, such as beam search, result in unstable
        sequence ordering in models like vLLM
        Parameters
        ----------
        input_ids
            The ids of the tokens of the existing sequences in a 2D tensor.
        logits
            The logits for the current generation step in a 2D tensor.
        Returns
        -------
        TensorType
            The processed logits as a 2D tensor.
        """
        ...

    def __call__(
        self, input_ids: TensorType, logits: TensorType
    ) -> TensorType:
        """Entrypoint for logits processors, this is the method that is
        called by the model.
        Because different models use different structures to store the
        input_ids and logits, we standardize their format to 2D tensors
        before calling the `process_logits` method. After processing, the
        logits are cast back to the original array library type before being
        returned.
        Parameters
        ----------
        input_ids
            The ids of the tokens of the existing sequences in a tensor.
        logits
            The logits for the current generation step in a tensor.
        Returns
        -------
        TensorType
            The processed logits as a tensor.
        """
        # if input_ids is 1D and logits is 2D with a single sequence,
        # reshape input_ids to 2D (needed for mlx-lm)
        if (
            len(self.tensor_adapter.shape(input_ids)) == 1
            and len(self.tensor_adapter.shape(logits)) == 2
            and self.tensor_adapter.shape(logits)[0] == 1
        ):
            input_ids = self.tensor_adapter.unsqueeze(input_ids)

        assert (
            self.tensor_adapter.shape(logits)[:-1]
            == self.tensor_adapter.shape(input_ids)[:-1]
        )

        # Guarantee passed as 2D Tensors, then covert back to original
        # (1D or 2D) shape
        if len(self.tensor_adapter.shape(logits)) == 2:
            processed_logits = self.process_logits(input_ids, logits)
        elif len(self.tensor_adapter.shape(logits)) == 1:
            processed_logits = self.tensor_adapter.squeeze(
                self.process_logits(
                    self.tensor_adapter.unsqueeze(input_ids),
                    self.tensor_adapter.unsqueeze(logits),
                ),
            )
        else:
            raise ValueError(
                f"Logits shape {self.tensor_adapter.shape(logits)} is not "
                + "supported"
            )

        return processed_logits

__call__(input_ids, logits)

Entrypoint for logits processors, this is the method that is called by the model. Because different models use different structures to store the input_ids and logits, we standardize their format to 2D tensors before calling the process_logits method. After processing, the logits are cast back to the original array library type before being returned.

Parameters:

Name Type Description Default
input_ids TensorType

The ids of the tokens of the existing sequences in a tensor.

required
logits TensorType

The logits for the current generation step in a tensor.

required

Returns:

Type Description
TensorType

The processed logits as a tensor.

Source code in outlines/processors/base_logits_processor.py
def __call__(
    self, input_ids: TensorType, logits: TensorType
) -> TensorType:
    """Entrypoint for logits processors, this is the method that is
    called by the model.
    Because different models use different structures to store the
    input_ids and logits, we standardize their format to 2D tensors
    before calling the `process_logits` method. After processing, the
    logits are cast back to the original array library type before being
    returned.
    Parameters
    ----------
    input_ids
        The ids of the tokens of the existing sequences in a tensor.
    logits
        The logits for the current generation step in a tensor.
    Returns
    -------
    TensorType
        The processed logits as a tensor.
    """
    # if input_ids is 1D and logits is 2D with a single sequence,
    # reshape input_ids to 2D (needed for mlx-lm)
    if (
        len(self.tensor_adapter.shape(input_ids)) == 1
        and len(self.tensor_adapter.shape(logits)) == 2
        and self.tensor_adapter.shape(logits)[0] == 1
    ):
        input_ids = self.tensor_adapter.unsqueeze(input_ids)

    assert (
        self.tensor_adapter.shape(logits)[:-1]
        == self.tensor_adapter.shape(input_ids)[:-1]
    )

    # Guarantee passed as 2D Tensors, then covert back to original
    # (1D or 2D) shape
    if len(self.tensor_adapter.shape(logits)) == 2:
        processed_logits = self.process_logits(input_ids, logits)
    elif len(self.tensor_adapter.shape(logits)) == 1:
        processed_logits = self.tensor_adapter.squeeze(
            self.process_logits(
                self.tensor_adapter.unsqueeze(input_ids),
                self.tensor_adapter.unsqueeze(logits),
            ),
        )
    else:
        raise ValueError(
            f"Logits shape {self.tensor_adapter.shape(logits)} is not "
            + "supported"
        )

    return processed_logits

__init__(tensor_library_name)

Parameters:

Name Type Description Default
tensor_library_name str

The name of the library to use to manipulate tensors. Possible values are "jax", "mlx", "numpy", "tensorflow" and "torch". You must choose the library that your model is using.

required
Source code in outlines/processors/base_logits_processor.py
def __init__(self, tensor_library_name: str):
    """
    Parameters
    ----------
    tensor_library_name
        The name of the library to use to manipulate tensors. Possible
        values are "jax", "mlx", "numpy", "tensorflow" and "torch". You
        must choose the library that your model is using.
    """
    # Temporary fix as torch raises a warning that can cause can an error
    # with python 3.12.
    if tensor_library_name == "torch":
        import torch._dynamo

        torch._dynamo.config.suppress_errors = True

    tensor_adapter_class = tensor_adapters.get(tensor_library_name)
    if tensor_adapter_class is None:
        raise NotImplementedError(
            f"Library {tensor_library_name} is not available"
        )
    self.tensor_adapter = tensor_adapter_class()  # type: ignore

process_logits(input_ids, logits) abstractmethod

Main method to implement for logits processors subclasses. This method applies a mask on the logits to bias the generation. It is called by the __call__ method that standardizes the shape of input_ids and logits to ensure they are 2D tensors. Elements to keep in mind when designing universal logits processors: - logits processors are only used once and never re-applied for a new sequence generator - Some models only pass output_ids, some models such as llamacpp and transformers prefix with input_ids - Some sampling methods, such as beam search, result in unstable sequence ordering in models like vLLM

Parameters:

Name Type Description Default
input_ids TensorType

The ids of the tokens of the existing sequences in a 2D tensor.

required
logits TensorType

The logits for the current generation step in a 2D tensor.

required

Returns:

Type Description
TensorType

The processed logits as a 2D tensor.

Source code in outlines/processors/base_logits_processor.py
@abstractmethod
def process_logits(
    self, input_ids: TensorType, logits: TensorType
) -> TensorType:
    """Main method to implement for logits processors subclasses.
    This method applies a mask on the logits to bias the generation.
    It is called by the `__call__` method that standardizes the shape of
    `input_ids` and `logits` to ensure they are 2D tensors.
    Elements to keep in mind when designing universal logits processors:
    - logits processors are only used once and never re-applied for a new
    sequence generator
    - Some models only pass output_ids, some models such as llamacpp and
    transformers prefix with input_ids
    - Some sampling methods, such as beam search, result in unstable
    sequence ordering in models like vLLM
    Parameters
    ----------
    input_ids
        The ids of the tokens of the existing sequences in a 2D tensor.
    logits
        The logits for the current generation step in a 2D tensor.
    Returns
    -------
    TensorType
        The processed logits as a 2D tensor.
    """
    ...

reset()

Reset the logits processor for a new generation

Only implement this method in subclasses if the logits processor needs to be reset for a new generation.

Source code in outlines/processors/base_logits_processor.py
def reset(self):
    """Reset the logits processor for a new generation

    Only implement this method in subclasses if the logits processor
    needs to be reset for a new generation.

    """
    pass

tensor_adapters

Library specific objects to manipulate tensors.

base

Base class for tensor adapters.

TensorAdapter

Bases: ABC

Abstract base class for tensor adapters.

This class defines the interface for tensor adapters that are used to manipulate tensors in different libraries. Concrete implementations of this class should provide specific implementations for each method as well as providing a library_name attribute.

TODO: Update the version of outlines-core used to receive plain arrays instead of torch tensors. In the meantime, implementations of this class must make sure that their full_like and concatenate methods can handle torch tensors.

Source code in outlines/processors/tensor_adapters/base.py
class TensorAdapter(ABC):
    """Abstract base class for tensor adapters.

    This class defines the interface for tensor adapters that are used to
    manipulate tensors in different libraries. Concrete implementations of
    this class should provide specific implementations for each method as
    well as providing a `library_name` attribute.

    TODO: Update the version of outlines-core used to receive plain arrays
    instead of torch tensors. In the meantime, implementations of this class
    must make sure that their `full_like` and `concatenate` methods can
    handle torch tensors.

    """
    library_name: str

    @abstractmethod
    def shape(self, tensor: TensorType) -> list[int]:
        """Get the shape of the tensor.

        Parameters
        ----------
        tensor
            The tensor to get the shape of.

        Returns
        -------
        list[int]
            The shape of the tensor. The list contains as many elements as
            there are dimensions in the tensor.

        """
        ...

    @abstractmethod
    def unsqueeze(self, tensor: TensorType) -> TensorType:
        """Add a dimension to the tensor at axis 0.

        Parameters
        ----------
        tensor
            The tensor to add a dimension to.

        Returns
        -------
        TensorType
            The tensor with an additional dimension.

        """
        ...

    @abstractmethod
    def squeeze(self, tensor: TensorType) -> TensorType:
        """Remove a dimension from the tensor at axis 0.

        Parameters
        ----------
        tensor
            The tensor to remove a dimension from.

        Returns
        -------
        TensorType
            The tensor with one less dimension.

        """
        ...

    @abstractmethod
    def to_list(self, tensor: TensorType) -> list:
        """Convert the tensor to a list.

        Parameters
        ----------
        tensor
            The tensor to convert to a list.

        Returns
        -------
        list
            The tensor as a list.

        """
        ...

    @abstractmethod
    def to_scalar(self, tensor: TensorType) -> Any:
        """Return the only element of the tensor.

        Parameters
        ----------
        tensor
            The tensor to return the only element of.

        Returns
        -------
        Any
            The only element of the tensor.

        """
        ...

    @abstractmethod
    def full_like(self, tensor: "torch.Tensor", fill_value: Any) -> TensorType: # type: ignore
        """Create a tensor with the same shape as the input tensor filled
        with a scalar value.

        ATTENTION: This method receives a torch tensor regardless of the
        library used.

        Parameters
        ----------
        tensor
            The tensor to create a new tensor with the same shape.
        fill_value
            The value to fill the new tensor with.

        Returns
        -------
        TensorType
            A tensor with the same shape as the input tensor filled with the
            specified value.

        """
        ...

    @abstractmethod
    def concatenate(
        self, tensors: list[Union["torch.Tensor", TensorType]]
    ) -> TensorType:
        """Concatenate a list of tensors along axis 0.

        ATTENTION: This method can either receive a list of torch tensors or
        a list of tensors from the library used.

        Parameters
        ----------
        tensors
            The list of tensors to concatenate.

        Returns
        -------
        TensorType
            The concatenated tensor.

        """
        ...

    @abstractmethod
    def get_device(self, tensor: TensorType) -> str:
        """Get the name of the tensor's device.

        Parameters
        ----------
        tensor
            The tensor to get the device of.

        Returns
        -------
        str
            The name of the tensor's device.

        """
        ...

    @abstractmethod
    def to_device(self, tensor: TensorType, device: str) -> TensorType:
        """Move the tensor to a specified device.

        Parameters
        ----------
        tensor
            The tensor to move to a specified device.
        device
            The name of the device to move the tensor to.

        Returns
        -------
        TensorType
            The tensor moved to the specified device.

        """
        ...

    @abstractmethod
    def boolean_ones_like(self, tensor: TensorType) -> TensorType:
        """Create a boolean ones tensor with the same shape as the input
        tensor.

        Parameters
        ----------
        tensor
            The tensor to create a boolean ones tensor with the same shape.

        Returns
        -------
        TensorType
            A boolean ones tensor with the same shape as the input tensor.

        """
        ...

    @abstractmethod
    def apply_mask(
        self, tensor: TensorType, mask: TensorType, value: Any
    ) -> TensorType:
        """Fill the elements of the tensor where the mask is True with the
        specified value.

        Parameters
        ----------
        tensor
            The tensor to fill.
        mask
            The mask to apply to the tensor.
        value
            The value to fill the tensor with.

        Returns
        -------
        TensorType
            The tensor with the mask applied.

        """
        ...

    @abstractmethod
    def argsort_descending(
        self, tensor: TensorType
    ) -> TensorType:
        """Return the indices that would sort the tensor in descending order
        along axis -1.

        Parameters
        ----------
        tensor
            The tensor to sort.

        Returns
        -------
        TensorType
            The indices that would sort the tensor in descending order along
            axis -1.

        """
        ...
apply_mask(tensor, mask, value) abstractmethod

Fill the elements of the tensor where the mask is True with the specified value.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to fill.

required
mask TensorType

The mask to apply to the tensor.

required
value Any

The value to fill the tensor with.

required

Returns:

Type Description
TensorType

The tensor with the mask applied.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def apply_mask(
    self, tensor: TensorType, mask: TensorType, value: Any
) -> TensorType:
    """Fill the elements of the tensor where the mask is True with the
    specified value.

    Parameters
    ----------
    tensor
        The tensor to fill.
    mask
        The mask to apply to the tensor.
    value
        The value to fill the tensor with.

    Returns
    -------
    TensorType
        The tensor with the mask applied.

    """
    ...
argsort_descending(tensor) abstractmethod

Return the indices that would sort the tensor in descending order along axis -1.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to sort.

required

Returns:

Type Description
TensorType

The indices that would sort the tensor in descending order along axis -1.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def argsort_descending(
    self, tensor: TensorType
) -> TensorType:
    """Return the indices that would sort the tensor in descending order
    along axis -1.

    Parameters
    ----------
    tensor
        The tensor to sort.

    Returns
    -------
    TensorType
        The indices that would sort the tensor in descending order along
        axis -1.

    """
    ...
boolean_ones_like(tensor) abstractmethod

Create a boolean ones tensor with the same shape as the input tensor.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to create a boolean ones tensor with the same shape.

required

Returns:

Type Description
TensorType

A boolean ones tensor with the same shape as the input tensor.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def boolean_ones_like(self, tensor: TensorType) -> TensorType:
    """Create a boolean ones tensor with the same shape as the input
    tensor.

    Parameters
    ----------
    tensor
        The tensor to create a boolean ones tensor with the same shape.

    Returns
    -------
    TensorType
        A boolean ones tensor with the same shape as the input tensor.

    """
    ...
concatenate(tensors) abstractmethod

Concatenate a list of tensors along axis 0.

ATTENTION: This method can either receive a list of torch tensors or a list of tensors from the library used.

Parameters:

Name Type Description Default
tensors list[Union[Tensor, TensorType]]

The list of tensors to concatenate.

required

Returns:

Type Description
TensorType

The concatenated tensor.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def concatenate(
    self, tensors: list[Union["torch.Tensor", TensorType]]
) -> TensorType:
    """Concatenate a list of tensors along axis 0.

    ATTENTION: This method can either receive a list of torch tensors or
    a list of tensors from the library used.

    Parameters
    ----------
    tensors
        The list of tensors to concatenate.

    Returns
    -------
    TensorType
        The concatenated tensor.

    """
    ...
full_like(tensor, fill_value) abstractmethod

Create a tensor with the same shape as the input tensor filled with a scalar value.

ATTENTION: This method receives a torch tensor regardless of the library used.

Parameters:

Name Type Description Default
tensor Tensor

The tensor to create a new tensor with the same shape.

required
fill_value Any

The value to fill the new tensor with.

required

Returns:

Type Description
TensorType

A tensor with the same shape as the input tensor filled with the specified value.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def full_like(self, tensor: "torch.Tensor", fill_value: Any) -> TensorType: # type: ignore
    """Create a tensor with the same shape as the input tensor filled
    with a scalar value.

    ATTENTION: This method receives a torch tensor regardless of the
    library used.

    Parameters
    ----------
    tensor
        The tensor to create a new tensor with the same shape.
    fill_value
        The value to fill the new tensor with.

    Returns
    -------
    TensorType
        A tensor with the same shape as the input tensor filled with the
        specified value.

    """
    ...
get_device(tensor) abstractmethod

Get the name of the tensor's device.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to get the device of.

required

Returns:

Type Description
str

The name of the tensor's device.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def get_device(self, tensor: TensorType) -> str:
    """Get the name of the tensor's device.

    Parameters
    ----------
    tensor
        The tensor to get the device of.

    Returns
    -------
    str
        The name of the tensor's device.

    """
    ...
shape(tensor) abstractmethod

Get the shape of the tensor.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to get the shape of.

required

Returns:

Type Description
list[int]

The shape of the tensor. The list contains as many elements as there are dimensions in the tensor.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def shape(self, tensor: TensorType) -> list[int]:
    """Get the shape of the tensor.

    Parameters
    ----------
    tensor
        The tensor to get the shape of.

    Returns
    -------
    list[int]
        The shape of the tensor. The list contains as many elements as
        there are dimensions in the tensor.

    """
    ...
squeeze(tensor) abstractmethod

Remove a dimension from the tensor at axis 0.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to remove a dimension from.

required

Returns:

Type Description
TensorType

The tensor with one less dimension.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def squeeze(self, tensor: TensorType) -> TensorType:
    """Remove a dimension from the tensor at axis 0.

    Parameters
    ----------
    tensor
        The tensor to remove a dimension from.

    Returns
    -------
    TensorType
        The tensor with one less dimension.

    """
    ...
to_device(tensor, device) abstractmethod

Move the tensor to a specified device.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to move to a specified device.

required
device str

The name of the device to move the tensor to.

required

Returns:

Type Description
TensorType

The tensor moved to the specified device.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def to_device(self, tensor: TensorType, device: str) -> TensorType:
    """Move the tensor to a specified device.

    Parameters
    ----------
    tensor
        The tensor to move to a specified device.
    device
        The name of the device to move the tensor to.

    Returns
    -------
    TensorType
        The tensor moved to the specified device.

    """
    ...
to_list(tensor) abstractmethod

Convert the tensor to a list.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to convert to a list.

required

Returns:

Type Description
list

The tensor as a list.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def to_list(self, tensor: TensorType) -> list:
    """Convert the tensor to a list.

    Parameters
    ----------
    tensor
        The tensor to convert to a list.

    Returns
    -------
    list
        The tensor as a list.

    """
    ...
to_scalar(tensor) abstractmethod

Return the only element of the tensor.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to return the only element of.

required

Returns:

Type Description
Any

The only element of the tensor.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def to_scalar(self, tensor: TensorType) -> Any:
    """Return the only element of the tensor.

    Parameters
    ----------
    tensor
        The tensor to return the only element of.

    Returns
    -------
    Any
        The only element of the tensor.

    """
    ...
unsqueeze(tensor) abstractmethod

Add a dimension to the tensor at axis 0.

Parameters:

Name Type Description Default
tensor TensorType

The tensor to add a dimension to.

required

Returns:

Type Description
TensorType

The tensor with an additional dimension.

Source code in outlines/processors/tensor_adapters/base.py
@abstractmethod
def unsqueeze(self, tensor: TensorType) -> TensorType:
    """Add a dimension to the tensor at axis 0.

    Parameters
    ----------
    tensor
        The tensor to add a dimension to.

    Returns
    -------
    TensorType
        The tensor with an additional dimension.

    """
    ...

jax

Tensor adapter for the jax library.

mlx

Tensor adapter for the mlx library.

numpy

Tensor adapter for the numpy library.

tensorflow

Tensor adapter for the tensorflow library.

torch

Tensor adapter for the torch library.