Improve outlines.processors, add integration tests to test_generate.py

benchmarks/bench_processors.py

@@ -0,0 +1,47 @@
+import mlx.core as mx
+import numpy as np
+import torch
+
+from outlines.processors import OutlinesLogitsProcessor
+
+
+class HalvingLogitsProcessor(OutlinesLogitsProcessor):
+    """Simply halve the passed logits"""
+
+    def process_logits(input_ids, logits):
+        return logits / 2
+
+
+# logits: (4, 30,000 ) dtype=float
+# input_ids shape: (4, 2048) dtype=int
+logit_processor_inputs = {
+    "torch": {
+        "logits": torch.rand((4, 30000), dtype=torch.float),
+        "input_ids": torch.randint(low=0, high=30000, size=(4, 2048), dtype=torch.int),
+    },
+    "numpy": {
+        "logits": np.random.rand(4, 30000).astype(np.float32),
+        "input_ids": np.random.randint(low=0, high=30000, size=(4, 2048)),
+    },
+    "mlx": {
+        "logits": mx.random.uniform(
+            low=-1e9, high=1e9, shape=(4, 30000), dtype=mx.float32
+        ),
+        "input_ids": mx.random.randint(
+            low=0, high=30000, shape=(4, 2048), dtype=mx.int32
+        ),
+    },
+}
+
+
+class LogitsProcessorBenchmark:
+    params = logit_processor_inputs.keys()
+
+    def setup(self, array_library):
+        self.logits_processor = HalvingLogitsProcessor()
+        lp_inputs = self.inputs[array_library]
+        self.logits = lp_inputs["logits"]
+        self.input_ids = lp_inputs["input_ids"]
+
+    def time_logits_processor(self, array_library):
+        self.logits_processor(self.logits, self.input_ids)

outlines/__init__.py

@@ -2,6 +2,7 @@
 import outlines.generate
 import outlines.grammars
 import outlines.models
+import outlines.processors
 import outlines.types
 from outlines.base import vectorize
 from outlines.caching import clear_cache, disable_cache, get_cache

outlines/models/mlxlm.py

@@ -9,7 +9,7 @@
     from transformers import PreTrainedTokenizer
 
     from outlines.generate.api import GenerationParameters, SamplingParameters
-    from outlines.processors import BaseLogitsProcessor
+    from outlines.processors import OutlinesLogitsProcessor
 
 
 class MLXLM:
@@ -120,7 +120,7 @@ def generate_step(
         temp: Optional[float],
         top_p: Optional[float],
         sampler: str,
-        logits_processor: "BaseLogitsProcessor",
+        logits_processor: "OutlinesLogitsProcessor",
     ) -> Generator[Tuple[int, float], None, None]:
         """
         Adapted from
@@ -135,7 +135,7 @@ def generate_step(
                 top_p (float, optional): Nulceus sampling, higher means model considers
                   more less likely words.
                 sampler (str): The sampler string defined by SequenceGeneratorAdapter
-                logits_processor (BaseLogitsProcessor): Augment logits before sampling.
+                logits_processor (OutlinesLogitsProcessor): Augment logits before sampling.
         """
         import mlx.core as mx
         import mlx_lm

outlines/processors/__init__.py

@@ -1,7 +1,7 @@
 from .structured import (
-    BaseLogitsProcessor,
     CFGLogitsProcessor,
     FSMLogitsProcessor,
     JSONLogitsProcessor,
+    OutlinesLogitsProcessor,
     RegexLogitsProcessor,
 )

outlines/processors/base_logits_processor.py

@@ -1,5 +1,5 @@
 from abc import abstractmethod
-from typing import List, Protocol, Union
+from typing import Any, List, Protocol, Union
 
 import numpy as np
 import torch
@@ -14,10 +14,11 @@ def is_mlx_array(logits):
     return isinstance(logits, mx.array)
 
 
-class BaseLogitsProcessor(Protocol):
+class OutlinesLogitsProcessor(Protocol):
     """
     Base class for logits processors which normalizes types of logits:
     - ndarray (used by llama-cpp-python), converted to torch.Tensor
+    - mlx.core.array (used by mlx-lm), converted to torch.Tensor
     - torch.Tensor (used by everything else)
 
     Normalization of types and conversion to torch.Tensor
@@ -29,50 +30,101 @@ class BaseLogitsProcessor(Protocol):
 
     @abstractmethod
     def process_logits(
-        self, input_ids: List[int], logits: torch.Tensor
+        self, input_ids: List[List[int]], logits: torch.Tensor
     ) -> torch.Tensor:
-        ...
+        """
+        input_ids and logits are always 2D tensors for handling a batch of sequences.
+
+        - input_ids -> List[List[tokens]]
+        - logits.shape[0] -> 2D_Tensor[logits]
+
+        Important 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
+        """
+        pass
 
+    @torch.no_grad()
     def __call__(
         self,
         input_ids: Union[NDArray[np.int64], List[int], torch.Tensor],
         logits: Union[NDArray[np.float32], torch.Tensor],
     ) -> Union[NDArray[np.int64], torch.Tensor]:
         """
         Apply logits processor
-        Unify type
-        - convert input_ids: either ndarray, List[int], or Tensor -> List[int]
-        - convert logits: either ndarray, mlx array, Tensor -> Tensor
-        Call process_logits() to perform business logic
+
+        1) Unify type
+        - convert input_ids: either ndarray, mlx array, List[int], or Tensor -> List[List[int]]
+        - convert logits: either ndarray, mlx array, or Tensor -> 2D float Tensor
+        2) Unify shape, ensure logits and input_ids are 2D
+        3) Call self.process_logits() to perform business logic
+        4) Cast logits back to original array library type
         """
-        with torch.no_grad():
-            if not isinstance(input_ids, list):
-                input_ids = input_ids.tolist()
-
-            if isinstance(logits, np.ndarray):
-                # Unify type, convert numpy array to Tensor
-                # from_numpy and .numpy() don't copy the data, it uses the same memory address
-                torch_logits = torch.from_numpy(logits)
-                processed_torch_logits = self.process_logits(input_ids, torch_logits)
-                return processed_torch_logits.detach().numpy()
-
-            elif isinstance(logits, torch.Tensor):
-                return self.process_logits(input_ids, logits)
-
-            elif is_mlx_array(logits):
-                # mlx -> torch -> mlx conversion docs:
-                # https://ml-explore.github.io/mlx/build/html/usage/numpy.html
-                import mlx.core as mx
-
-                torch_logits = torch.from_dlpack(logits)
-                processed_torch_logits = self.process_logits(input_ids, torch_logits)
-
-                # numpy doesn't support bfloat16, mlx doesn't support direct conversion from torch
-                logits_float32_numpy = processed_torch_logits.float().numpy()
-                return mx.array(logits_float32_numpy)
-
-            else:
-                raise TypeError(
-                    "LogitsProcessor must be called with either np.NDArray"
-                    ", torch.Tensor, or mlx.core.array typed logits"
-                )
+
+        # ensure logits are torch Tensors
+        torch_logits = self._to_torch(logits)
+
+        assert torch_logits.shape[:-1] == self._to_torch(input_ids).shape[:-1]
+
+        # ensure input_ids are List
+        if not isinstance(input_ids, list):
+            input_ids = input_ids.tolist()  # compatible with numpy, torch, and mlx
+
+        # Guarantee passed as 2D Tensors, then covert back to original (1D or 2D) shape
+        if len(torch_logits.shape) == 2:
+            processed_logits = self.process_logits(input_ids, logits)
+        elif len(torch_logits.shape) == 1:
+            processed_logits = self.process_logits(
+                [input_ids], torch_logits.unsqueeze(0)
+            ).squeeze(0)
+
+        # return logits as passed array type
+        return self._from_torch(processed_logits, type(logits))
+
+    @staticmethod
+    def _to_torch(tensor_like: Any) -> torch.Tensor:
+        """Convert various types to torch.Tensor."""
+        if isinstance(tensor_like, torch.Tensor):
+            return tensor_like
+
+        elif isinstance(tensor_like, np.ndarray):
+            return torch.from_numpy(tensor_like)
+
+        elif isinstance(tensor_like, list):
+            return torch.tensor(tensor_like)
+
+        elif is_mlx_array(tensor_like):
+            # mlx -> torch -> mlx conversion docs:
+            # https://ml-explore.github.io/mlx/build/html/usage/numpy.html
+            return torch.from_dlpack(tensor_like)
+
+        else:
+            raise TypeError(
+                "LogitsProcessor must be called with either np.NDArray"
+                ", torch.Tensor, list, or mlx.core.array typed logits"
+            )
+
+    @staticmethod
+    def _from_torch(tensor: torch.Tensor, target_type: Any) -> Any:
+        """Convert torch.Tensor to the specified target type."""
+        if target_type == torch.Tensor:
+            return tensor
+
+        elif target_type == np.ndarray:
+            return tensor.detach().numpy()
+
+        elif target_type == list:
+            return tensor.detach().tolist()
+
+        elif target_type == "mlx_array":
+            import mlx.core as mx
+
+            # numpy doesn't support bfloat16, mlx doesn't support direct conversion from torch
+            return mx.array(tensor.float().numpy())
+
+        else:
+            raise RuntimeError(
+                "Failed to convert torch tensors back to original dtype. {tensor}"
+                f"tensor={tensor}, target_type={target_type}"
+            )

outlines/processors/structured.py

@@ -24,7 +24,7 @@
 limitations under the License.
 """
 import math
-from typing import TYPE_CHECKING, List, Optional, Type, Union
+from typing import TYPE_CHECKING, Dict, List, Optional, Type, Union
 
 import numpy as np
 import torch
@@ -35,13 +35,13 @@
 from outlines.fsm.json_schema import build_regex_from_schema
 from outlines.integrations.utils import convert_json_schema_to_str
 
-from .base_logits_processor import BaseLogitsProcessor
+from .base_logits_processor import OutlinesLogitsProcessor
 
 if TYPE_CHECKING:
     from outlines.models.tokenizer import Tokenizer
 
 
-class FSMLogitsProcessor(BaseLogitsProcessor):
+class FSMLogitsProcessor(OutlinesLogitsProcessor):
     """Bias generation using a finite state machine.
 
     Attributes
@@ -63,12 +63,13 @@ def __init__(self, tokenizer: "Tokenizer", fsm: Guide):
             The finite state machine which is used to bias the logits.
         """
         self.tokenizer = tokenizer
-        self._fsm_state = 0
+        self._fsm_states: Dict[int, int] = {}
         self.fsm: Guide = fsm
         self._is_first_token = True
+        self._seq_start_idx: Optional[int] = None
 
     def process_logits(
-        self, input_ids: List[int], logits: torch.Tensor
+        self, input_ids: List[List[int]], logits: torch.Tensor
     ) -> NDArray[np.float32]:
         """Use the FSM to bias the logits before sampling the next token.
 
@@ -84,17 +85,31 @@ def process_logits(
         torch.Tensor
             The biased logits.
         """
+        sequence_states: List[int] = []  # vector of states corresponding to `input_ids`
+
         if self._is_first_token:
             self._is_first_token = False
+            self._seq_start_idx = len(input_ids[0])
+
+            self._fsm_states = {hash(tuple([])): 0}
+            sequence_states = [0] * len(input_ids)
+
         else:
-            last_token = input_ids[-1]
-            self._fsm_state = self.fsm.get_next_state(self._fsm_state, last_token)
+            for seq_ids in input_ids:
+                prev_state_key = hash(tuple(seq_ids[self._seq_start_idx : -1]))
+                prev_state = self._fsm_states[prev_state_key]
 
-        allowed_tokens = self.fsm.get_next_instruction(self._fsm_state).tokens
-        allowed_tokens = torch.tensor(allowed_tokens, device=logits.device)
+                curr_state_key = hash(tuple(seq_ids[self._seq_start_idx :]))
+                curr_state = self.fsm.get_next_state(prev_state, seq_ids[-1])
+
+                self._fsm_states[curr_state_key] = curr_state
+                sequence_states.append(curr_state)
 
         mask = torch.full_like(logits, -math.inf)
-        mask[allowed_tokens] = logits[allowed_tokens]
+        for i, fsm_state in enumerate(sequence_states):
+            allowed_tokens = self.fsm.get_next_instruction(fsm_state).tokens
+            mask[i, allowed_tokens] = logits[i, allowed_tokens]
+
         return mask
 
     def copy(self) -> "FSMLogitsProcessor":