sglang 0.4.3.post2__py3-none-any.whl → 0.4.3.post3__py3-none-any.whl

sglang/srt/constrained/llguidance_backend.py

@@ -0,0 +1,151 @@
+# Copyright 2023-2024 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Constrained decoding with llguidance backend."""
+
+import json
+import os
+from typing import List, Optional, Tuple
+
+import llguidance
+import llguidance.hf
+import llguidance.torch
+import torch
+from llguidance.gbnf_to_lark import any_to_lark
+
+from sglang.srt.constrained.base_grammar_backend import (
+    BaseGrammarBackend,
+    BaseGrammarObject,
+)
+
+
+class GuidanceGrammar(BaseGrammarObject):
+    def __init__(
+        self, llguidance_tokenizer: llguidance.LLTokenizer, serialized_grammar: str
+    ):
+        self.llguidance_tokenizer = llguidance_tokenizer
+        self.serialized_grammar = serialized_grammar
+
+        # TODO: add support for fast-forward tokens in the future
+        self.ll_interpreter = llguidance.LLInterpreter(
+            self.llguidance_tokenizer,
+            self.serialized_grammar,
+            enable_backtrack=False,
+            enable_ff_tokens=False,
+            log_level=int(os.environ.get("LLGUIDANCE_LOG_LEVEL", "1")),
+        )
+        self.pending_ff_tokens: list[int] = []
+        self.finished = False
+        self.bitmask = None
+
+    def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
+        if len(self.pending_ff_tokens) > 0:
+            s = self.llguidance_tokenizer.decode_str(self.pending_ff_tokens)
+            ff_tokens = self.pending_ff_tokens
+            self.pending_ff_tokens = []
+            return (ff_tokens, s)
+
+        return None
+
+    def jump_forward_str_state(self, helper: Tuple[List[int], str]) -> Tuple[str, int]:
+        return "", -1
+
+    def jump_and_retokenize(
+        self, old_output_ids: List[int], new_output_ids: List[int], next_state: int
+    ):
+        pass
+
+    def accept_token(self, token: int):
+        backtrack, ff_tokens = self.ll_interpreter.commit_token(token)
+        if len(ff_tokens) > 0 and backtrack == 0:
+            # first token is last generated token
+            ff_tokens = ff_tokens[1:]
+            self.pending_ff_tokens.extend(ff_tokens)
+
+    def fill_vocab_mask(self, vocab_mask: torch.Tensor, idx: int) -> None:
+        if len(self.pending_ff_tokens) > 0:
+            # if we have pending fast-forward tokens,
+            # just return them immediately
+            ff_token = self.pending_ff_tokens.pop(0)
+            vocab_mask[idx, :] = 0
+            vocab_mask[idx, ff_token // 32] = 1 << (ff_token % 32)
+            return
+
+        if self.ll_interpreter.has_pending_stop():
+            self.finished = True
+
+        llguidance.torch.fill_next_token_bitmask(self.ll_interpreter, vocab_mask, idx)
+
+    def allocate_vocab_mask(
+        self, vocab_size: int, batch_size: int, device
+    ) -> torch.Tensor:
+        if self.bitmask is None or self.bitmask.shape[0] < batch_size:
+            # only create bitmask when batch gets larger
+            self.bitmask = llguidance.torch.allocate_token_bitmask(
+                batch_size, self.llguidance_tokenizer.vocab_size
+            )
+            bitmask = self.bitmask
+        else:
+            bitmask = self.bitmask[:batch_size]
+
+        return bitmask
+
+    @staticmethod
+    def move_vocab_mask(vocab_mask: torch.Tensor, device) -> torch.Tensor:
+        return vocab_mask.to(device, non_blocking=True)
+
+    @staticmethod
+    def apply_vocab_mask(logits: torch.Tensor, vocab_mask: torch.Tensor) -> None:
+        llguidance.torch.apply_token_bitmask_inplace(logits, vocab_mask)
+
+    def copy(self):
+        return GuidanceGrammar(
+            llguidance_tokenizer=self.llguidance_tokenizer,
+            serialized_grammar=self.serialized_grammar,
+        )
+
+
+class GuidanceBackend(BaseGrammarBackend):
+    def __init__(self, tokenizer, whitespace_pattern: Optional[str] = None):
+        super().__init__()
+
+        self.tokenizer = tokenizer
+        self.whitespace_flexible = (
+            True if whitespace_pattern == "whitespace_flexible" else False
+        )
+        self.llguidance_tokenizer = llguidance.hf.from_tokenizer(self.tokenizer, None)
+
+    def _from_serialized(self, serialized_grammar) -> GuidanceGrammar:
+        return GuidanceGrammar(
+            llguidance_tokenizer=self.llguidance_tokenizer,
+            serialized_grammar=serialized_grammar,
+        )
+
+    def dispatch_json(self, key_string: str) -> GuidanceGrammar:
+        json_schema = key_string
+        compiler = llguidance.JsonCompiler(whitespace_flexible=self.whitespace_flexible)
+        serialized_grammar = compiler.compile(json_schema)
+        return self._from_serialized(serialized_grammar)
+
+    def dispatch_regex(self, key_string: str) -> GuidanceGrammar:
+        compiler = llguidance.RegexCompiler()
+        serialized_grammar = compiler.compile(regex=key_string)
+        return self._from_serialized(serialized_grammar)
+
+    def dispatch_ebnf(self, key_string: str) -> GuidanceGrammar:
+        compiler = llguidance.LarkCompiler()
+        serialized_grammar = compiler.compile(any_to_lark(key_string))
+        return self._from_serialized(serialized_grammar)
+
+    def dispatch_structural_tag(self, key_string: str):
+        return super().dispatch_structural_tag(key_string)

sglang/srt/constrained/outlines_backend.py

@@ -28,17 +28,11 @@ from sglang.srt.constrained.base_grammar_backend import (
     BaseGrammarObject,
 )
 from sglang.srt.constrained.outlines_jump_forward import OutlinesJumpForwardMap
-from sglang.srt.utils import is_hip
 
-is_hip_ = is_hip()
-
-if is_hip_:
+try:
+    from outlines.fsm.json_schema import build_regex_from_schema
+except ImportError:
     from outlines_core.fsm.json_schema import build_regex_from_schema
-else:
-    try:
-        from outlines.fsm.json_schema import build_regex_from_schema
-    except ImportError:
-        from outlines_core.fsm.json_schema import build_regex_from_schema
 
 
 logger = logging.getLogger(__name__)
@@ -121,7 +115,6 @@ class OutlinesGrammarBackend(BaseGrammarBackend):
         self,
         tokenizer,
         whitespace_pattern: bool,
-        allow_jump_forward: bool,
     ):
         super().__init__()
 
@@ -146,27 +139,9 @@ class OutlinesGrammarBackend(BaseGrammarBackend):
             self.outlines_tokenizer.vocabulary = (
                 self.outlines_tokenizer.tokenizer.get_vocab()
             )
-        self.allow_jump_forward = allow_jump_forward
         self.whitespace_pattern = whitespace_pattern
 
-    def init_value_impl(self, key: Tuple[str, str]) -> OutlinesGrammar:
-        key_type, key_string = key
-        if key_type == "json":
-            try:
-                regex = build_regex_from_object(
-                    key_string,
-                    whitespace_pattern=self.whitespace_pattern,
-                )
-            except (NotImplementedError, json.decoder.JSONDecodeError) as e:
-                logger.warning(
-                    f"Skip invalid json_schema: json_schema={key_string}, {e=}"
-                )
-                return None
-        elif key_type == "regex":
-            regex = key_string
-        else:
-            raise ValueError(f"Invalid key_type: {key_type}")
-
+    def _compile_regex(self, regex: str) -> Optional[OutlinesGrammar]:
         try:
             if hasattr(RegexGuide, "from_regex"):
                 # outlines >= 0.1.1
@@ -178,12 +153,28 @@ class OutlinesGrammarBackend(BaseGrammarBackend):
             logger.warning(f"skip invalid regex schema: {regex=}, {e=}")
             return None
 
-        if self.allow_jump_forward:
-            jump_forward_map = OutlinesJumpForwardMap(regex)
-        else:
-            jump_forward_map = None
+        jump_forward_map = None
         return OutlinesGrammar(guide, jump_forward_map)
 
+    def dispatch_ebnf(self, key_string: str):
+        return super().dispatch_ebnf(key_string)
+
+    def dispatch_structural_tag(self, key_string: str):
+        return super().dispatch_structural_tag(key_string)
+
+    def dispatch_json(self, key_string: str):
+        try:
+            regex = build_regex_from_object(
+                key_string,
+                whitespace_pattern=self.whitespace_pattern,
+            )
+        except (NotImplementedError, json.decoder.JSONDecodeError) as e:
+            logger.warning(f"Skip invalid json_schema: json_schema={key_string}, {e=}")
+        return self._compile_regex(regex)
+
+    def dispatch_regex(self, key_string: str):
+        return self._compile_regex(key_string)
+
 
 def build_regex_from_object(
     object: Union[str, BaseModel, Dict], whitespace_pattern: Optional[str] = None

sglang/srt/constrained/xgrammar_backend.py

@@ -13,15 +13,16 @@
 # ==============================================================================
 """Constrained decoding with xgrammar backend."""
 
+import json
 import logging
-from typing import List, Tuple
+from typing import List, Optional, Tuple, Union
 
 import torch
 from xgrammar import (
     CompiledGrammar,
-    Grammar,
     GrammarCompiler,
     GrammarMatcher,
+    StructuralTagItem,
     TokenizerInfo,
     allocate_token_bitmask,
     apply_token_bitmask_inplace,
@@ -41,17 +42,22 @@ MAX_ROLLBACK_TOKENS = 200
 class XGrammarGrammar(BaseGrammarObject):
 
     def __init__(
-        self, matcher: GrammarMatcher, vocab_size: int, ctx: CompiledGrammar
+        self,
+        matcher: GrammarMatcher,
+        vocab_size: int,
+        ctx: CompiledGrammar,
+        override_stop_tokens: Optional[Union[List[int], int]],
     ) -> None:
         self.matcher = matcher
         self.vocab_size = vocab_size
         self.ctx = ctx
+        self.override_stop_tokens = override_stop_tokens
         self.finished = False
 
     def accept_token(self, token: int):
         assert self.matcher.accept_token(token)
 
-    def try_jump_forward(self, tokenizer) -> Tuple[List[int], str]:
+    def try_jump_forward(self, tokenizer) -> Optional[Tuple[List[int], str]]:
         s = self.matcher.find_jump_forward_string()
         if s:
             return [], s
@@ -95,8 +101,14 @@ class XGrammarGrammar(BaseGrammarObject):
         apply_token_bitmask_inplace(logits, vocab_mask)
 
     def copy(self):
-        matcher = GrammarMatcher(self.ctx, max_rollback_tokens=MAX_ROLLBACK_TOKENS)
-        return XGrammarGrammar(matcher, self.vocab_size, self.ctx)
+        matcher = GrammarMatcher(
+            self.ctx,
+            max_rollback_tokens=MAX_ROLLBACK_TOKENS,
+            override_stop_tokens=self.override_stop_tokens,
+        )
+        return XGrammarGrammar(
+            matcher, self.vocab_size, self.ctx, self.override_stop_tokens
+        )
 
 
 class XGrammarGrammarBackend(BaseGrammarBackend):
@@ -110,42 +122,61 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
         tokenizer_info = TokenizerInfo.from_huggingface(
             tokenizer, vocab_size=vocab_size
         )
+        override_stop_tokens = None
+
         self.grammar_compiler = GrammarCompiler(tokenizer_info=tokenizer_info)
         self.vocab_size = vocab_size
+        self.override_stop_tokens = override_stop_tokens
 
-    def init_value_impl(self, key: Tuple[str, str]) -> XGrammarGrammar:
-
-        key_type, key_string = key
-        if key_type == "json":
-            try:
-                if key_string == "$$ANY$$":
-                    ctx = self.grammar_compiler.compile_builtin_json_grammar()
-                else:
-                    ctx = self.grammar_compiler.compile_json_schema(schema=key_string)
-            except RuntimeError as e:
-                logging.warning(
-                    f"Skip invalid json_schema: json_schema={key_string}, {e=}"
-                )
-                return None
-        elif key_type == "ebnf":
-            try:
-                ctx = self.grammar_compiler.compile_grammar(key_string)
-            except RuntimeError as e:
-                logging.warning(f"Skip invalid ebnf: ebnf={key_string}, {e=}")
-                return None
-        elif key_type == "regex":
-            try:
-                ctx = self.grammar_compiler.compile_grammar(
-                    Grammar.from_regex(key_string)
-                )
-            except RuntimeError as e:
-                logging.warning(f"Skip invalid regex: regex={key_string}, {e=}")
-                return None
-        else:
-            raise ValueError(f"Invalid key_type: {key_type}")
-
+    def _from_context(self, ctx: CompiledGrammar) -> XGrammarGrammar:
         matcher = GrammarMatcher(ctx, max_rollback_tokens=MAX_ROLLBACK_TOKENS)
-        return XGrammarGrammar(matcher, self.vocab_size, ctx)
+        return XGrammarGrammar(matcher, self.vocab_size, ctx, self.override_stop_tokens)
+
+    def dispatch_json(self, key_string: str) -> Optional[XGrammarGrammar]:
+        try:
+            if key_string == "$$ANY$$":
+                ctx = self.grammar_compiler.compile_builtin_json_grammar()
+            else:
+                ctx = self.grammar_compiler.compile_json_schema(schema=key_string)
+        except RuntimeError as e:
+            logging.warning(f"Skip invalid json_schema: json_schema={key_string}, {e=}")
+            return None
+        return self._from_context(ctx)
+
+    def dispatch_ebnf(self, key_string: str) -> Optional[XGrammarGrammar]:
+        try:
+            ctx = self.grammar_compiler.compile_grammar(key_string)
+        except RuntimeError as e:
+            logging.warning(f"Skip invalid ebnf: ebnf={key_string}, {e=}")
+            return None
+        return self._from_context(ctx)
+
+    def dispatch_regex(self, key_string: str) -> Optional[XGrammarGrammar]:
+        try:
+            ctx = self.grammar_compiler.compile_regex(key_string)
+        except RuntimeError as e:
+            logging.warning(f"Skip invalid regex: regex={key_string}, {e=}")
+            return None
+        return self._from_context(ctx)
+
+    def dispatch_structural_tag(self, key_string: str) -> Optional[XGrammarGrammar]:
+        try:
+            structural_tag = json.loads(key_string)
+            tags = [
+                StructuralTagItem(
+                    begin=structure["begin"],
+                    schema=json.dumps(structure["schema"]),
+                    end=structure["end"],
+                )
+                for structure in structural_tag["structures"]
+            ]
+            ctx = self.grammar_compiler.compile_structural_tag(
+                tags, structural_tag["triggers"]
+            )
+        except RuntimeError as e:
+            logging.warning(f"Skip invalid regex: regex={key_string}, {e=}")
+            return None
+        return self._from_context(ctx)
 
     def reset(self):
         if self.grammar_compiler: