Compress mask store by only storing viable terminal sequences

Author: shubhamugare

syncode/dfa_mask_store.py

@@ -1,5 +1,6 @@
 from collections import defaultdict
 import copy, os, pickle
+import time
 import interegular
 import torch
 import regex
@@ -155,7 +156,6 @@ def incomplete_case_lookup(self, dfa_state: DFAState) -> Any:
             if dfa_state in self._exact_lookup:
                 return self._exact_lookup[dfa_state]
             else:
-                print(f"Warning: Exact lookup not found for {dfa_state} in the DFA mask store. This could be an error.", flush=True) 
                 return self._overapprox_lookup[dfa_state]
         raise ValueError(f"Invalid mode: {self._mode}")
 
@@ -167,9 +167,6 @@ def store_overapprox_lookup(self, dfa_state: DFAState, mask: torch.Tensor):
 
     def complete_case_lookup(self, dfa_state: DFAState) -> Any:
         assert isinstance(dfa_state, DFAState)
-        if dfa_state not in self._exact_lookup:
-            # FIXME: This is bit strange and need to be checked more carefully
-            return self._overapprox_lookup[dfa_state]
         return self._exact_lookup[dfa_state]
 
     def add_exact_lookup(self, dfa_state: DFAState, token):
@@ -294,7 +291,8 @@ def __init__(self,
                  special_token_ids: Iterable=[], 
                  indentation: bool=True,
                  mode='grammar_strict', # 'grammar_strict' or 'grammar_mask'
-                 ignore_terminals: Iterable[str]=[]
+                 ignore_terminals: Iterable[str]=[],
+                 parse_table=None
                  ):
         self._vocab = vocab
         self.special_token_ids = special_token_ids  
@@ -308,7 +306,11 @@ def __init__(self,
         # Iterate through each pair of DFA state and next terminals and store the overapproximate tokens
         self._lookup_table = LookupTable(vocab, special_token_ids, indentation=indentation, mode=mode)
         terminal_names = [terminal.name for terminal in terminals]
-        self._store_overapproximate_tokens(terminal_names, vocab)
+
+        followings_terminas_map = None
+        if parse_table is not None:
+            followings_terminas_map = self._compute_following_terminals_map(terminal_names, parse_table)
+        self._store_overapproximate_tokens(terminal_names, vocab, followings_terminas_map)
 
         self.indentation = indentation       
 
@@ -325,7 +327,14 @@ def set_ignore_whitespace(self, terminals: Iterable[TerminalDef], ignore_termina
         return ignore_whitespace
 
     @staticmethod
-    def load_dfa_mask_store(grammar: Grammar, tokenizer, use_cache=True, logger=common.EmptyLogger(), mode='grammar_strict'):
+    def load_dfa_mask_store(
+        grammar: Grammar, 
+        tokenizer, 
+        use_cache=True, 
+        logger=None, 
+        mode='grammar_strict',
+        parse_table=None
+    ):
         '''
         Loads the dfa for the given language and tokenizer. If the dfa is not cached, it is created and cached. 
         '''
@@ -351,7 +360,17 @@ def load_dfa_mask_store(grammar: Grammar, tokenizer, use_cache=True, logger=comm
         simplifications = grammar.simplifications()
         os.makedirs(dfa_dir, exist_ok=True)
 
-        mask_store = DFAMaskStore(base_parser.terminals, vocab, simplifications=simplifications, special_token_ids=[tokenizer.eos_token_id], mode=mode, ignore_terminals=base_parser.ignore_tokens)
+        start_time = time.time()
+        mask_store = DFAMaskStore(
+            base_parser.terminals, 
+            vocab, 
+            simplifications=simplifications, 
+            special_token_ids=[tokenizer.eos_token_id], 
+            mode=mode, 
+            ignore_terminals=base_parser.ignore_tokens, 
+            parse_table=parse_table
+            )
+        print(f"Time taken to create DFA mask store: {time.time() - start_time} seconds", flush=True)
 
         pickle.dump(mask_store, open(dfa_path, 'wb'))
         return mask_store
@@ -360,12 +379,42 @@ def _get_default_mask(self) -> torch.Tensor:
         mask = torch.zeros(len(self._vocab), dtype=torch.bool)
         return mask
 
-    def _store_overapproximate_tokens(self, terminals: Iterable[str], vocab: Iterable[str]):
+    def _compute_following_terminals_map(self, terminals: Iterable[str], parse_table) -> defaultdict:
+        """
+        From terminals, filter out terminals that cannot follow the current terminal
+        according to the grammar.
+
+        If in the parsing table Action[cur_terminal, parser_state] = 'shift, new_parser_state' then next terminals
+        are the terminals that are legal in new_parser_state.
+        """
+        following_terminals_map = defaultdict(set)
+        terminals_set = set(terminals)
+
+        # We iterate through each cur_terminal:    
+        for cur_terminal in terminals:
+            # We iterate through each parser_state:
+            for _, row in parse_table.states.items():
+                if cur_terminal in row:
+                    action = row[cur_terminal]
+                    # -> If we see a shift action to new_parser_state
+                    if str(action[0]) == 'Shift':
+                        new_parser_state = action[1]
+                        for next_terminal in parse_table.states[new_parser_state]:
+                            # Lark parse_table stores non-terminals and terminals together
+                            if next_terminal in terminals_set:
+                                # -> -> we add the terminals that are legal in new_parser_state
+                                following_terminals_map[cur_terminal].add(next_terminal)
+        
+        return following_terminals_map
+
+
+    def _store_overapproximate_tokens(self, terminals: Iterable[str], vocab: Iterable[str], followings_terminas_map: dict=None):
         """
         Stores the overapproximate tokens for each dfa state and next terminals
         """
         all_dfa_states = self._dfas.states()
         pbar = tqdm(total=len(all_dfa_states))
+
         for dfa_state in all_dfa_states:
             for token_idx, token in enumerate(vocab):
                 is_special_token = token_idx in self.special_token_ids
@@ -375,12 +424,18 @@ def _store_overapproximate_tokens(self, terminals: Iterable[str], vocab: Iterabl
                         self._lookup_table.dfa_state_and_next_terminal_to_tokens_add(
                             dfa_state, '$END', token_idx)
                 else:
-                    self._process_regular_tokens(terminals, dfa_state, token_idx, token)
+                    if followings_terminas_map is not None and dfa_state.terminal in followings_terminas_map:
+                        following_terminals = followings_terminas_map[dfa_state.terminal]
+                    else:
+                        following_terminals = terminals
+
+                    self._process_regular_tokens(following_terminals, dfa_state, token_idx, token)
 
             pbar.update(1)
 
     def _process_regular_tokens(self, terminals, dfa_state, token_idx, token):
         remainder = token.replace('\t', '    ')
+
         is_valid, remainder = self._dfas.consume_prefix(dfa_state, remainder)
         if is_valid:
             if remainder == '':
@@ -455,12 +510,10 @@ def _lookup_next_tokens(self, dfa_states: Iterable[DFAState], r: ParseResult) ->
                             elif len(accept_sequence) == 2:
                                 overapprox_token_ids |= self._lookup_next_tokens_for_dfa_state(dfa_state, accept_sequence[1])
                             elif len(accept_sequence) == 3:
-                                # This is useful in under-approximating `grammar_strict` mode as they help improve the precision of SynCode
-                                if self._mode == 'grammar_strict':
-                                    # If the DFA state is a final state we can jump to the start of next terminal
-                                    if self._dfas.is_final(dfa_state): 
-                                        ignore_init_state = self._dfas.initial(accept_sequence[1])
-                                        overapprox_token_ids |= self._lookup_next_tokens_for_dfa_state(ignore_init_state, accept_sequence[2])
+                                # If the DFA state is a final state we can jump to the start of next terminal
+                                if self._dfas.is_final(dfa_state): 
+                                    ignore_init_state = self._dfas.initial(accept_sequence[1])
+                                    overapprox_token_ids |= self._lookup_next_tokens_for_dfa_state(ignore_init_state, accept_sequence[2])
                             else:
                                 raise ValueError(f"Invalid accept sequence: {accept_sequence}")
         return overapprox_token_ids

syncode/grammar_decoder.py

@@ -49,18 +49,19 @@ def __init__(self,
         # Ignore whitespace tokens
         self._ignore_whitespace = self._get_ignore_whitespace(self.grammar)
 
+        # Create parsers
+        self.inc_parser: Iterator[IncrementalParser] = create_parser(self.grammar, logger=self.logger, parser=parser, ignore_whitespace=self._ignore_whitespace)
+
         # Load dfa mask store
         self.dfa_mask_store = DFAMaskStore.load_dfa_mask_store(
                                     grammar=self.grammar, 
                                     tokenizer=self.tokenizer, 
                                     use_cache=use_cache, 
                                     logger=self.logger,
                                     mode=mode,
+                                    parse_table=self.inc_parser.base_parser.parser.parser._parse_table
                                     )
 
-        # Create parsers
-        self.inc_parser: Iterator[IncrementalParser] = create_parser(self.grammar, logger=self.logger, parser=parser, ignore_whitespace=self._ignore_whitespace)
-
 
     def _log_current_status(self, partial_code, r: ParseResult):
         self.logger.log_code('Partial code', partial_code)