cleanup lm-format-enforcer cache in teardown step

dottxt-ai · Oct 16, 2024 · 08f10af · 08f10af
1 parent 5bc6379
commit 08f10af
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/benchmark_lfe.py b/src/benchmark_lfe.py
@@ -37,6 +37,9 @@ def time_lfe(self, _, regex_name):
             for i in range(len(regex_sample_tokens)):
                 _ = token_enforcer.get_allowed_tokens(regex_sample_tokens[: i + 1])
 
+    def teardown(self, *args):
+        del self.tokenizer_data
+
 
 class LMFormatEnforcerJsonSchema:
     params = [models, json_cases.keys()]
@@ -66,3 +69,6 @@ def time_lfe(self, _, json_schema_name):
             json_sample_tokens = self.tokenizer.encode(json_sample)
             for i in range(len(json_sample_tokens)):
                 _ = token_enforcer.get_allowed_tokens(json_sample_tokens[: i + 1])
+
+    def teardown(self, *args):
+        del self.tokenizer_data