[Feature] integrate Structural Tag in xgrammar backend for function calling (#3566)

Co-authored-by: shuaills <shishuaiuoe@gmail.com>
This commit is contained in:
mlmz
2025-02-28 15:33:41 +08:00
committed by GitHub
parent eec3f6d1eb
commit bac414ab53
5 changed files with 44 additions and 1 deletions

View File

@@ -13,6 +13,7 @@
# ==============================================================================
"""Constrained decoding with xgrammar backend."""
import json
import logging
from typing import List, Tuple
@@ -21,6 +22,7 @@ from xgrammar import (
CompiledGrammar,
GrammarCompiler,
GrammarMatcher,
StructuralTagItem,
TokenizerInfo,
allocate_token_bitmask,
apply_token_bitmask_inplace,
@@ -138,6 +140,23 @@ class XGrammarGrammarBackend(BaseGrammarBackend):
except RuntimeError as e:
logging.warning(f"Skip invalid regex: regex={key_string}, {e=}")
return None
elif key_type == "structural_tag":
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
else:
raise ValueError(f"Invalid key_type: {key_type}")