diff --git a/src/benchmark_lfe.py b/src/benchmark_lfe.py index bded073..1138395 100644 --- 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 index 60fb01d..5c3b183 100644 --- 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)