correct redundant test runs

dottxt-ai · Oct 17, 2024 · a4bf704 · a4bf704
1 parent a5adbe4
commit a4bf704
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 4 deletions.
diff --git a/src/benchmark_lfe.py b/src/benchmark_lfe.py
@@ -54,9 +54,13 @@ def time_lfe_first_token(self, _, regex_name):
         self._get_first_token(enforcer)
 
 
-class LMFormatEnforcerRegexRunTime(LMFormatEnforcerRegex):
+class LMFormatEnforcerRegexRunTime(LMFormatEnforcerBenchmark):
     """Class which warms-up enforcer in setup steps"""
 
+    params = [models, regex_cases.keys()]
+    param_names = ["model", "regex_name"]
+    timeout = 1200
+
     def setup(self, model, regex_name):
         samples = regex_cases[regex_name]["samples"]
         self.do_setup(model, samples)
@@ -92,9 +96,13 @@ def time_lfe_first_token(self, _, json_schema_name):
         self._get_first_token(enforcer)
 
 
-class LMFormatEnforcerJsonSchemaRunTime(LMFormatEnforcerJsonSchema):
+class LMFormatEnforcerJsonSchemaRunTime(LMFormatEnforcerBenchmark):
     """Class which warms-up enforcer in setup steps"""
 
+    params = [models, json_cases.keys()]
+    param_names = ["model", "json_schema_name"]
+    timeout = 600
+
     def setup(self, model, json_schema_name):
         samples = json_cases[json_schema_name]["samples"]
         self.do_setup(model, samples)

diff --git a/src/benchmark_outlines.py b/src/benchmark_outlines.py
@@ -71,9 +71,13 @@ def time_outlines_first_token(self, _, regex_name):
         self._get_first_token(guide)
 
 
-class OutlinesRegexRunTime(OutlinesRegex):
+class OutlinesRegexRunTime(OutlinesBenchmark):
     """Class which warms-up Guide in setup steps"""
 
+    params = [models, regex_cases.keys()]
+    param_names = ["model", "regex_name"]
+    timeout = 1200
+
     def setup(self, model, regex_name):
         samples = regex_cases[regex_name]["samples"]
         self.do_setup(model, samples)
@@ -111,9 +115,15 @@ def time_outlines_first_token(self, _, json_schema_name):
         self._get_first_token(guide)
 
 
-class OutlinesJsonSchemaRunTime(OutlinesJsonSchema):
+class OutlinesJsonSchemaRunTime(OutlinesBenchmark):
     """Class which warms-up Guide in setup steps"""
 
+    json_from_regex_fn = lambda self, schema: build_regex_from_schema(schema)
+
+    params = [models, json_cases.keys()]
+    param_names = ["model", "json_schema_name"]
+    timeout = 1200
+
     def setup(self, model, json_schema_name):
         samples = json_cases[json_schema_name]["samples"]
         self.do_setup(model, samples)