enable generate.fsm with llamacpp by using outlines.processors

outlines/fsm/guide.py

@@ -184,13 +184,7 @@ def __init__(self, regex_string: str, tokenizer: "Tokenizer"):
         ) = create_states_mapping(regex_string, tokenizer)
         self.eos_token_id = tokenizer.eos_token_id
         self.final_states = fsm_finals | {-1}
-
-        # cache returned masks token masks
-        # this increases performance of the mask substantially
-        self.states_to_token_mask = {
-            state: torch.tensor(list(next_tokens_to_end_states.keys()))
-            for state, next_tokens_to_end_states in self.states_to_token_maps.items()
-        }
+        self._cache_state_to_token_tensor()
 
     def get_next_instruction(self, state: int) -> Instruction:
         """Return the next instruction for guided generation.
@@ -285,8 +279,19 @@ def create_states_mapping_from_interegular_fsm(
             from_interegular_instance.empty_token_ids,
         ) = create_states_mapping_from_interegular_fsm(interegular_fsm)
         from_interegular_instance.eos_token_id = tokenizer.eos_token_id
+        from_interegular_instance._cache_state_to_token_tensor()
         return from_interegular_instance
 
+    def _cache_state_to_token_tensor(self):
+        """
+        cache state -> token int tensor
+        this increases performance of mask construction substantially
+        """
+        self.states_to_token_mask = {
+            state: torch.tensor(list(next_tokens_to_end_states.keys()))
+            for state, next_tokens_to_end_states in self.states_to_token_maps.items()
+        }
+
     def is_final_state(self, state: int) -> bool:
         """Determine whether the current state of the guide is a final state."""
         return state in self.final_states

outlines/generate/fsm.py

@@ -1,14 +1,31 @@
+from functools import singledispatch
+
 import interegular
 
 from outlines.fsm.guide import RegexGuide
-from outlines.generate.api import SequenceGenerator
+from outlines.generate.api import SequenceGenerator, SequenceGeneratorAdapter
+from outlines.models import MLXLM, LlamaCpp, Transformers
 from outlines.samplers import Sampler, multinomial
 
 
+@singledispatch
 def fsm(
     model, fsm: interegular.fsm.FSM, sampler: Sampler = multinomial()
 ) -> SequenceGenerator:
     fsm = RegexGuide.from_interegular_fsm(fsm, model.tokenizer)
     device = model.device
     generator = SequenceGenerator(fsm, model, sampler, device)
     return generator
+
+
+@fsm.register(MLXLM)
+@fsm.register(Transformers)
+@fsm.register(LlamaCpp)
+def fsm_unified(
+    model, fsm: interegular.fsm.FSM, sampler: Sampler = multinomial()
+) -> SequenceGeneratorAdapter:
+    from outlines.processors import FSMLogitsProcessor
+
+    fsm = RegexGuide.from_interegular_fsm(fsm, model.tokenizer)
+    logits_processor = FSMLogitsProcessor(tokenizer=model.tokenizer, fsm=fsm)
+    return SequenceGeneratorAdapter(model, logits_processor, sampler)

outlines/generate/regex.py

@@ -41,6 +41,7 @@ def regex(model, regex_str: str, sampler: Sampler = multinomial()):
 
 @regex.register(MLXLM)
 @regex.register(Transformers)
+@regex.register(LlamaCpp)
 def regex_unified(
     model,
     regex_str: str,
@@ -52,18 +53,6 @@ def regex_unified(
     return SequenceGeneratorAdapter(model, logits_processor, sampler)
 
 
-@regex.register(LlamaCpp)
-def regex_llamacpp(
-    model: LlamaCpp,
-    regex_str: str,
-    sampler: Sampler = multinomial(),
-):
-    from outlines.integrations.llamacpp import RegexLogitsProcessor
-
-    logits_processor = RegexLogitsProcessor(regex_str, llm=model.model)
-    return SequenceGeneratorAdapter(model, logits_processor, sampler)
-
-
 @regex.register(VLLM)
 def regex_vllm(
     model: VLLM,

outlines/generate/text.py

@@ -38,6 +38,7 @@ def text(model, sampler: Sampler = multinomial()) -> SequenceGenerator:
 
 @text.register(MLXLM)
 @text.register(Transformers)
+@text.register(LlamaCpp)
 def text_unified(model, sampler: Sampler = multinomial()):
     return SequenceGeneratorAdapter(model, None, sampler)
 
@@ -47,11 +48,6 @@ def text_vllm(model: VLLM, sampler: Sampler = multinomial()):
     return SequenceGeneratorAdapter(model, None, sampler)
 
 
-@text.register(LlamaCpp)
-def text_llamacpp(model: LlamaCpp, sampler: Sampler = multinomial()):
-    return SequenceGeneratorAdapter(model, None, sampler)
-
-
 @text.register(OpenAI)
 def text_openai(model: OpenAI, sampler: Sampler = multinomial()) -> OpenAI:
     if not isinstance(sampler, multinomial):

outlines/models/llamacpp.py

@@ -141,6 +141,10 @@ class LlamaCpp:
     def __init__(self, model: "Llama"):
         self.model = model
 
+    @property
+    def tokenizer(self):
+        return LlamaCppTokenizer(self.model)
+
     def prepare_generation_parameters(
         self,
         generation_parameters: GenerationParameters,

outlines/models/transformers.py

@@ -182,10 +182,6 @@ def forward(
 
         return output.logits, output.past_key_values
 
-    @property
-    def device(self):
-        return self.model.device
-
     def __call__(
         self,
         input_ids: "torch.LongTensor",

tests/generate/test_generate.py

@@ -101,6 +101,17 @@ def test_generate_text_stream(request, model_fixture):
             assert isinstance(token, str)
 
 
+@pytest.mark.parametrize("pattern", REGEX_PATTERNS)
+@pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES)
+def test_generate_fsm(request, model_fixture, pattern):
+    import interegular
+
+    model = request.getfixturevalue(model_fixture)
+    generator = generate.fsm(model, interegular.parse_pattern(pattern).to_fsm())
+    res = generator("test")
+    assert re.fullmatch(pattern, res) is not None, res
+
+
 @pytest.mark.parametrize("pattern", REGEX_PATTERNS)
 @pytest.mark.parametrize("model_fixture", ALL_MODEL_FIXTURES)
 def test_generate_regex(request, model_fixture, pattern):