Cache creation of states_token_mapping in the __init__ of RegexFSM

outlines/caching.py

@@ -1,20 +1,74 @@
+import asyncio
+import hashlib
 import os
 from typing import Callable, Optional
 
-from perscache import Cache, NoCache
-from perscache.serializers import JSONSerializer
-from perscache.storage import LocalFileStorage
+import cloudpickle
+from diskcache import Cache
 
 home_dir = os.path.expanduser("~")
 cache_dir = os.environ.get("OUTLINES_CACHE_DIR", f"{home_dir}/.cache/outlines")
-memory = Cache(serializer=JSONSerializer(), storage=LocalFileStorage(cache_dir))
-
-
-def cache(ignore: Optional[str] = None):
-    def cache_fn(fn: Callable):
-        return memory.cache(ignore=ignore)(fn)
+memory = Cache(cache_dir, eviction_policy="none", cull_limit=0)
+_caching_enabled = True
+
+
+def hash_arguments(*args, **kwargs) -> str:
+    """Create a hash out of the args and kwargs provided"""
+    result = hashlib.md5()
+    for item in list(args) + sorted(kwargs.items()):
+        result.update(cloudpickle.dumps(item))
+    return result.hexdigest()
+
+
+def cache(key_function: Optional[Callable] = None):
+    """Caching decorator for memoizing function calls.
+    The cache key is created based on the values returned by the key_function callable
+    if provided or based on the arguments of the decorated function directly otherwise
+    Parameters
+    ----------
+    key_function
+      A callable function used to generate a unique key for each function call. It's
+      called with the arguments of the decorated function as arguments
+    Returns
+    -------
+      A decorator function that can be applied to other functions.
+    """
 
-    return cache_fn
+    def decorator(cached_function: Callable):
+        def wrapper(*args, **kwargs):
+            if not _caching_enabled:
+                return cached_function(*args, **kwargs)
+            if key_function:
+                key_args = key_function(*args, **kwargs)
+                cache_key = hash_arguments(*key_args)
+            else:
+                cache_key = hash_arguments(*args, **kwargs)
+            if cache_key in memory:
+                return memory[cache_key]
+            result = cached_function(*args, **kwargs)
+            memory[cache_key] = result
+            return result
+
+        async def async_wrapper(*args, **kwargs):
+            if not _caching_enabled:
+                return await cached_function(*args, **kwargs)
+            if key_function:
+                key_args = key_function(*args, **kwargs)
+                cache_key = hash_arguments(*key_args)
+            else:
+                cache_key = hash_arguments(*args, **kwargs)
+            if cache_key in memory:
+                return memory[cache_key]
+            result = await cached_function(*args, **kwargs)
+            memory[cache_key] = result
+            return result
+
+        if asyncio.iscoroutinefunction(cached_function):
+            return async_wrapper
+        else:
+            return wrapper
+
+    return decorator
 
 
 def get_cache():
@@ -51,11 +105,11 @@ def disable_cache():
     >>> cache.disable()
 
     """
-    global memory
-    memory = NoCache()
+    global _caching_enabled
+    _caching_enabled = False
 
 
 def clear_cache():
     """Erase the cache completely."""
     global memory
-    memory.storage.clear()
+    memory.clear()

outlines/fsm/fsm.py

@@ -1,9 +1,10 @@
-from typing import TYPE_CHECKING, List, NewType, Protocol
+from typing import TYPE_CHECKING, List, NewType, Protocol, Tuple
 
 import interegular
 from lark import Lark
 
 # from outlines.fsm.parsing import PartialLark
+from outlines.caching import cache
 from outlines.fsm.regex import create_fsm_index_tokenizer, make_deterministic_fsm
 
 if TYPE_CHECKING:
@@ -97,28 +98,49 @@ def copy(self) -> "StopAtTokenFSM":
 class RegexFSM(FSM):
     """FSM to generate text that is in the language of a regular expression."""
 
-    def __init__(self, regex_string: str, tokenizer: "Tokenizer"):
-        regex_pattern = interegular.parse_pattern(regex_string)
-        regex_fsm, _ = make_deterministic_fsm(regex_pattern.to_fsm().reduce())
+    def __init__(
+        self,
+        regex_string: str,
+        tokenizer: "Tokenizer",
+    ):
+        @cache()
+        def create_states_mapping(
+            regex_string: str, cacheable_vocabulary: Tuple[Tuple[str, int]]
+        ) -> Tuple[dict, set, set]:
+            """Create the variables related to the mapping between states and tokens
+            The parameters of the function are used for caching purpose
+            """
+            regex_pattern = interegular.parse_pattern(regex_string)
+            regex_fsm, _ = make_deterministic_fsm(regex_pattern.to_fsm().reduce())
+            (
+                states_to_token_maps,
+                empty_token_ids,
+            ) = create_fsm_index_tokenizer(regex_fsm, tokenizer)
+
+            # We make sure that it is possible to generate strings in the language
+            # of the regular expression with the tokens present in the model's
+            # vocabulary.
+            if not any(
+                regex_fsm.finals.intersection(v.values())
+                for v in states_to_token_maps.values()
+            ):
+                raise ValueError(
+                    "The vocabulary does not allow us to build a sequence that matches the input regex"
+                )
+
+            final_states = regex_fsm.finals | {
+                -1
+            }  # Include the EOS token in final states
+            return states_to_token_maps, empty_token_ids, final_states
+
         (
             self.states_to_token_maps,
             self.empty_token_ids,
-        ) = create_fsm_index_tokenizer(regex_fsm, tokenizer)
-
-        # We make sure that it is possible to generate strings in the language
-        # of the regular expression with the tokens present in the model's
-        # vocabulary.
-        if not any(
-            regex_fsm.finals.intersection(v.values())
-            for v in self.states_to_token_maps.values()
-        ):
-            raise ValueError(
-                "The vocabulary does not allow us to build a sequence that matches the input regex"
-            )
-
-        self.final_states = regex_fsm.finals | {
-            -1
-        }  # Include the EOS token in final states
+            self.final_states,
+        ) = create_states_mapping(
+            regex_string, tuple(sorted(tokenizer.vocabulary.items()))
+        )
+        self.num_tokens_generated = 0
         self.vocabulary = tokenizer.vocabulary.values()
         self.end_token_id = tokenizer.eos_token_id
 

pyproject.toml

@@ -28,7 +28,8 @@ dependencies = [
    "lark",
    "nest_asyncio",
    "numpy",
-   "perscache",
+   "cloudpickle",
+   "diskcache",
    "pydantic>=2.0",
    "scipy",
    "torch>=2.1",
@@ -105,7 +106,8 @@ module = [
     "mamba_ssm.*",
     "nest_asyncio",
     "numpy.*",
-    "perscache.*",
+    "cloudpickle.*",
+    "diskcache.*",
     "pydantic.*",
     "pytest",
     "referencing.*",

tests/fsm/test_fsm.py

@@ -22,6 +22,7 @@ def test_regex_vocabulary_error():
     class MockTokenizer:
         vocabulary = {"a": 1}
         special_tokens = {"eos"}
+        eos_token_id = 3
 
         def convert_token_to_string(self, token):
             return token

tests/test_cache.py

@@ -1,8 +1,7 @@
 import os
 import tempfile
-from pathlib import Path
 
-import perscache
+import diskcache
 import pytest
 
 
@@ -35,19 +34,18 @@ def test_cache(refresh_environment):
         import outlines
 
         memory = outlines.get_cache()
-        assert memory.storage.location == Path(tempdir)
+        assert memory.directory == tempdir
 
         yield outlines.caching.cache()
 
-        memory.storage.clear()
+        memory.clear()
 
 
 def test_get_cache(test_cache):
     import outlines
 
     memory = outlines.get_cache()
-    assert isinstance(memory, perscache.Cache)
-    assert isinstance(memory.storage, perscache.storage.LocalFileStorage)
+    assert isinstance(memory, diskcache.Cache)
 
     # If the cache is enabled then the size
     # of `store` should not increase the