Add library specific handling of tensors in logits processors #1498

RobinPicard · 2025-03-15T18:04:41Z

This PR addresses issue #1445

The PR creates classes of type TensorHandler for various tensor libraries. Such classes are used by the logit processors to avoid having to manipulate the input_its/logits tensors itself. This solution allows us not to require the user to download all tensor libraries supported by outlines and to treat the tensor in their native type without having to turn them into torch tensors. We also modify the local models to possess an attribute tensor_library_name to indicate to the logits processor what TensorHandler implementation to use.

Thanks to those change, we remove the mandatory dependencies on torch and numpy.

outlines/models/base.py

outlines/processors/tensor_handlers/__init__.py

rlouf · 2025-03-17T14:36:14Z

I was originally skeptical of the tensor handler the approach but it's growing on me. I would use the name tensor_adapter though.

Thanks to those change, we remove the mandatory dependencies on torch and numpy.

🎉

RobinPicard · 2025-03-18T08:36:33Z

I changed the name from tensor_handler to tensor_adapter @rlouf

rlouf · 2025-03-22T17:51:22Z

outlines/models/llamacpp.py


-    def __init__(self, model: "Llama"):
+    def __init__(self, model: "Llama", tensor_library_name: Optional[str] = None):


Isn't tensor_library_name for llama-cpp-python necessarily numpy?

rlouf · 2025-03-22T17:51:40Z

outlines/models/mlxlm.py


    def __init__(
        self,
        model: "nn.Module",
        tokenizer: "PreTrainedTokenizer",
+        tensor_library_name: Optional[str] = None,


Isn't it necessarily mlx?

rlouf · 2025-03-22T17:53:24Z

outlines/models/transformers.py

@@ -163,11 +163,13 @@ def format_output_type(self, output_type):

 class Transformers(Model):


We may be able to do this automatically. It looks like all models implemented implemented with JAX inherit from FlaxPretrainedModel

rlouf · 2025-03-22T17:55:39Z

outlines/models/vllm.py

@@ -139,8 +144,8 @@ def generate_stream(self, model_input, output_type, **inference_kwargs):
        )


-def from_vllm(model: "LLM") -> VLLM:
-    return VLLM(model)
+def from_vllm(model: "LLM", tensor_library_name: Optional[str] = None) -> VLLM:


Isn't it necessarily torch? I don't think vLLM supports any other backend.

rlouf · 2025-03-22T17:56:20Z

pyproject.toml

@@ -40,7 +39,6 @@ dependencies = [
   "typing_extensions",
   "iso3166",
   "airportsdata",
-   "torch",


rlouf

Looks good overall. The main comment I have is that I don't think that any library other than transformers supports more than one backend, and even for transformers there may be a way to determine the backend automatically.

RobinPicard · 2025-03-23T22:24:24Z

Yes, you're right. I completely removed the possibility of specifying the tensor library to use when initiating a model as there's no case in which it's necessary. As a result I could simplify the logic the tensor_library_name attribute (just a simple attribute without distinct default value now).

rlouf · 2025-03-23T22:39:48Z

outlines/models/transformers.py

@@ -185,11 +185,18 @@ def __init__(
            the `transformers` API for tokenizers.

        """
+        from transformers import FlaxPreTrainedModel


It looks like this requires flax to be installed :/ Let's see if we can do this another way.

Can't we just add flax to the test dependencies? We already have jax there

It will still fail whenever users try to initialise a transformers model. We could use try/except for the checks (also add TFPretrainedModel).

Ah yes, right

I've added a TensorAdapter for tensorflow

rlouf · 2025-03-25T06:27:36Z

Do we have tests that initialize a Transformers instance for each of the backend transformers supports? If not that would be why coverage is failing.

RobinPicard · 2025-03-25T08:23:10Z

Yes. The issue was that there are lines in the Transformers model that are executed only if some libraries are not available. I added a coverage exclusion rule for them.

rlouf reviewed Mar 15, 2025

View reviewed changes

outlines/models/base.py Outdated Show resolved Hide resolved

rlouf reviewed Mar 15, 2025

View reviewed changes

outlines/processors/tensor_handlers/__init__.py Outdated Show resolved Hide resolved

RobinPicard force-pushed the remove_mandatory_dependency_torch branch 2 times, most recently from 6304811 to 6217d08 Compare March 16, 2025 14:36

RobinPicard marked this pull request as ready for review March 16, 2025 15:02

RobinPicard requested a review from rlouf March 16, 2025 15:03

RobinPicard self-assigned this Mar 16, 2025

RobinPicard added this to the 1.0 milestone Mar 16, 2025

RobinPicard linked an issue Mar 16, 2025 that may be closed by this pull request

Create library-specific versions of logits processors #1445

Closed

RobinPicard force-pushed the remove_mandatory_dependency_torch branch from 6217d08 to 8762761 Compare March 18, 2025 08:17

rlouf force-pushed the remove_mandatory_dependency_torch branch from 8762761 to 733eeed Compare March 18, 2025 13:19

RobinPicard force-pushed the remove_mandatory_dependency_torch branch 5 times, most recently from 71dc231 to 6b250b4 Compare March 19, 2025 20:56

rlouf reviewed Mar 22, 2025

View reviewed changes

pyproject.toml

@@ -40,7 +39,6 @@ dependencies = [

"typing_extensions",

"iso3166",

"airportsdata",

"torch",

Copy link

Member

rlouf Mar 22, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🥳

rlouf requested changes Mar 22, 2025

View reviewed changes

RobinPicard force-pushed the remove_mandatory_dependency_torch branch from 6b250b4 to 407eddf Compare March 23, 2025 22:22

RobinPicard force-pushed the remove_mandatory_dependency_torch branch from 407eddf to 2aaa146 Compare March 23, 2025 22:34

rlouf reviewed Mar 23, 2025

View reviewed changes

RobinPicard force-pushed the remove_mandatory_dependency_torch branch from 2aaa146 to 1d00c9a Compare March 23, 2025 22:46

Create the TensorAdapter classes

bba69d7

Modify the logits processors to use the tensor adapters

a324020

RobinPicard force-pushed the remove_mandatory_dependency_torch branch 2 times, most recently from 3318a42 to 0951993 Compare March 24, 2025 09:45

Modify models and generators to implement the use of tensor adapters

a9208a9

RobinPicard force-pushed the remove_mandatory_dependency_torch branch 2 times, most recently from 60f96bb to ae46cff Compare March 24, 2025 14:30

RobinPicard requested a review from rlouf March 24, 2025 14:51

rlouf approved these changes Mar 25, 2025

View reviewed changes

Add coverage exclusion for mlx, vllm and serve

ed9b572

RobinPicard force-pushed the remove_mandatory_dependency_torch branch from ae46cff to ed9b572 Compare March 25, 2025 08:07

RobinPicard merged commit 246dd80 into dottxt-ai:v1.0 Mar 25, 2025
5 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add library specific handling of tensors in logits processors #1498

Add library specific handling of tensors in logits processors #1498

RobinPicard commented Mar 15, 2025 •

edited

Loading

rlouf commented Mar 17, 2025

RobinPicard commented Mar 18, 2025

rlouf Mar 22, 2025

rlouf Mar 22, 2025

rlouf Mar 22, 2025

rlouf Mar 22, 2025

rlouf Mar 22, 2025

rlouf left a comment

RobinPicard commented Mar 23, 2025

rlouf Mar 23, 2025

RobinPicard Mar 23, 2025 •

edited

Loading

rlouf Mar 23, 2025

RobinPicard Mar 24, 2025

RobinPicard Mar 24, 2025

rlouf commented Mar 25, 2025

RobinPicard commented Mar 25, 2025 •

edited

Loading


		def __init__(self, model: "Llama"):
		def __init__(self, model: "Llama", tensor_library_name: Optional[str] = None):

		@@ -163,11 +163,13 @@ def format_output_type(self, output_type):

		class Transformers(Model):

Add library specific handling of tensors in logits processors #1498

Add library specific handling of tensors in logits processors #1498

Conversation

RobinPicard commented Mar 15, 2025 • edited Loading

rlouf commented Mar 17, 2025

RobinPicard commented Mar 18, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rlouf left a comment

Choose a reason for hiding this comment

RobinPicard commented Mar 23, 2025

Choose a reason for hiding this comment

RobinPicard Mar 23, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rlouf commented Mar 25, 2025

RobinPicard commented Mar 25, 2025 • edited Loading

RobinPicard commented Mar 15, 2025 •

edited

Loading

RobinPicard Mar 23, 2025 •

edited

Loading

RobinPicard commented Mar 25, 2025 •

edited

Loading