add some syntactic sugar to structural tags.

Signed-off-by: Yuchuan <[email protected]>

Author: Seven-Streams

python/xgrammar/compiler.py

@@ -12,7 +12,7 @@
     _convert_schema_to_str,
     _get_structural_tag_str_from_args,
 )
-from .structural_tag import StructuralTag
+from .structural_tag import Format, StructuralTag
 from .tokenizer_info import TokenizerInfo
 
 
@@ -220,7 +220,7 @@ def compile_regex(self, regex: str) -> CompiledGrammar:
 
     @overload
     def compile_structural_tag(
-        self, structural_tag: Union[StructuralTag, str, Dict[str, Any]]
+        self, structural_tag: Union[StructuralTag, str, Dict[str, Any], Format]
     ) -> CompiledGrammar: ...
 
     @overload
@@ -246,8 +246,9 @@ def compile_structural_tag(self, *args, **kwargs) -> CompiledGrammar:
 
         Parameters
         ----------
-        structural_tag : Union[StructuralTag, str, Dict[str, Any]]
+        structural_tag : Union[StructuralTag, str, Dict[str, Any], Format]
             The structural tag either as a StructuralTag object, or a JSON string or a dictionary.
+            If the input is a format enum, it will be converted to a StructuralTag object automatically.
 
         tags : List[StructuralTagItem]
             (Deprecated) The structural tags. Use StructuralTag class instead.

python/xgrammar/grammar.py

@@ -1,13 +1,13 @@
 """This module provides classes representing grammars."""
 
 import json
-from typing import Any, Dict, List, Optional, Tuple, Type, Union, overload
+from typing import Any, Dict, List, Optional, Tuple, Type, Union, get_args, overload
 
 from pydantic import BaseModel
 from typing_extensions import deprecated
 
 from .base import XGRObject, _core
-from .structural_tag import StructuralTag, StructuralTagItem
+from .structural_tag import Format, StructuralTag, StructuralTagItem
 
 
 def _convert_instance_to_str(instance: Union[str, Dict[str, Any], StructuralTag]) -> str:
@@ -111,8 +111,12 @@ def _get_structural_tag_str_from_args(args: List[Any], kwargs: Dict[str, Any]) -
         When the arguments are invalid.
     """
     if len(args) == 1:
+        possible_formats = get_args(Format)
         if isinstance(args[0], (str, dict, StructuralTag)):
             return _convert_instance_to_str(args[0])
+        elif any(isinstance(args[0], fmt) for fmt in possible_formats):
+            structural_tag = StructuralTag(format=args[0])
+            return _convert_instance_to_str(structural_tag)
         else:
             raise TypeError("Invalid argument type for from_structural_tag")
     elif len(args) == 2 and isinstance(args[0], list) and isinstance(args[1], list):
@@ -285,7 +289,7 @@ def from_regex(regex_string: str, *, print_converted_ebnf: bool = False) -> "Gra
     @overload
     @staticmethod
     def from_structural_tag(
-        structural_tag: Union[StructuralTag, str, Dict[str, Any]]
+        structural_tag: Union[StructuralTag, str, Dict[str, Any], Format]
     ) -> "Grammar": ...
 
     @overload
@@ -311,8 +315,9 @@ def from_structural_tag(*args, **kwargs) -> "Grammar":
 
         Parameters
         ----------
-        structural_tag : Union[StructuralTag, str, Dict[str, Any]]
+        structural_tag : Union[StructuralTag, str, Dict[str, Any], Format]
             The structural tag either as a StructuralTag object, or a JSON string or a dictionary.
+            If the input is a format enum, it will be converted to a StructuralTag object automatically.
 
         tags : List[StructuralTagItem]
             (Deprecated) The structural tags. Use StructuralTag class instead.

tests/python/test_structural_tag_converter.py

@@ -1,15 +1,16 @@
 import sys
 import time
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Tuple, Union
 
 import pytest
 from transformers import AutoTokenizer
 
 import xgrammar as xgr
+from xgrammar.structural_tag import Format, StructuralTag
 from xgrammar.testing import _is_grammar_accept_string
 
 PROFILER_ON = True
-tokenizer_id = "meta-llama/Llama-3.1-8B-Instruct"
+tokenizer_id = "meta-llama/Meta-Llama-3-8B-Instruct"
 
 
 class Profiler:
@@ -22,13 +23,16 @@ def __init__(self, tokenizer_id: str):
             self.tokenizer_info, max_threads=16, cache_enabled=False
         )
 
-    def profile_stag(self, structural_tag_format: Dict[str, Any], instance: str):
-        structural_tag = {"type": "structural_tag", "format": structural_tag_format}
+    def profile_stag(
+        self, structural_tag: Union[Dict[str, Any], StructuralTag, Format], instance: str
+    ):
+        if isinstance(structural_tag, Dict):
+            structural_tag = {"type": "structural_tag", "format": structural_tag}
         time_begin = time.monotonic_ns()
         compiled_grammar = self.compiler.compile_structural_tag(structural_tag)
         time_end = time.monotonic_ns()
         compiler_duration = time_end - time_begin
-        print(f"Compiling structural tag {structural_tag_format}")
+        print(f"Compiling structural tag {structural_tag}")
         print(f"Compile time: {compiler_duration / 1000 / 1000} ms")
         matcher = xgr.GrammarMatcher(compiled_grammar)
         token_bitmask = xgr.allocate_token_bitmask(1, self.tokenizer_info.vocab_size)
@@ -69,6 +73,33 @@ def check_stag_with_instance(
         profiler.profile_stag(structural_tag_format, instance)
 
 
+def check_stag_with_instance_in_stag_style(
+    structural_tag: StructuralTag,
+    instance: str,
+    is_accepted: bool = True,
+    debug_print: bool = False,
+):
+    stag_grammar = xgr.Grammar.from_structural_tag(structural_tag)
+    accepted = _is_grammar_accept_string(stag_grammar, instance, debug_print=debug_print)
+    assert accepted == is_accepted
+    if PROFILER_ON:
+        profiler.profile_stag(structural_tag, instance)
+
+
+def check_stag_with_instance_in_format_style(
+    structural_tag_format: Format,
+    instance: str,
+    is_accepted: bool = True,
+    debug_print: bool = False,
+):
+
+    stag_grammar = xgr.Grammar.from_structural_tag(structural_tag_format)
+    accepted = _is_grammar_accept_string(stag_grammar, instance, debug_print=debug_print)
+    assert accepted == is_accepted
+    if PROFILER_ON:
+        profiler.profile_stag(structural_tag_format, instance)
+
+
 const_string_stag_grammar = [
     (
         {"type": "const_string", "value": "Hello!"},
@@ -1959,5 +1990,16 @@ def test_from_structural_tag_with_structural_tag_instance(
     assert _is_grammar_accept_string(grammar, instance) == is_accepted
 
 
+@pytest.mark.parametrize(
+    "stag_format, instance, is_accepted", basic_structural_tags_instance_is_accepted
+)
+def test_from_structural_tag_with_format_instance(
+    stag_format: xgr.structural_tag.Format, instance: str, is_accepted: bool
+):
+    stag = xgr.StructuralTag(format=stag_format)
+    grammar = xgr.Grammar.from_structural_tag(stag)
+    assert _is_grammar_accept_string(grammar, instance) == is_accepted
+
+
 if __name__ == "__main__":
     pytest.main(sys.argv)