Fix JSON inference example

Author: turboderp

examples/inference_json.py

@@ -5,7 +5,7 @@
 from exllamav2 import ExLlamaV2, ExLlamaV2Config, ExLlamaV2Cache, ExLlamaV2Tokenizer
 from exllamav2.generator import ExLlamaV2DynamicGenerator
 from exllamav2.generator.filters import ExLlamaV2PrefixFilter
-from lmformatenforcer.integrations.exllamav2 import ExLlamaV2TokenEnforcerFilter
+from inference_json_lmfe_wrapper import ExLlamaV2TokenEnforcerFilter
 from lmformatenforcer import JsonSchemaParser
 from pydantic import BaseModel, conlist
 from typing import Literal
@@ -61,7 +61,7 @@ class Superhero(BaseModel):
     filters.append(None)
     prompts.append(p)
     filters.append([
-        ExLlamaV2TokenEnforcerFilter(schema_parser, tokenizer),
+        ExLlamaV2TokenEnforcerFilter(model, tokenizer, schema_parser),
         ExLlamaV2PrefixFilter(model, tokenizer, ["{", " {"])
     ])
 

examples/inference_json_lmfe_wrapper.py

@@ -0,0 +1,47 @@
+
+import sys, os
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
+from exllamav2 import ExLlamaV2, ExLlamaV2Tokenizer
+from exllamav2.generator.filters import ExLlamaV2Filter
+from functools import lru_cache
+from lmformatenforcer.integrations.exllamav2 import build_token_enforcer_tokenizer_data
+from lmformatenforcer import TokenEnforcer, CharacterLevelParser
+from typing import List
+
+
+# Temporary wrapper for lm-format-enforcer, until the integration in LMFE itself is updated
+
+
+@lru_cache(10)
+def _get_lmfe_tokenizer_data(tokenizer: ExLlamaV2Tokenizer):
+    return build_token_enforcer_tokenizer_data(tokenizer)
+
+
+class ExLlamaV2TokenEnforcerFilter(ExLlamaV2Filter):
+
+    token_sequence: List[int]
+
+    def __init__(
+        self,
+        model: ExLlamaV2,
+        tokenizer: ExLlamaV2Tokenizer,
+        character_level_parser: CharacterLevelParser,
+    ):
+        super().__init__(model, tokenizer)
+        tokenizer_data = _get_lmfe_tokenizer_data(tokenizer)
+        self.token_enforcer = TokenEnforcer(tokenizer_data, character_level_parser)
+        self.token_sequence = []
+
+    def begin(self, prefix_str: str) -> None:
+        self.token_sequence = []
+
+    def feed(self, token) -> None:
+        self.token_sequence.append(int(token[0][0]))
+
+    def next(self):
+        allowed_tokens = self.token_enforcer.get_allowed_tokens(self.token_sequence)
+        return sorted(allowed_tokens), []
+
+    def use_background_worker(self):
+        return True