[Fix] Apply the GrammarNormalizer to grammars generated from stags. (#429)

This PR applies `GrammarNormalizer` to grammars generated from stags to
improve the efficiency at runtime.

---------

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

Author: Seven-Streams

cpp/structural_tag.cc

@@ -1020,7 +1020,7 @@ Result<Grammar, StructuralTagError> StructuralTagToGrammar(const std::string& st
   if (result.IsErr()) {
     return ResultErr(std::move(result).UnwrapErr());
   }
-  return ResultOk(std::move(result).Unwrap());
+  return ResultOk(GrammarNormalizer::Apply(std::move(result).Unwrap()));
 }
 
 }  // namespace xgrammar

tests/python/test_grammar_matcher_structural_tag.py

@@ -42,22 +42,22 @@ class Schema(BaseModel):
 basic_integer ::= (("0") | (basic_integer_1 [1-9] [0-9]*))
 basic_string ::= (("\"" basic_string_sub)) (=(root_part_0 [ \n\t]* "}"))
 root_part_0 ::= (([ \n\t]* "," [ \n\t]* "\"arg2\"" [ \n\t]* ":" [ \n\t]* basic_integer)) (=([ \n\t]* "}"))
-root ::= (("{" [ \n\t]* "\"arg1\"" [ \n\t]* ":" [ \n\t]* basic_string root_part_0 [ \n\t]* "}"))
+root ::= (("{" [ \n\t]* "\"arg1\"" [ \n\t]* ":" [ \n\t]* basic_string root_part_0 [ \n\t]* "}")) (=("</function>"))
 basic_integer_1 ::= ("" | ("-")) (=([1-9] [0-9]*))
 basic_escape_1 ::= (([\"\\/bfnrt]) | ("u" [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9])) (=(basic_string_sub_1))
 basic_string_sub_1 ::= (("\"") | ([^\0-\x1f\"\\\r\n] basic_string_sub_1) | ("\\" basic_escape_1 basic_string_sub_1)) (=([ \n\t]* [,}\]:]))
 basic_integer_2 ::= (("0") | (basic_integer_1_1 [1-9] [0-9]*))
 basic_string_1 ::= (("\"" basic_string_sub_1)) (=(root_part_0_1 [ \n\t]* "}"))
 root_part_0_1 ::= (([ \n\t]* "," [ \n\t]* "\"arg2\"" [ \n\t]* ":" [ \n\t]* basic_integer_2)) (=([ \n\t]* "}"))
-root_1 ::= (("{" [ \n\t]* "\"arg1\"" [ \n\t]* ":" [ \n\t]* basic_string_1 root_part_0_1 [ \n\t]* "}"))
+root_1 ::= (("{" [ \n\t]* "\"arg1\"" [ \n\t]* ":" [ \n\t]* basic_string_1 root_part_0_1 [ \n\t]* "}")) (=("</function>"))
 basic_integer_1_1 ::= ("" | ("-")) (=([1-9] [0-9]*))
 basic_escape_2 ::= (([\"\\/bfnrt]) | ("u" [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9] [A-Fa-f0-9])) (=(basic_string_sub_2))
 basic_string_sub_2 ::= (("\"") | ([^\0-\x1f\"\\\r\n] basic_string_sub_2) | ("\\" basic_escape_2 basic_string_sub_2)) (=([ \n\t]* [,}\]:]))
 basic_number ::= ((basic_number_7 basic_number_3 basic_number_6)) (=(root_part_0_2 [ \n\t]* "}"))
 basic_string_2 ::= (("\"" basic_string_sub_2))
 root_prop_1 ::= (("[" [ \n\t]* basic_string_2 root_prop_1_1 [ \n\t]* "]") | ("[" [ \n\t]* "]"))
 root_part_0_2 ::= (([ \n\t]* "," [ \n\t]* "\"arg4\"" [ \n\t]* ":" [ \n\t]* root_prop_1)) (=([ \n\t]* "}"))
-root_2 ::= (("{" [ \n\t]* "\"arg3\"" [ \n\t]* ":" [ \n\t]* basic_number root_part_0_2 [ \n\t]* "}"))
+root_2 ::= (("{" [ \n\t]* "\"arg3\"" [ \n\t]* ":" [ \n\t]* basic_number root_part_0_2 [ \n\t]* "}")) (=("</function>"))
 basic_number_1 ::= ("" | ("-")) (=([1-9] [0-9]*))
 basic_number_2 ::= (([0-9] basic_number_2) | ([0-9]))
 basic_number_3 ::= ("" | ("." basic_number_2)) (=(basic_number_6))
@@ -98,7 +98,6 @@ class Schema2(BaseModel):
     triggers = ["<function=f", "<function=g"]
 
     grammar = xgr.Grammar.from_structural_tag(tags, triggers)
-
     assert str(grammar) == expected_grammar_test_structural_tag
 
     accepted_inputs = [

tests/python/test_structural_tag_converter.py

@@ -72,8 +72,7 @@ def check_stag_with_instance(
 const_string_stag_grammar = [
     (
         {"type": "const_string", "value": "Hello!"},
-        r"""const_string ::= (("Hello!"))
-root ::= ((const_string))
+        r"""root ::= (("Hello!"))
 """,
     )
 ]
@@ -178,8 +177,7 @@ def test_qwen_parameter_xml_format(
                 {"type": "json_schema", "json_schema": {"type": "number"}},
             ],
         },
-        r"""const_string ::= (("Hello!"))
-basic_number ::= ((basic_number_7 basic_number_3 basic_number_6))
+        r"""basic_number ::= ((basic_number_7 basic_number_3 basic_number_6))
 root ::= ((basic_number))
 basic_number_1 ::= ("" | ("-")) (=([1-9] [0-9]*))
 basic_number_2 ::= (([0-9] basic_number_2) | ([0-9]))
@@ -188,7 +186,7 @@ def test_qwen_parameter_xml_format(
 basic_number_5 ::= (([0-9] basic_number_5) | ([0-9]))
 basic_number_6 ::= ("" | ([eE] basic_number_4 basic_number_5))
 basic_number_7 ::= (("0") | (basic_number_1 [1-9] [0-9]*)) (=(basic_number_3 basic_number_6))
-sequence ::= ((const_string root))
+sequence ::= (("Hello!" root))
 root_1 ::= ((sequence))
 """,
     )
@@ -222,8 +220,7 @@ def test_sequence_format(
                 {"type": "json_schema", "json_schema": {"type": "number"}},
             ],
         },
-        r"""const_string ::= (("Hello!"))
-basic_number ::= ((basic_number_7 basic_number_3 basic_number_6))
+        r"""basic_number ::= ((basic_number_7 basic_number_3 basic_number_6))
 root ::= ((basic_number))
 basic_number_1 ::= ("" | ("-")) (=([1-9] [0-9]*))
 basic_number_2 ::= (([0-9] basic_number_2) | ([0-9]))
@@ -232,7 +229,7 @@ def test_sequence_format(
 basic_number_5 ::= (([0-9] basic_number_5) | ([0-9]))
 basic_number_6 ::= ("" | ([eE] basic_number_4 basic_number_5))
 basic_number_7 ::= (("0") | (basic_number_1 [1-9] [0-9]*)) (=(basic_number_3 basic_number_6))
-or ::= ((const_string) | (root))
+or ::= (("Hello!") | (root))
 root_1 ::= ((or))
 """,
     )
@@ -266,7 +263,7 @@ def test_or_format(
             "end": "END",
         },
         r"""basic_number ::= ((basic_number_7 basic_number_3 basic_number_6))
-root ::= ((basic_number))
+root ::= ((basic_number)) (=("END"))
 basic_number_1 ::= ("" | ("-")) (=([1-9] [0-9]*))
 basic_number_2 ::= (([0-9] basic_number_2) | ([0-9]))
 basic_number_3 ::= ("" | ("." basic_number_2)) (=(basic_number_6))
@@ -334,8 +331,7 @@ def test_any_text_format(
 any_text_only_stag_grammar = [
     (
         {"type": "any_text"},
-        r"""any_text ::= (([\0-\U0010ffff]*))
-root ::= ((any_text))
+        r"""root ::= (([\0-\U0010ffff]*))
 """,
     )
 ]
@@ -370,8 +366,8 @@ def _get_triggered_tag_format(at_least_one: bool, stop_after_first: bool):
     (
         0,
         _get_triggered_tag_format(at_least_one=False, stop_after_first=False),
-        r"""const_string ::= (("L1"))
-const_string_1 ::= (("L2"))
+        r"""const_string ::= (("L1")) (=("A"))
+const_string_1 ::= (("L2")) (=("A"))
 triggered_tags_group ::= (("1" const_string "A") | ("2" const_string_1 "A"))
 triggered_tags ::= TagDispatch(
   ("A", triggered_tags_group),
@@ -388,7 +384,7 @@ def _get_triggered_tag_format(at_least_one: bool, stop_after_first: bool):
         r"""const_string ::= (("L1"))
 const_string_1 ::= (("L2"))
 triggered_tags_group ::= (("1" const_string "A") | ("2" const_string_1 "A"))
-triggered_tags_first ::= (("A1" const_string "A") | ("A2" const_string_1 "A"))
+triggered_tags_first ::= (("A1" const_string "A") | ("A2" const_string_1 "A")) (=(triggered_tags_sub))
 triggered_tags_sub ::= TagDispatch(
   ("A", triggered_tags_group),
   stop_eos=true,
@@ -402,8 +398,8 @@ def _get_triggered_tag_format(at_least_one: bool, stop_after_first: bool):
     (
         2,
         _get_triggered_tag_format(at_least_one=False, stop_after_first=True),
-        r"""const_string ::= (("L1"))
-const_string_1 ::= (("L2"))
+        r"""const_string ::= (("L1")) (=("A"))
+const_string_1 ::= (("L2")) (=("A"))
 triggered_tags_group ::= (("1" const_string "A") | ("2" const_string_1 "A"))
 triggered_tags ::= TagDispatch(
   ("A", triggered_tags_group),
@@ -417,8 +413,8 @@ def _get_triggered_tag_format(at_least_one: bool, stop_after_first: bool):
     (
         3,
         _get_triggered_tag_format(at_least_one=True, stop_after_first=True),
-        r"""const_string ::= (("L1"))
-const_string_1 ::= (("L2"))
+        r"""const_string ::= (("L1")) (=("A"))
+const_string_1 ::= (("L2")) (=("A"))
 triggered_tags ::= (("A1" const_string "A") | ("A2" const_string_1 "A"))
 root ::= ((triggered_tags))
 """,
@@ -466,8 +462,7 @@ def test_triggered_tag_format(
                 }
             ],
         },
-        r"""const_string ::= (("[TEXT]"))
-triggered_tags_group ::= (("" const_string "<end>"))
+        r"""triggered_tags_group ::= (("[TEXT]" "<end>"))
 triggered_tags ::= TagDispatch(
   ("<start>", triggered_tags_group),
   stop_eos=true,
@@ -527,8 +522,8 @@ def _get_triggered_tag_with_outside_tag(at_least_one: bool, stop_after_first: bo
     (
         0,
         _get_triggered_tag_with_outside_tag(at_least_one=False, stop_after_first=False),
-        r"""const_string ::= (("L1"))
-const_string_1 ::= (("L2"))
+        r"""const_string ::= (("L1")) (=("A"))
+const_string_1 ::= (("L2")) (=("A"))
 triggered_tags_group ::= (("1" const_string "A") | ("2" const_string_1 "A"))
 triggered_tags ::= TagDispatch(
   ("A", triggered_tags_group),
@@ -546,7 +541,7 @@ def _get_triggered_tag_with_outside_tag(at_least_one: bool, stop_after_first: bo
         r"""const_string ::= (("L1"))
 const_string_1 ::= (("L2"))
 triggered_tags_group ::= (("1" const_string "A") | ("2" const_string_1 "A"))
-triggered_tags_first ::= (("A1" const_string "A") | ("A2" const_string_1 "A"))
+triggered_tags_first ::= (("A1" const_string "A") | ("A2" const_string_1 "A")) (=(triggered_tags_sub))
 triggered_tags_sub ::= TagDispatch(
   ("A", triggered_tags_group),
   stop_eos=false,
@@ -561,8 +556,8 @@ def _get_triggered_tag_with_outside_tag(at_least_one: bool, stop_after_first: bo
     (
         2,
         _get_triggered_tag_with_outside_tag(at_least_one=False, stop_after_first=True),
-        r"""const_string ::= (("L1"))
-const_string_1 ::= (("L2"))
+        r"""const_string ::= (("L1")) (=("A"))
+const_string_1 ::= (("L2")) (=("A"))
 triggered_tags_group ::= (("1" const_string "A") | ("2" const_string_1 "A"))
 triggered_tags ::= TagDispatch(
   ("A", triggered_tags_group),
@@ -577,9 +572,9 @@ def _get_triggered_tag_with_outside_tag(at_least_one: bool, stop_after_first: bo
     (
         3,
         _get_triggered_tag_with_outside_tag(at_least_one=True, stop_after_first=True),
-        r"""const_string ::= (("L1"))
-const_string_1 ::= (("L2"))
-triggered_tags_sub ::= (("A1" const_string "A") | ("A2" const_string_1 "A"))
+        r"""const_string ::= (("L1")) (=("A"))
+const_string_1 ::= (("L2")) (=("A"))
+triggered_tags_sub ::= (("A1" const_string "A") | ("A2" const_string_1 "A")) (=("end"))
 triggered_tags ::= ((triggered_tags_sub "end"))
 tag ::= (("begin" triggered_tags))
 root ::= ((tag))
@@ -632,47 +627,47 @@ def _get_tags_with_separator_format(at_least_one: bool, stop_after_first: bool):
     (
         0,
         _get_tags_with_separator_format(at_least_one=False, stop_after_first=False),
-        r"""const_string ::= (("L1"))
+        r"""const_string ::= (("L1")) (=("A"))
 tag ::= (("A1" const_string "A"))
-const_string_1 ::= (("L2"))
+const_string_1 ::= (("L2")) (=("A"))
 tag_1 ::= (("A2" const_string_1 "A"))
 tags_with_separator_tags ::= ((tag) | (tag_1))
-tags_with_separator_sub ::= (("AA" tags_with_separator_tags tags_with_separator_sub) | "")
-tags_with_separator ::= ((tags_with_separator_tags tags_with_separator_sub) | "")
+tags_with_separator_sub ::= ("" | ("AA" tags_with_separator_tags tags_with_separator_sub))
+tags_with_separator ::= ("" | (tags_with_separator_tags tags_with_separator_sub))
 root ::= ((tags_with_separator))
 """,
     ),
     (
         1,
         _get_tags_with_separator_format(at_least_one=True, stop_after_first=False),
-        r"""const_string ::= (("L1"))
+        r"""const_string ::= (("L1")) (=("A"))
 tag ::= (("A1" const_string "A"))
-const_string_1 ::= (("L2"))
+const_string_1 ::= (("L2")) (=("A"))
 tag_1 ::= (("A2" const_string_1 "A"))
 tags_with_separator_tags ::= ((tag) | (tag_1))
-tags_with_separator_sub ::= (("AA" tags_with_separator_tags tags_with_separator_sub) | "")
+tags_with_separator_sub ::= ("" | ("AA" tags_with_separator_tags tags_with_separator_sub))
 tags_with_separator ::= ((tags_with_separator_tags tags_with_separator_sub))
 root ::= ((tags_with_separator))
 """,
     ),
     (
         2,
         _get_tags_with_separator_format(at_least_one=False, stop_after_first=True),
-        r"""const_string ::= (("L1"))
+        r"""const_string ::= (("L1")) (=("A"))
 tag ::= (("A1" const_string "A"))
-const_string_1 ::= (("L2"))
+const_string_1 ::= (("L2")) (=("A"))
 tag_1 ::= (("A2" const_string_1 "A"))
 tags_with_separator_tags ::= ((tag) | (tag_1))
-tags_with_separator ::= ((tags_with_separator_tags) | "")
+tags_with_separator ::= ("" | (tags_with_separator_tags))
 root ::= ((tags_with_separator))
 """,
     ),
     (
         3,
         _get_tags_with_separator_format(at_least_one=True, stop_after_first=True),
-        r"""const_string ::= (("L1"))
+        r"""const_string ::= (("L1")) (=("A"))
 tag ::= (("A1" const_string "A"))
-const_string_1 ::= (("L2"))
+const_string_1 ::= (("L2")) (=("A"))
 tag_1 ::= (("A2" const_string_1 "A"))
 tags_with_separator_tags ::= ((tag) | (tag_1))
 tags_with_separator ::= ((tags_with_separator_tags))
@@ -729,9 +724,9 @@ def _get_tags_with_separator_format_with_outside_tag(at_least_one: bool, stop_af
         _get_tags_with_separator_format_with_outside_tag(
             at_least_one=False, stop_after_first=False
         ),
-        r"""const_string ::= (("L1"))
+        r"""const_string ::= (("L1")) (=("A"))
 tag ::= (("A1" const_string "A"))
-const_string_1 ::= (("L2"))
+const_string_1 ::= (("L2")) (=("A"))
 tag_1 ::= (("A2" const_string_1 "A"))
 tags_with_separator_tags ::= ((tag) | (tag_1))
 tags_with_separator_sub ::= (("AA" tags_with_separator_tags tags_with_separator_sub) | ("end"))
@@ -743,9 +738,9 @@ def _get_tags_with_separator_format_with_outside_tag(at_least_one: bool, stop_af
     (
         1,
         _get_tags_with_separator_format_with_outside_tag(at_least_one=True, stop_after_first=False),
-        r"""const_string ::= (("L1"))
+        r"""const_string ::= (("L1")) (=("A"))
 tag ::= (("A1" const_string "A"))
-const_string_1 ::= (("L2"))
+const_string_1 ::= (("L2")) (=("A"))
 tag_1 ::= (("A2" const_string_1 "A"))
 tags_with_separator_tags ::= ((tag) | (tag_1))
 tags_with_separator_sub ::= (("AA" tags_with_separator_tags tags_with_separator_sub) | ("end"))
@@ -757,11 +752,11 @@ def _get_tags_with_separator_format_with_outside_tag(at_least_one: bool, stop_af
     (
         2,
         _get_tags_with_separator_format_with_outside_tag(at_least_one=False, stop_after_first=True),
-        r"""const_string ::= (("L1"))
+        r"""const_string ::= (("L1")) (=("A"))
 tag ::= (("A1" const_string "A"))
-const_string_1 ::= (("L2"))
+const_string_1 ::= (("L2")) (=("A"))
 tag_1 ::= (("A2" const_string_1 "A"))
-tags_with_separator_tags ::= ((tag) | (tag_1))
+tags_with_separator_tags ::= ((tag) | (tag_1)) (=("end"))
 tags_with_separator ::= ((tags_with_separator_tags "end") | ("end"))
 tag_2 ::= (("begin" tags_with_separator))
 root ::= ((tag_2))
@@ -770,11 +765,11 @@ def _get_tags_with_separator_format_with_outside_tag(at_least_one: bool, stop_af
     (
         3,
         _get_tags_with_separator_format_with_outside_tag(at_least_one=True, stop_after_first=True),
-        r"""const_string ::= (("L1"))
+        r"""const_string ::= (("L1")) (=("A"))
 tag ::= (("A1" const_string "A"))
-const_string_1 ::= (("L2"))
+const_string_1 ::= (("L2")) (=("A"))
 tag_1 ::= (("A2" const_string_1 "A"))
-tags_with_separator_tags ::= ((tag) | (tag_1))
+tags_with_separator_tags ::= ((tag) | (tag_1)) (=("end"))
 tags_with_separator ::= ((tags_with_separator_tags "end"))
 tag_2 ::= (("begin" tags_with_separator))
 root ::= ((tag_2))
@@ -1049,13 +1044,12 @@ def test_compound_format(
             },
             "end": "<end>",
         },
-        r"""const_string ::= (("[TEXT]"))
-any_text ::= TagDispatch(
+        r"""any_text ::= TagDispatch(
   stop_eos=false,
   stop_str=("<end>"),
   loop_after_dispatch=false
 )
-sequence ::= ((const_string any_text))
+sequence ::= (("[TEXT]" any_text))
 tag ::= (("<start>" sequence))
 root ::= ((tag))
 """,
@@ -1105,22 +1099,21 @@ def test_compound_format(
   stop_str=("<end2>"),
   loop_after_dispatch=false
 )
-triggered_tags_group ::= ((">" any_text ""))
-triggered_tags_first ::= (("<start2>" any_text ""))
+triggered_tags_group ::= ((">" any_text))
+triggered_tags_first ::= (("<start2>" any_text)) (=(triggered_tags_sub))
 triggered_tags_sub ::= TagDispatch(
   ("<start2", triggered_tags_group),
   stop_eos=false,
   stop_str=("<end>"),
   loop_after_dispatch=true
 )
 triggered_tags ::= ((triggered_tags_first triggered_tags_sub))
-const_string ::= (("[TEXT2]"))
 any_text_1 ::= TagDispatch(
   stop_eos=false,
   stop_str=("<end>"),
   loop_after_dispatch=false
 )
-sequence ::= ((const_string any_text_1))
+sequence ::= (("[TEXT2]" any_text_1))
 any_text_2 ::= TagDispatch(
   stop_eos=false,
   stop_str=("<end3>"),
@@ -1192,30 +1185,28 @@ def test_compound_format(
   stop_str=("<end2>"),
   loop_after_dispatch=false
 )
-triggered_tags_group ::= ((">" any_text ""))
-triggered_tags_first ::= (("<start2>" any_text ""))
+triggered_tags_group ::= ((">" any_text))
+triggered_tags_first ::= (("<start2>" any_text)) (=(triggered_tags_sub))
 triggered_tags_sub ::= TagDispatch(
   ("<start2", triggered_tags_group),
   stop_eos=true,
   stop_str=(),
   loop_after_dispatch=true
 )
 triggered_tags ::= ((triggered_tags_first triggered_tags_sub))
-const_string ::= (("[TEXT]"))
 any_text_1 ::= (([\0-\U0010ffff]*))
-sequence ::= ((const_string any_text_1))
+sequence ::= (("[TEXT]" any_text_1))
 any_text_2 ::= TagDispatch(
   stop_eos=false,
   stop_str=("<end3>"),
   loop_after_dispatch=false
 )
 tag ::= (("<start3>" any_text_2))
 tags_with_separator_tags ::= ((tag))
-tags_with_separator_sub ::= (("<sep>" tags_with_separator_tags tags_with_separator_sub) | "")
+tags_with_separator_sub ::= ("" | ("<sep>" tags_with_separator_tags tags_with_separator_sub))
 tags_with_separator ::= ((tags_with_separator_tags tags_with_separator_sub))
-const_string_1 ::= (("[TEXT2]"))
 any_text_3 ::= (([\0-\U0010ffff]*))
-sequence_1 ::= ((const_string_1 any_text_3))
+sequence_1 ::= (("[TEXT2]" any_text_3))
 or ::= ((tags_with_separator) | (sequence_1))
 or_1 ::= ((triggered_tags) | (sequence) | (or))
 root ::= ((or_1))