Merge pull request #190 from structuredllm/gemma

Fix Phi-4 issue

Author: shubhamugare

syncode/language_model.py

@@ -111,7 +111,8 @@ def generate_grammar_constrained_completion(
             stop_criteria = []
 
         # Generate completions
-        if self.opp and (gen_mode == GenerationMode.SAMPLE or gen_mode == GenerationMode.GREEDY_SEARCH) and batch_size == 1: # Use our own implementation for greedy search and sampling
+        if self.opp and (gen_mode == GenerationMode.SAMPLE or gen_mode == GenerationMode.GREEDY_SEARCH) and batch_size == 1: 
+            # Use our own implementation for greedy search and sampling
             generated_ids = self._generate(
                 inputs, 
                 gen_config, 
@@ -239,6 +240,11 @@ def _generate(
             # (the clone itself is always small)
             next_token_scores = outputs.logits[:, -1, :].to(copy=True, dtype=torch.float32, device=token_ids.device)
 
+            if len(next_token_scores.shape) == 3:
+                # FIXME: This is a strange behaviour for some models like Phi-4
+                # We expect next_token_scores to be of shape (batch_size, vocab_size)
+                next_token_scores = next_token_scores[:, -1, :]
+
             if grammar_decoder is not None:
                 next_token = self._get_next_token(gen_mode, token_ids, logits_processor, next_token_scores)
                 is_valid = grammar_decoder.is_valid(token_ids, next_token)

syncode/mask_store/byte_tokenizer.py

@@ -188,6 +188,13 @@ def __init__(self, tokenizer, vocab_type=None):
         # Cache special token IDs as a set for faster lookups
         self.special_token_ids = set(getattr(tokenizer, "all_special_ids", []))
 
+        # Added tokens are typically special tokens
+        # if added_tokens_decoder is not None self.tokenizer.added_tokens_decoder.keys()
+        # to special_token_ids
+        if hasattr(tokenizer, "added_tokens_decoder"):
+            self.special_token_ids.update(tokenizer.added_tokens_decoder.keys())
+
+
     @classmethod
     def from_pretrained(cls, model_id, vocab_type=None):
         """