Logging and Sampling outlines.processors

docs/reference/logits_processors.md

@@ -0,0 +1,51 @@
+# Logits Processors
+
+TODO: Explanation of what logits processors do
+
+TODO: List of `OutlinesLogitsProcessor`
+
+TODO: Example of how to implement TemperatureLogitsProcessor
+
+TODO: using logits processors with models directly vs using with outlines
+
+TODO: Explanation of pipelines
+
+TODO: Link to log logits
+
+## Using Logits Processors in Outlines
+
+TODO Explanation
+
+```
+import outlines
+```
+
+
+## Chaining Logits Processors
+
+```
+import outlines
+import outlines.processors as processors
+
+model = outlines.models.llamacpp(
+    repo_id="M4-ai/TinyMistral-248M-v2-Instruct-GGUF",
+    filename="TinyMistral-248M-v2-Instruct.Q4_K_M.gguf"
+)
+
+# Create a chained logits processor
+logits_processor = (
+    processors.sequence_logging(model.tokenizer) |  # Log the generated sequence
+	processors.logits_logging(model.tokenizer) |  # Log the raw logits
+	processors.regex(r"[0-9]*", model.tokenizer) |  # Restrict the logits to match the pattern
+	processors.temperature(0.5) |  # Set temperature to 0.5
+	processors.logits_logging(model.tokenizer)  # Log the restricted, temperature-augmentent, sampled logits
+)
+
+generator = outlines.generate.base(model, logits_process)
+generator("What is your favorite number? ")
+```
+
+Output:
+```
+TODO
+```

outlines/generate/base.py

@@ -0,0 +1,22 @@
+from functools import singledispatch
+
+from outlines.generate.api import SequenceGeneratorAdapter
+from outlines.models import MLXLM, OpenAI
+from outlines.processors import OutlinesLogitsProcessor
+from outlines.samplers import Sampler, multinomial
+
+
+@singledispatch
+def base(
+    model, logits_processor: OutlinesLogitsProcessor, sampler: Sampler = multinomial()
+):
+    return SequenceGeneratorAdapter(model, logits_processor, sampler)
+
+
+@base.register(OpenAI)
+def base_openai(
+    model: OpenAI,
+    logits_processor: OutlinesLogitsProcessor,
+    sampler: Sampler = multinomial(),
+) -> Exception:
+    raise NotImplementedError("The OpenAI API does not support logits processing.")

outlines/processors/__init__.py

@@ -1,7 +1,34 @@
+from .base_logits_processor import ChainedLogitsProcessor, OutlinesLogitsProcessor
+from .logging import LogitsLoggingLogitsProcessor, SequenceLoggingLogitsProcessor
+from .sampling import (
+    FrequencyPenaltyLogitsProcessor,
+    MinPLogitsProcessor,
+    NoRepeatNGramLogitsProcessor,
+    PresencePenaltyLogitsProcessor,
+    QuadraticSmoothingLogitsProcessor,
+    RepetitionPenaltyLogitsProcessor,
+    TemperatureLogitsProcessor,
+    TFSLogitsProcessor,
+    TopKLogitsProcessor,
+    TopPLogitsProcessor,
+)
 from .structured import (
     CFGLogitsProcessor,
     GuideLogitsProcessor,
     JSONLogitsProcessor,
-    OutlinesLogitsProcessor,
     RegexLogitsProcessor,
 )
+
+# aliases for convenience
+chained = ChainedLogitsProcessor
+
+cfg = CFGLogitsProcessor
+guide = GuideLogitsProcessor
+json = JSONLogitsProcessor
+regex = RegexLogitsProcessor
+
+temperature = TemperatureLogitsProcessor
+min_p = MinPLogitsProcessor
+
+sequence_logging = SequenceLoggingLogitsProcessor
+logits_logging = LogitsLoggingLogitsProcessor

outlines/processors/base_logits_processor.py

@@ -1,5 +1,6 @@
+import reprlib
 from abc import abstractmethod
-from typing import TYPE_CHECKING, List, Protocol, Type, Union
+from typing import TYPE_CHECKING, List, Protocol, Type, Union, runtime_checkable
 
 import numpy as np
 import torch
@@ -20,6 +21,7 @@ def is_mlx_array_type(array_type):
     return issubclass(array_type, mx.array)
 
 
+@runtime_checkable
 class OutlinesLogitsProcessor(Protocol):
     """
     Base class for logits processors which normalizes types of logits:
@@ -133,3 +135,35 @@ def _from_torch(tensor: torch.Tensor, target_type: Type) -> Array:
             raise TypeError(
                 f"Failed to convert torch tensors to target_type `{target_type}`"
             )
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}({reprlib.repr(self.__dict__)})"
+
+    def __or__(self, other):
+        if not isinstance(other, OutlinesLogitsProcessor):
+            raise ValueError(
+                "Can only chain with another OutlinesLogitsProcessor instance."
+            )
+        return ChainedLogitsProcessor([self, other])
+
+
+class ChainedLogitsProcessor(OutlinesLogitsProcessor):
+    """Handle chaining two logits processors to process logits sequentially"""
+
+    processors: List[OutlinesLogitsProcessor]
+
+    def __init__(self, processors: List[OutlinesLogitsProcessor]):
+        self.processors = []
+        for processor in processors:
+            if isinstance(processor, ChainedLogitsProcessor):
+                self.processors.extend(processor.processors)
+            else:
+                self.processors.append(processor)
+
+    def process_logits(
+        self, input_ids: List[List[int]], logits: torch.Tensor
+    ) -> torch.Tensor:
+        result = logits
+        for processor in self.processors:
+            result = processor.process_logits(input_ids, result)
+        return result

outlines/processors/logging.py

@@ -0,0 +1,64 @@
+import logging
+import sys
+import warnings
+from typing import List
+
+import torch
+
+from .base_logits_processor import OutlinesLogitsProcessor
+
+
+class LogitsLoggingLogitsProcessor(OutlinesLogitsProcessor):
+    """Handle chaining two logits processors to process logits sequentially"""
+
+    def __init__(self, tokenizer, top_n=8, logger=None, warn=True):
+        self.tokenizer = tokenizer
+        self.top_n = top_n
+        if logger is not None:
+            self.logger = logger
+        else:
+            self.logger = logging.getLogger("logits_logger")
+            self.logger.setLevel(logging.info)
+            self.logger.addHandler(logging.StreamHandler(sys.stderr))
+        if warn:
+            warnings.warn(
+                "Do not use LoggingLogitsProcessor in production, it slows down generation."
+            )
+
+    def process_logits(
+        self, input_ids: List[List[int]], logits: torch.Tensor
+    ) -> torch.Tensor:
+        # all token probs for the current batch
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        # top candidate tokens probs
+        top_indices = torch.topk(probs, self.top_n).indices
+        top_indices = [
+            set(row.tolist()) | {self.tokenizer.eos_token_id} for row in top_indices
+        ]
+        batch_top_probs = [
+            {token_idx: probs[batch_num, token_idx] for token_idx in row_indices}
+            for batch_num, row_indices in enumerate(top_indices)
+        ]
+        self.logger.info(batch_top_probs)
+        return logits
+
+
+class SequenceLoggingLogitsProcessor(OutlinesLogitsProcessor):
+    def __init__(self, tokenizer, top_n=8, logger=None, warn=True):
+        self.tokenizer = tokenizer
+        if logger is not None:
+            self.logger = logger
+        else:
+            self.logger = logging.getLogger("sequence_logger")
+            self.logger.setLevel(logging.info)
+            self.logger.addHandler(logging.StreamHandler(sys.stderr))
+        if warn:
+            warnings.warn(
+                "Do not use SequenceLoggingProcessor in production, it slows down generation."
+            )
+
+    def process_logits(
+        self, input_ids: List[List[int]], logits: torch.Tensor
+    ) -> torch.Tensor:
+        self.logger.info(self.tokenizer.decode(input_ids))
+        return logits