From 15b0121a7da9a4816c4667e54808703a67ab897f Mon Sep 17 00:00:00 2001 From: Ari Weinstein Date: Fri, 19 Apr 2024 12:46:21 -0400 Subject: [PATCH] Fix slowdown after generating backslash escape sequences --- lmformatenforcer/jsonschemaparser.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/lmformatenforcer/jsonschemaparser.py b/lmformatenforcer/jsonschemaparser.py index 7df02e4..84f482a 100644 --- a/lmformatenforcer/jsonschemaparser.py +++ b/lmformatenforcer/jsonschemaparser.py @@ -119,8 +119,13 @@ def can_end(self) -> bool: return all(parser.can_end() for parser in self.object_stack) def shortcut_key(self) -> Optional[Hashable]: - if self.object_stack: - current_parser = self.object_stack[-1] + current_parsers = [] + for parser in reversed(self.object_stack): + current_parsers.append(parser) + if not parser.can_end(): + break + + for current_parser in current_parsers: if isinstance(current_parser, StringParsingState): if not current_parser.allowed_strings and current_parser.seen_opening_quote and not current_parser.seen_closing_quote and not current_parser.regex_parser: # Performance optimization: When we are parsing a string that is not from a list of allowed strings, most tokens