-
Notifications
You must be signed in to change notification settings - Fork 399
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Improve
outlines.processors
, add integration tests to test_generate…
….py (#998) A lot of these fixes were intended for #966 however that's blocked until there's a new `transformers` release. These improvements are general to all models and will enable PRs resolving #806 and #965 # Structure of `OutlinesLogitsProcessor` The goal is to create a base class which allows a logits processors to be implemented once and used for any `outlines.models` inference library. To accomplish this we must normalize the input array. It must have a consistent type (`torch.Tensor`) and consistent dimensionality (2). We can normalize both of these simply, and without any copy operations. `mlx.core.array`, `numpy.array`, and `torch.Tensor` all support [pythons array standard `__dlpack__`](https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__dlpack__.html). This standard allows for casting between array types without copying. `torch.Tensor` is the only input type which cannot always be cast to any other type because torch tensors may live in GPU memory. Therefore, we cast all arrays to `torch.Tensor`, implement logits processors using torch methods, and convert back to the original array type in `OutlinesLogitsProcessor`. See docstring of `OutlinesLogitsProcessor.__call__()` for more details. # Detailed Changes - Rename `BaseLogitsProcessor` to `OutlinesLogitsProcessor` - Ensure `OutlinesLogitsProcessor.process_logits()` is always passed a 2D batch request with `torch.Tensor` logits and `List` input_ids. Also clean up code to be more readable in `OutlinesLogitsProcessor__call__()` - Ensure `FSMLogitsProcessor` allows unstable sequence ordering (beam search in transformers and vLLM change the order of sequences) - Update `tests/generate/test_generate.py` to cover more permutations of - regex / text - batch / single - greedy / multinomial / beam search - `stream()` / `generate()` - Ensure performance stability with difference array libraries through `benchmark_processors.py`
- Loading branch information
Showing
7 changed files
with
332 additions
and
83 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,52 @@ | ||
import mlx.core as mx | ||
import numpy as np | ||
import torch | ||
|
||
from outlines.processors import OutlinesLogitsProcessor | ||
|
||
|
||
def is_mlx_lm_allowed(): | ||
try: | ||
import mlx.core as mx | ||
except ImportError: | ||
return False | ||
return mx.metal.is_available() | ||
|
||
|
||
class HalvingLogitsProcessor(OutlinesLogitsProcessor): | ||
"""Simply halve the passed logits""" | ||
|
||
def process_logits(self, input_ids, logits): | ||
return logits / 2 | ||
|
||
|
||
class LogitsProcessorBenchmark: | ||
params = ["torch", "numpy"] | ||
if mx.metal.is_available(): | ||
params += ["mlx"] | ||
|
||
def setup(self, array_library): | ||
self.logits_processor = HalvingLogitsProcessor() | ||
|
||
# logits: (4, 30,000 ) dtype=float | ||
# input_ids shape: (4, 2048) dtype=int | ||
if array_library == "torch": | ||
self.logits = torch.rand((4, 30000), dtype=torch.float) | ||
self.input_ids = torch.randint( | ||
low=0, high=30000, size=(4, 2048), dtype=torch.int | ||
) | ||
elif array_library == "numpy": | ||
self.logits = np.random.rand(4, 30000).astype(np.float32) | ||
self.input_ids = np.random.randint(low=0, high=30000, size=(4, 2048)) | ||
elif array_library == "mlx": | ||
self.logits = mx.random.uniform( | ||
low=-1e9, high=1e9, shape=(4, 30000), dtype=mx.float32 | ||
) | ||
self.input_ids = mx.random.randint( | ||
low=0, high=30000, shape=(4, 2048), dtype=mx.int32 | ||
) | ||
else: | ||
raise ValueError | ||
|
||
def time_logits_processor(self, array_library): | ||
self.logits_processor(self.input_ids, self.logits) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,7 +1,7 @@ | ||
from .structured import ( | ||
BaseLogitsProcessor, | ||
CFGLogitsProcessor, | ||
FSMLogitsProcessor, | ||
JSONLogitsProcessor, | ||
OutlinesLogitsProcessor, | ||
RegexLogitsProcessor, | ||
) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.