Add different logging system

Author: shubhamugare

syncode/__init__.py

@@ -1,3 +1,6 @@
 from syncode.infer import Syncode
 from grammar_decoder import SyncodeLogitsProcessor
 from parsers.grammars import Grammar
+import common
+
+common.setup_logging()

syncode/common.py

@@ -1,4 +1,6 @@
+import logging
 import os
+import sys
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
@@ -36,6 +38,46 @@ def get_output_path(model_name, grammar, dataset, num_samples, mode):
         os.makedirs(out_dir, exist_ok=True)
         return out_dir,out_path
 
+# This is the setup for Python logging
+def setup_logging(level=None):
+    """
+    Configure the root logger for both application and test usage.
+    
+    This function is safe to call multiple times - it will only configure
+    logging once to avoid duplicate handlers.
+    
+    Args:
+        level: Override the logging level. If None, uses the LOG_LEVEL 
+               environment variable or defaults to INFO.
+    
+    Returns:
+        The root logger
+    """ 
+    # Determine the logging level
+    if level is None:
+        # Get level from environment or default to INFO
+        level_name = os.environ.get('LOG_LEVEL', 'INFO')
+        level = getattr(logging, level_name.upper(), logging.INFO)
+    
+    # Get the root logger
+    root_logger = logging.getLogger()
+    
+    # Clear any existing handlers to avoid duplicates
+    for handler in root_logger.handlers[:]:
+        root_logger.removeHandler(handler)
+    
+    # Set the logging level
+    root_logger.setLevel(level)
+    
+    # Create a stdout handler
+    handler = logging.StreamHandler(sys.stdout)
+    formatter = logging.Formatter('[%(asctime)s-%(name)s] - %(message)s')
+    handler.setFormatter(formatter)
+    root_logger.addHandler(handler)
+        
+    return root_logger
+
+
 class Logger:
     """
     Logger class for logging the output of the model

syncode/grammar_decoder.py

@@ -7,6 +7,9 @@
 from syncode.parsers import create_parser, create_base_parser
 from syncode.mask_store.mask_store import MaskStore
 from syncode.parsers.grammars import Grammar
+import logging
+logger = logging.getLogger(__name__)
+
 
 # Set to True for debugging
 DEBUG = True
@@ -18,15 +21,16 @@ class SyncodeLogitsProcessor(LogitsProcessor):
     Args:
         grammar (str): The grammar to use for parsing e.g. "python".
         tokenizer (PreTrainedTokenizer): The tokenizer to use for decoding.
-        logger (common.Logger): The logger to use for logging.
         use_cache (bool, optional): Whether to use the cache. Defaults to True.
         parse_output_only (bool, optional): Whether to parse the prompt. Defaults to False.
+        num_samples (int, optional): The number of sequences to generate. Defaults to 1.
         dev_mode (bool, optional): Whether to run in development mode. Defaults to False.
+        parser (str, optional): The parser to use. Defaults to 'lalr'.
+        mode (str, optional): The mode to use. Defaults to 'grammar_mask'.
     """
     def __init__(self, 
         grammar: Grammar, 
         tokenizer: PreTrainedTokenizer, 
-        logger: common.Logger=common.EmptyLogger(), 
         use_cache=True,
         parse_output_only=True, 
         num_samples=1,
@@ -38,7 +42,6 @@ def __init__(self,
         self.byte_tokenizer = ByteTokenizer(tokenizer)
 
         self.grammar = grammar
-        self.logger = logger
         self.dev_mode = dev_mode
         self.batch_size = num_samples
         self.parse_failed = False
@@ -55,23 +58,17 @@ def __init__(self,
         self._ignore_whitespace = self._get_ignore_whitespace(self.grammar)
 
         # Create parser
-        self.inc_parser: IncrementalParser = create_parser(self.grammar, logger=self.logger, parser=parser, ignore_whitespace=self._ignore_whitespace)
+        self.inc_parser: IncrementalParser = create_parser(self.grammar, parser=parser, ignore_whitespace=self._ignore_whitespace)
 
         # Load dfa mask store
         self.dfa_mask_store = MaskStore.init_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,
                                     )
-
+        
 
-    def _log_current_status(self, partial_code, r: ParseResult):
-        self.logger.log_code('Partial code', partial_code)
-        self.logger.log(repr(r))
-
     def _get_ignore_whitespace(self, grammar):
         """
         Check if the grammar allows whitespace tokens to be ignored.
@@ -158,11 +155,7 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> to
             res, skip = self._parse_partial_code(idx, partial_code, remainder_bytes, accepted_generation=True)
             if skip: continue
 
-            accept_mask = self.dfa_mask_store.get_accept_mask(res, logger=self.logger)
-
-            if DEBUG: 
-                self._log_current_status(partial_code, res)
-                greedy_token = self.tokenizer.decode(scores[idx].argmax(dim=-1)) 
+            accept_mask = self.dfa_mask_store.get_accept_mask(res)
 
             if torch.sum(accept_mask) != 0: # If there are acceptable tokens for the current partial code 
                 if len(scores[idx]) > len(accept_mask):
@@ -172,11 +165,8 @@ def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> to
                     accept_mask = accept_mask[: len(scores[idx])]
                 scores[idx] = scores[idx].masked_fill(~accept_mask.to(scores.device), -float("inf"))
             else: # Otherwise, report the error and mask no tokens
-                self.logger.log('No acceptable tokens for the current partial code!')
-                self._log_current_status(partial_code, res)
-
-            # For debugging - remove later
-            if DEBUG: self._debug_greedy(scores, idx, partial_code, res, greedy_token)
+                logger.debug('No acceptable tokens for the current partial code!')
+                logger.debug(repr(res))
 
         return scores
 
@@ -239,28 +229,6 @@ def update_valid_state(self, partial_code: str, idx: int, r: ParseResult):
                 if accept_seq[0] == '$END' or accept_seq[0] == 'EOF':
                     self.last_valid_state[idx] = len(partial_code) - len(r.remainder)
 
-    def _debug_greedy(self, scores, idx, partial_code, r, greedy_token):
-        greedy_grammar_token = self.tokenizer.decode(scores[idx].argmax(dim=-1))
-        if greedy_token != greedy_grammar_token:
-            self._log_greedy_difference(greedy_grammar_token, partial_code, r, greedy_token)
-
-    def _log_greedy_difference(self, greedy_grammar_token, partial_code, r, greedy_token):
-        self.logger.log_check(f"Greedy token and greedy grammar-based token do not match!")
-        self.logger.log(f"Greedy token: {repr(greedy_token)}")
-        self.logger.log(f"Greedy grammar-based token: {repr(greedy_grammar_token)}")
-        self._log_current_status(partial_code, r)
-    
-    def print_debug(self):
-        print('-'*50)
-        print('Parsed terminals:')
-
-        name_to_pattern = {}
-        for term in self.inc_parser.base_parser.terminals:
-            name_to_pattern[term.name] = term.pattern
-
-        for token in self.inc_parser.parsed_lexer_tokens:
-            print(f"(type: {name_to_pattern[token.type]} | value: '{token.value}')")
-        print('-'*50)
 
     @staticmethod
     def _bytes_to_string(byte_sequence: bytes) -> tuple[str, bytes]:

syncode/infer.py

@@ -56,8 +56,6 @@ class Syncode:
         new_mask_store (bool, optional): Use new DFA mask store. Defaults to False.
         
         dev_mode (bool, optional): Development mode. Defaults to False.
-
-        log_level (int, optional): Log level. Defaults to 2. 0 for no logs, 1 for minimal logs, 2 for all logs including time.
         
         opp (bool, optional): Whether to use opportunistic generation. Defaults to True.
     """
@@ -70,7 +68,6 @@ def __init__(
         grammar: Optional[str] = None,
         parse_output_only: bool = True,
         dev_mode: bool = False,
-        log_level: int = 1,
         new_mask_store: bool = False,
         parser: Literal["lr", "lalr"] = "lalr",
         seed: Optional[int] = None,
@@ -91,7 +88,6 @@ def __init__(
         self.num_samples = kwargs.get('num_return_sequences', 1)
         self.new_mask_store = new_mask_store
         self.parser = parser
-        self.log_level = log_level
 
         # Set seed
         if seed is not None:

syncode/mask_store/fsm_set.py

@@ -1,7 +1,9 @@
+import time
 import interegular
 from typing import Any, Optional, Tuple, Iterable, Dict
 from syncode.mask_store.byte_fsm import ByteFSM
-
+import logging
+logger = logging.getLogger(__name__)
 
 class JointFSMState:
     """
@@ -27,6 +29,7 @@ class FSMSet:
     Uses external ByteFSM for regex matching.
     """
     def __init__(self, terminals: Iterable['MockTerminalDef'], simplifications: Dict[str, str] = {}):
+        start_time = time.time()
         self._terminals_to_byte_fsm: Dict[str, ByteFSM] = {}  # Store ByteFSM instances
         self.anything_else = interegular.fsm.anything_else
         self._simplifications: Dict[str, str] = simplifications
@@ -41,6 +44,7 @@ def __init__(self, terminals: Iterable['MockTerminalDef'], simplifications: Dict
             # This handles the regex pattern matching
             byte_fsm = ByteFSM(terminal_regex)
             self._terminals_to_byte_fsm[terminal.name] = byte_fsm
+        logger.info(f"FSMs initialized in {time.time() - start_time:.2f} seconds")
 
     def states(self):
         """Returns all possible DFA states for all terminals."""

syncode/mask_store/lookup_table.py

@@ -4,22 +4,38 @@
 import regex
 from syncode.mask_store.mask_store import JointFSMState
 from syncode.parse_result import IndentationConstraint
-from typing import Any, Tuple, Iterable, Dict
+from typing import Any, Tuple, Iterable, Dict, Union
+import logging
+logger = logging.getLogger(__name__)
+
 
 class LookupTable:
     """
     Stores the overapproximate tokens
     """
-    def __init__(self, vocab: Iterable[str], special_token_ids: Iterable[int], indentation=False, mode='grammar_mask'):
+    def __init__(
+            self, 
+            vocab: Iterable[str], 
+            eos_token_id: int,
+            special_token_ids: Iterable[int], 
+            indent=False, 
+            mode='grammar_mask'
+        ):
         self._fsm_state_and_next_terminal_to_tokens: defaultdict = defaultdict(list)
         self._overapprox_lookup: Dict[JointFSMState, Any] = {}
         self._exact_lookup: dict = {}
         self._mode = mode
         self._vocab: Iterable[str] = vocab
-        self.indentation = indentation
+        self.indent = indent
+
+        # In the default mask, add all tokens that are special tokens except the EOS token
+        self._default_mask = torch.zeros(len(vocab), dtype=torch.bool)
+        for token_id in special_token_ids:
+            if token_id != eos_token_id:
+                self._default_mask[token_id] = 1
 
-        self._default_mask = self._get_default_mask(special_token_ids)
-        if indentation:
+        if indent:
+            logger.info("Indentation mode enabled")
             self._whitespace_tokens_map: defaultdict = defaultdict(list)
             self._indentation_to_tokens_map: defaultdict = defaultdict(list)
             self._create_indentation_to_tokens_map()
@@ -83,18 +99,14 @@ def convert_lookups_from_list_to_mask(self):
             self._exact_lookup[key] = self._list_to_mask(val)
 
         # TODO: move this logic to the lookup table
-        if self.indentation:
+        if self.indent:
             for key, val in self._whitespace_tokens_map.items():
                 self._whitespace_tokens_map[key] = self._list_to_mask(val)
             for key, val in self._indentation_to_tokens_map.items():
                 self._indentation_to_tokens_map[key] = self._list_to_mask(val)
 
-    def _get_default_mask(self, special_token_ids=None) -> torch.Tensor:
-        if special_token_ids is not None:
-            mask = torch.zeros(len(self._vocab), dtype=torch.bool)
-        else:
-            mask = copy.deepcopy(self._default_mask)
-        return mask
+    def _get_default_mask(self) -> torch.Tensor:
+        return copy.deepcopy(self._default_mask)
 
     def _create_indentation_to_tokens_map(self):
         """
@@ -107,15 +119,36 @@ def _create_indentation_to_tokens_map(self):
             else:
                 self._indentation_to_tokens_map[indent].append(token_idx)
 
-    def _get_indent_type(self, s: str) -> Tuple[bool, int]:
-        m = regex.match(r'[\t ]+', s, partial=True)
+    def _get_indent_type(self, s: Union[str, bytes]) -> Tuple[bool, int]:
+        """
+        Determine the indentation type and level from a string or bytes input.
+        
+        Args:
+            s (Union[str, bytes]): The input string or bytes to analyze
+            
+        Returns:
+            Tuple[bool, int]: A tuple containing:
+                - bool: Whether the input is entirely whitespace
+                - int: The indentation level (spaces + 4*tabs)
+        """
+        # Convert bytes to string if needed
+        if isinstance(s, bytes):
+            try:
+                s_str = s.decode('utf-8')
+            except UnicodeDecodeError:
+                # Handle decode errors by returning default values
+                return False, 0
+        else:
+            s_str = s
+            
+        m = regex.match(r'[\t ]+', s_str, partial=True)
         full_match = False
         if m != None:
             start, end = m.start(), m.end()
-            if end == len(s):
+            if end == len(s_str):
                 full_match = True
-            return full_match, s[start: end].count(' ') + 4*s[start: end].count('\t')
-        return False, 0   
+            return full_match, s_str[start: end].count(' ') + 4*s_str[start: end].count('\t')
+        return False, 0 
 
     def get_indentation_tokens(self, indent_constraint: IndentationConstraint, get_list=False):
         """