sglang 0.3.5__py3-none-any.whl → 0.3.5.post2__py3-none-any.whl

sglang/srt/constrained/grammar.py

@@ -1,190 +0,0 @@
-"""
-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.
-"""
-
-"""Cache for the compressed finite state machine."""
-import logging
-from typing import List, Optional, Tuple, Union
-
-import torch
-
-from sglang.srt.constrained import GrammarMatcher, RegexGuide
-from sglang.srt.constrained.bnf_cache import BNFCache
-from sglang.srt.constrained.fsm_cache import FSMCache
-from sglang.srt.constrained.jump_forward import JumpForwardCache, JumpForwardMap
-
-# from sglang.srt.managers.schedule_batch import Req
-
-logger = logging.getLogger(__name__)
-
-INIT_INCREMENTAL_DETOKENIZATION_OFFSET = 5
-
-
-class XGrammarJump:
-    pass
-
-
-class JumpHelper:
-    data: Union[List, str]
-    state: int
-    suffix_ids: List[int]
-
-    def __init__(
-        self, data: Union[List, str] = "", state: int = -1, suffix_ids=[]
-    ) -> None:
-        self.data = data
-        self.state = state
-        self.suffix_ids = suffix_ids
-
-    def can_jump(self):
-        return len(self.data) > 0
-
-
-class Grammar:
-    grammar: Union[GrammarMatcher, Tuple[RegexGuide, int]]
-    jump_map: Union[XGrammarJump, JumpForwardMap, None]
-
-    def __init__(
-        self,
-        grammar: Union[GrammarMatcher, Tuple[RegexGuide, int]],
-        jump_map: Union[XGrammarJump, JumpForwardMap, None],
-    ) -> None:
-        self.grammar = grammar
-        self.jump_map = jump_map
-
-    def accept_token(self, token: int):
-        if isinstance(self.grammar, GrammarMatcher):
-            assert self.grammar.accept_token(token)
-        else:
-            guide, state = self.grammar
-            self.grammar = guide, guide.get_next_state(state, token)
-
-    def try_jump(self, tokenizer) -> JumpHelper:
-        if isinstance(self.jump_map, XGrammarJump):
-            assert isinstance(self.grammar, GrammarMatcher)
-            return JumpHelper(self.grammar.find_jump_forward_string())
-        elif isinstance(self.jump_map, JumpForwardMap):
-            assert isinstance(self.grammar, Tuple)
-
-            _, state = self.grammar
-            jump_forward_bytes = self.jump_map.jump_forward_byte(state)
-            if jump_forward_bytes is None or len(jump_forward_bytes) == 0:
-                return JumpHelper()  # can't jump
-
-            # preprocess the jump forward string
-            suffix_bytes = []
-            continuation_range = range(0x80, 0xC0)
-            cur_state = state
-            while (
-                len(jump_forward_bytes)
-                and jump_forward_bytes[0][0] in continuation_range
-            ):
-                # continuation bytes
-                byte_edge = jump_forward_bytes.pop(0)
-                suffix_bytes.append(byte_edge[0])
-                cur_state = byte_edge[1]
-
-            suffix_tokens = [f"<0x{hex(b)[2:].upper()}>" for b in suffix_bytes]
-            suffix_ids = tokenizer.convert_tokens_to_ids(suffix_tokens)
-            return JumpHelper(suffix_ids, cur_state, suffix_bytes)
-        else:
-            return JumpHelper()  # can't jump
-
-    def jump_forward_str_state(self, helper: JumpHelper) -> Tuple[str, int]:
-        if isinstance(helper.data, str):
-            return helper.data, -1
-        else:
-            assert isinstance(self.jump_map, JumpForwardMap)
-            return self.jump_map.jump_forward_symbol(helper.state)
-
-    def jump_and_retokenize(
-        self, old_output_ids: List[int], new_output_ids: List[int], next_state: int
-    ):
-        if isinstance(self.grammar, GrammarMatcher):
-            k = 0
-            for i, old_id in enumerate(old_output_ids):
-                if old_id == new_output_ids[i]:
-                    k = i + 1
-                else:
-                    break
-
-            # rollback to the last token that is the same
-            if k < len(old_output_ids):
-                self.grammar.rollback(len(old_output_ids) - k)
-
-            for i in range(k, len(new_output_ids)):
-                assert self.grammar.accept_token(new_output_ids[i])
-        else:
-            self.grammar = self.grammar[0], next_state
-
-    def fill_vocab_mask(self, vocab_mask: torch.Tensor, vocab_size: int):
-        if isinstance(self.grammar, GrammarMatcher):
-            # Note that this bitmask is a bitset, not bool
-            bitmask = self.grammar.find_next_token_bitmask()
-            # Mask the tokens that are not allowed
-            vocab_mask[
-                self.grammar.get_rejected_tokens_from_bitmask(bitmask, vocab_size)
-            ] = 1
-        else:
-            guide, state = self.grammar
-            vocab_mask.fill_(1)
-            vocab_mask[guide.get_next_instruction(state).tokens] = 0
-
-
-class GrammarCache:
-    grammar_cache: Union[BNFCache, FSMCache]
-    jump_cache: Union[XGrammarJump, JumpForwardCache, None]
-
-    def __init__(
-        self,
-        tokenizer_path,
-        tokenizer_args_dict,
-        skip_tokenizer_init=False,
-        whitespace_patterns=None,
-        backend=None,
-        allow_jump=False,
-    ):
-        if backend == "xgrammar":
-            self.grammar_cache = BNFCache(
-                tokenizer_path=tokenizer_path,
-                tokenizer_args_dict=tokenizer_args_dict,
-                skip_tokenizer_init=skip_tokenizer_init,
-                whitespace_patterns=whitespace_patterns,
-            )
-            self.jump_cache = XGrammarJump() if allow_jump else None
-        else:
-            assert backend == "outlines"
-            self.grammar_cache = FSMCache(
-                tokenizer_path=tokenizer_path,
-                tokenizer_args_dict=tokenizer_args_dict,
-                skip_tokenizer_init=skip_tokenizer_init,
-                constrained_json_whitespace_pattern=whitespace_patterns,
-                enable=True,
-            )
-            self.jump_cache = JumpForwardCache() if allow_jump else None
-
-    def query(self, key: Tuple[str, str], vocab_size: int) -> Grammar:
-        if isinstance(self.grammar_cache, BNFCache):
-            assert not isinstance(self.jump_cache, JumpForwardCache)
-            return Grammar(self.grammar_cache.query(key, vocab_size), self.jump_cache)
-        else:
-            jump_map = None
-            guide, regex = self.grammar_cache.query(key)
-            if isinstance(self.jump_cache, JumpForwardCache):
-                jump_map = self.jump_cache.query(regex)
-            return Grammar((guide, 0), jump_map)
-
-    def reset(self):
-        if isinstance(self.grammar_cache, FSMCache):
-            self.grammar_cache.reset()
-        if isinstance(self.jump_cache, JumpForwardCache):
-            self.jump_cache.reset()

sglang/srt/constrained/jump_forward.py

@@ -1,203 +0,0 @@
-"""
-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.
-"""
-
-"""
-Faster constrained decoding.
-Reference: https://lmsys.org/blog/2024-02-05-compressed-fsm/
-"""
-
-import dataclasses
-import logging
-from collections import defaultdict
-
-import interegular
-import outlines.caching
-from interegular import InvalidSyntax
-
-from sglang.srt.constrained import (
-    FSMInfo,
-    disk_cache,
-    make_byte_level_fsm,
-    make_deterministic_fsm,
-)
-from sglang.srt.constrained.base_tool_cache import BaseToolCache
-
-IP_REGEX = r"((25[0-5]|2[0-4]\d|[01]?\d\d?)\.){3}(25[0-5]|2[0-4]\d|[01]?\d\d?)"
-
-logger = logging.getLogger(__name__)
-
-
-@dataclasses.dataclass
-class JumpEdge:
-    symbol: str = None
-    symbol_next_state: int = None
-    byte: int = None
-    byte_next_state: int = None
-
-
-class JumpForwardMap:
-    def __init__(self, regex_string):
-        @disk_cache()
-        def _init_state_to_jump_forward(regex_string):
-            try:
-                regex_pattern = interegular.parse_pattern(regex_string)
-            except InvalidSyntax as e:
-                logger.warning(f"skip invalid regex: {regex_string}, {e=}")
-                self.state_to_jump_forward = None
-                return
-
-            byte_fsm = make_byte_level_fsm(
-                regex_pattern.to_fsm().reduce(), keep_utf8=True
-            )
-            regex_fsm, _ = make_deterministic_fsm(byte_fsm)
-
-            fsm_info: FSMInfo = regex_fsm.fsm_info
-
-            symbol_to_id = fsm_info.alphabet_symbol_mapping
-            id_to_symbol = {}
-            for symbol, id_ in symbol_to_id.items():
-                id_to_symbol.setdefault(id_, []).append(symbol)
-
-            transitions = fsm_info.transitions
-
-            outgoings_ct = defaultdict(int)
-            # NOTE(lsyin): Final states can lead to terminate, so they have one outgoing edge naturally
-            for s in fsm_info.finals:
-                outgoings_ct[s] = 1
-
-            state_to_jump_forward = {}
-            for (state, id_), next_state in transitions.items():
-                if id_ == fsm_info.alphabet_anything_value:
-                    # Arbitrarily symbol cannot be recognized as jump forward
-                    continue
-
-                symbols = id_to_symbol[id_]
-                for c in symbols:
-                    if len(c) > 1:
-                        # Skip byte level transitions like c = "5E"
-                        continue
-
-                    outgoings_ct[state] += 1
-                    if outgoings_ct[state] > 1:
-                        if state in state_to_jump_forward:
-                            del state_to_jump_forward[state]
-                        break
-
-                    state_to_jump_forward[state] = JumpEdge(
-                        symbol=c,
-                        symbol_next_state=next_state,
-                    )
-
-            # Process the byte level jump forward
-            outgoings_ct = defaultdict(int)
-            for s in fsm_info.finals:
-                outgoings_ct[s] = 1
-
-            for (state, id_), next_state in transitions.items():
-                if id_ == fsm_info.alphabet_anything_value:
-                    continue
-                symbols = id_to_symbol[id_]
-                for c in symbols:
-                    byte_ = None
-                    if len(c) == 1 and ord(c) < 0x80:
-                        # ASCII character
-                        byte_ = ord(c)
-                    elif len(c) > 1:
-                        # FIXME: This logic is due to the leading \x00
-                        # https://github.com/outlines-dev/outlines/pull/930
-                        byte_ = int(symbols[0][1:], 16)
-
-                    if byte_ is not None:
-                        outgoings_ct[state] += 1
-                        if outgoings_ct[state] > 1:
-                            if state in state_to_jump_forward:
-                                del state_to_jump_forward[state]
-                            break
-                        e = state_to_jump_forward.get(state, JumpEdge())
-                        e.byte = byte_
-                        e.byte_next_state = next_state
-                        state_to_jump_forward[state] = e
-
-            return state_to_jump_forward
-
-        self.state_to_jump_forward = _init_state_to_jump_forward(regex_string)
-
-    def jump_forward_symbol(self, state):
-        jump_forward_str = ""
-        next_state = state
-        while state in self.state_to_jump_forward:
-            e = self.state_to_jump_forward[state]
-            if e.symbol is None:
-                break
-            jump_forward_str += e.symbol
-            next_state = e.symbol_next_state
-            state = next_state
-
-        return jump_forward_str, next_state
-
-    def jump_forward_byte(self, state):
-        if state not in self.state_to_jump_forward:
-            return None
-
-        jump_forward_bytes = []
-        next_state = None
-        while state in self.state_to_jump_forward:
-            e = self.state_to_jump_forward[state]
-            assert e.byte is not None and e.byte_next_state is not None
-            jump_forward_bytes.append((e.byte, e.byte_next_state))
-            next_state = e.byte_next_state
-            state = next_state
-
-        return jump_forward_bytes
-
-    def is_jump_forward_symbol_state(self, state):
-        return (
-            state in self.state_to_jump_forward
-            and self.state_to_jump_forward[state].symbol is not None
-        )
-
-
-class JumpForwardCache(BaseToolCache):
-    def __init__(self):
-        super().__init__()
-
-    def init_value(self, regex):
-        forward_map = JumpForwardMap(regex)
-        if forward_map.state_to_jump_forward:
-            return forward_map
-        else:
-            return None
-
-
-def test_main(regex_string):
-    jump_forward_map = JumpForwardMap(regex_string)
-    for state, e in jump_forward_map.state_to_jump_forward.items():
-        if e.symbol is not None:
-            jump_forward_str, next_state = jump_forward_map.jump_forward_symbol(state)
-            print(f"{state} -> {next_state}", jump_forward_str)
-        bytes_ = jump_forward_map.jump_forward_byte(state)
-        print(f"{state} -> {bytes_[-1][1]}", [hex(b) for b, _ in bytes_])
-
-
-if __name__ == "__main__":
-    import outlines
-
-    outlines.caching.clear_cache()
-    test_main(r"The google's DNS sever address is " + IP_REGEX)
-    test_main(r"霍格沃茨特快列车|霍比特人比尔博")
-    # 霍格: \xe9\x9c\x8d \xe6\xa0\xbc ...
-    # 霍比: \xe9\x9c\x8d \xe6\xaf\x94 ...
-
-    test_main(r"[-+]?[0-9]+[ ]*")