Create TextLogitsProcessor class

Author: RobinPicard

outlines/generate/text.py

@@ -40,12 +40,18 @@ def text(model, sampler: Sampler = multinomial()) -> SequenceGenerator:
 @text.register(Transformers)
 @text.register(LlamaCpp)
 def text_unified(model, sampler: Sampler = multinomial()):
-    return SequenceGeneratorAdapter(model, None, sampler)
+    from outlines.processors import TextLogitsProcessor
+
+    logits_processor = TextLogitsProcessor(model.tokenizer)
+    return SequenceGeneratorAdapter(model, logits_processor, sampler)
 
 
 @text.register(VLLM)
 def text_vllm(model: VLLM, sampler: Sampler = multinomial()):
-    return SequenceGeneratorAdapter(model, None, sampler)
+    from outlines.integrations.vllm import TextLogitsProcessor
+
+    logits_processor = TextLogitsProcessor(model)
+    return SequenceGeneratorAdapter(model, logits_processor, sampler)
 
 
 @text.register(OpenAI)

outlines/integrations/llamacpp.py

@@ -33,7 +33,7 @@
 from numpy.typing import NDArray
 from pydantic import BaseModel
 
-from outlines.fsm.guide import CFGGuide, Guide, RegexGuide
+from outlines.fsm.guide import CFGGuide, Guide, RegexGuide, StopAtEOSGuide
 from outlines.fsm.json_schema import build_regex_from_schema
 from outlines.integrations.utils import convert_json_schema_to_str
 from outlines.models.llamacpp import LlamaCppTokenizer
@@ -104,6 +104,30 @@ def copy(self) -> "LogitsProcessor":
         return LogitsProcessor(tokenizer=self.tokenizer, fsm=self.fsm.copy())
 
 
+class TextLogitsProcessor(LogitsProcessor):
+    """Bias vLLM generation for free text (required because of prompt alignment).
+
+    Attributes
+    ----------
+    tokenizer
+        The tokenizer used to convert tokens to ids.
+    fsm
+        The finite state machine which is used to bias the logits.
+    """
+
+    def __init__(self, llm: "Llama"):
+        """Compile the FSM that drives the regex-guided generation.
+
+        Parameters
+        ----------
+        llm
+            The Llama model.
+        """
+        tokenizer = LlamaCppTokenizer(model=llm)
+        fsm = StopAtEOSGuide(tokenizer)
+        super().__init__(tokenizer=tokenizer, fsm=fsm)
+
+
 class RegexLogitsProcessor(LogitsProcessor):
     """Bias LlamaCpp generation based on a regular expression.
 

outlines/integrations/vllm.py

@@ -32,39 +32,41 @@
 import torch
 from pydantic import BaseModel
 
-from outlines.fsm.guide import RegexGuide
+from outlines.fsm.guide import Guide, RegexGuide, StopAtEOSGuide
 from outlines.fsm.json_schema import build_regex_from_schema
 from outlines.integrations.utils import adapt_tokenizer, convert_json_schema_to_str
 
 if TYPE_CHECKING:
     from vllm import LLM
 
+    from outlines.models.tokenizer import Tokenizer
 
-class RegexLogitsProcessor:
-    """Bias vLLM generation based on a regular expression.
+
+class FSMLogitsProcessor:
+    """Bias vLLM generation based on a FSM.
 
     Attributes
     ----------
     fsm
         The finite state machine which is used to bias the logits.
     """
 
-    def __init__(self, regex_string: str, llm: "LLM"):
+    def __init__(self, fsm: Guide):
         """Compile the FSM that drives the regex-structured generation.
 
         Parameters
         ----------
-        regex_string
-            A string that represents a regular expression.
-        llm
-            The vLLM model.
+        fsm
+            Guide.
 
-        Raises
-        ------
-        ValueError
-            If the provided LLM instance in `RegexLogitsProcessor` neither has a
-            `tokenizer` attribute or a `get_tokenizer` method.
         """
+        self.fsm = fsm
+        self.mask_cache: Dict[int, torch.Tensor] = {}
+        self._fsm_state: DefaultDict[int, int] = defaultdict(int)
+
+    @staticmethod
+    def get_llm_tokenizer(llm: "LLM") -> "Tokenizer":
+        """Give the tokenizer attached to the LLM provided"""
         if hasattr(llm, "get_tokenizer"):
             tokenizer = llm.get_tokenizer()
         elif hasattr(llm, "tokenizer"):
@@ -74,13 +76,10 @@ def __init__(self, regex_string: str, llm: "LLM"):
                 tokenizer = llm.tokenizer
         else:
             raise ValueError(
-                "The provided LLM instance in `RegexLogitsProcessor` neither has a "
+                "The provided LLM instance in `FSMLogitsProcessor` neither has a "
                 "`tokenizer` attribute or a `get_tokenizer` method."
             )
-        tokenizer = adapt_tokenizer(tokenizer=tokenizer)
-        self.mask_cache: Dict[int, torch.Tensor] = {}
-        self.fsm = RegexGuide(regex_string, tokenizer)
-        self._fsm_state: DefaultDict[int, int] = defaultdict(int)
+        return adapt_tokenizer(tokenizer=tokenizer)
 
     def __call__(self, input_ids: List[int], scores: torch.Tensor) -> torch.Tensor:
         """Use the FSM to bias the logits before sampling the next token.
@@ -125,6 +124,64 @@ def __call__(self, input_ids: List[int], scores: torch.Tensor) -> torch.Tensor:
         return biased_scores
 
 
+class TextLogitsProcessor(FSMLogitsProcessor):
+    """Bias vLLM generation for free text (required because of prompt alignment).
+
+    Attributes
+    ----------
+    fsm
+        The finite state machine which is used to bias the logits.
+    """
+
+    def __init__(self, llm: "LLM"):
+        """Compile the FSM that drives the regex-structured generation.
+
+        Parameters
+        ----------
+        llm
+            The vLLM model.
+
+        Raises
+        ------
+        ValueError
+            If the provided LLM instance in `TextLogitsProcessor` neither has a
+            `tokenizer` attribute or a `get_tokenizer` method.
+        """
+        tokenizer = self.get_llm_tokenizer(llm)
+        fsm = StopAtEOSGuide(tokenizer)
+        super().__init__(fsm=fsm)
+
+
+class RegexLogitsProcessor(FSMLogitsProcessor):
+    """Bias vLLM generation based on a regular expression.
+
+    Attributes
+    ----------
+    fsm
+        The finite state machine which is used to bias the logits.
+    """
+
+    def __init__(self, regex_string: str, llm: "LLM"):
+        """Compile the FSM that drives the regex-structured generation.
+
+        Parameters
+        ----------
+        regex_string
+            A string that represents a regular expression.
+        llm
+            The vLLM model.
+
+        Raises
+        ------
+        ValueError
+            If the provided LLM instance in `RegexLogitsProcessor` neither has a
+            `tokenizer` attribute or a `get_tokenizer` method.
+        """
+        tokenizer = self.get_llm_tokenizer(llm)
+        fsm = RegexGuide(regex_string, tokenizer)
+        super().__init__(fsm=fsm)
+
+
 class JSONLogitsProcessor(RegexLogitsProcessor):
     """Bias vLLM generation based on a JSON schema.
 

outlines/processors/__init__.py

@@ -4,4 +4,5 @@
     JSONLogitsProcessor,
     OutlinesLogitsProcessor,
     RegexLogitsProcessor,
+    TextLogitsProcessor,
 )

outlines/processors/structured.py

@@ -29,7 +29,7 @@
 import torch
 from pydantic import BaseModel
 
-from outlines.fsm.guide import CFGGuide, Guide, RegexGuide
+from outlines.fsm.guide import CFGGuide, Guide, RegexGuide, StopAtEOSGuide
 from outlines.fsm.json_schema import build_regex_from_schema
 from outlines.integrations.utils import convert_json_schema_to_str
 
@@ -115,6 +115,29 @@ def copy(self) -> "FSMLogitsProcessor":
         return FSMLogitsProcessor(tokenizer=self.tokenizer, fsm=self.fsm.copy())
 
 
+class TextLogitsProcessor(FSMLogitsProcessor):
+    """Bias generation for free text (required because of prompt alignment).
+
+    Attributes
+    ----------
+    tokenizer
+        The tokenizer used to convert tokens to ids.
+    fsm
+        The finite state machine which is used to bias the logits.
+    """
+
+    def __init__(self, tokenizer: "Tokenizer"):
+        """Compile the FSM that drives the regex-guided generation.
+
+        Parameters
+        ----------
+        tokenizer
+            An Outlines tokenizer.
+        """
+        fsm = StopAtEOSGuide(tokenizer)
+        super().__init__(tokenizer=tokenizer, fsm=fsm)
+
+
 class RegexLogitsProcessor(FSMLogitsProcessor):
     """Bias generation based on a regular expression.