Merge branch 'introduce-asv-ci-workflow-883'

lapp0 · May 26, 2024 · 592c884 · 592c884
2 parents c699d23 + 359b385
commit 592c884
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 3 deletions.
diff --git a/benchmarks/bench_json_schema.py b/benchmarks/bench_json_schema.py
@@ -5,7 +5,11 @@
 from outlines.fsm.guide import RegexGuide  # noqa: E402
 from outlines.fsm.json_schema import build_regex_from_schema  # noqa: E402
 
-from .common import ensure_numba_compiled, setup_tokenizer  # noqa: E402
+from .common import (  # noqa: E402
+    clear_outlines_cache,
+    ensure_numba_compiled,
+    setup_tokenizer,
+)
 
 simple_schema = """{
         "$defs": {
@@ -70,6 +74,7 @@ class JsonSchemaBenchmark:
     params = schemas.keys()
 
     def setup(self, schema_name):
+        clear_outlines_cache()
         self.tokenizer = setup_tokenizer()
         self.schema = schemas[schema_name]
         ensure_numba_compiled(self.tokenizer)

diff --git a/benchmarks/bench_numba_compile.py b/benchmarks/bench_numba_compile.py
@@ -5,13 +5,14 @@
 
 import outlines
 
-from .common import setup_tokenizer
+from .common import clear_outlines_cache, setup_tokenizer
 
 outlines.disable_cache()
 
 
 class NumbaCompileBenchmark:
     def setup(self):
+        clear_outlines_cache()
         from outlines.fsm import regex
 
         self.tokenizer = setup_tokenizer()

diff --git a/benchmarks/bench_regex_guide.py b/benchmarks/bench_regex_guide.py
@@ -1,6 +1,6 @@
 import outlines
 
-from .common import ensure_numba_compiled, setup_tokenizer
+from .common import clear_outlines_cache, ensure_numba_compiled, setup_tokenizer
 
 outlines.disable_cache()
 
@@ -23,6 +23,7 @@ class RegexGuideBenchmark:
     params = regex_samples.keys()
 
     def setup(self, pattern_name):
+        clear_outlines_cache()
         self.tokenizer = setup_tokenizer()
         ensure_numba_compiled(self.tokenizer)
         self.pattern = regex_samples[pattern_name]
@@ -35,6 +36,7 @@ class MemoryRegexGuideBenchmark:
     params = ["simple_phone", "complex_span_constrained_relation_extraction"]
 
     def setup(self, pattern_name):
+        clear_outlines_cache()
         self.tokenizer = setup_tokenizer()
         ensure_numba_compiled(self.tokenizer)
         self.pattern = regex_samples[pattern_name]

diff --git a/benchmarks/common.py b/benchmarks/common.py
@@ -1,9 +1,14 @@
 from transformers import AutoTokenizer
 
+import outlines.caching
 from outlines.fsm.guide import RegexGuide
 from outlines.models.transformers import TransformerTokenizer
 
 
+def clear_outlines_cache():
+    outlines.caching.clear_cache()
+
+
 def setup_tokenizer():
     tokenizer = AutoTokenizer.from_pretrained("gpt2")
     return TransformerTokenizer(tokenizer)