95 lines
3 KiB
Python
95 lines
3 KiB
Python
"""GBNF — Convert AST to GBNF grammar format for llama.cpp constrained decoding."""
|
|
|
|
from .grammar import (
|
|
Symbol, Concat, Alt, Plus, Optional, Star, Epsilon, Empty,
|
|
)
|
|
|
|
|
|
def _needs_group(node):
|
|
"""Check if a node needs parentheses in GBNF output."""
|
|
return isinstance(node, (Alt, Concat))
|
|
|
|
|
|
def _node_to_gbnf(node):
|
|
"""Convert AST node to GBNF fragment string."""
|
|
if isinstance(node, Symbol):
|
|
escaped = node.value.replace('\\', '\\\\').replace('"', '\\"')
|
|
return f'"{escaped}"'
|
|
if isinstance(node, (Epsilon, Empty)):
|
|
return ''
|
|
if isinstance(node, Concat):
|
|
parts = []
|
|
for child in node.parts:
|
|
frag = _node_to_gbnf(child)
|
|
if isinstance(child, Alt):
|
|
frag = f'({frag})'
|
|
parts.append(frag)
|
|
return ' '.join(p for p in parts if p)
|
|
if isinstance(node, Alt):
|
|
parts = [_node_to_gbnf(child) for child in node.parts]
|
|
return ' | '.join(p for p in parts if p)
|
|
if isinstance(node, Plus):
|
|
frag = _node_to_gbnf(node.child)
|
|
if _needs_group(node.child):
|
|
return f'({frag})+'
|
|
return f'{frag}+'
|
|
if isinstance(node, Optional):
|
|
frag = _node_to_gbnf(node.child)
|
|
if _needs_group(node.child):
|
|
return f'({frag})?'
|
|
return f'{frag}?'
|
|
if isinstance(node, Star):
|
|
frag = _node_to_gbnf(node.child)
|
|
if _needs_group(node.child):
|
|
return f'({frag})*'
|
|
return f'{frag}*'
|
|
return ''
|
|
|
|
|
|
def to_gbnf(node):
|
|
"""Convert AST node to a GBNF grammar string.
|
|
|
|
Args:
|
|
node: Grammar AST node
|
|
|
|
Returns:
|
|
GBNF grammar string with a single 'root' rule.
|
|
"""
|
|
if node is None or isinstance(node, Empty):
|
|
return 'root ::= ""'
|
|
if isinstance(node, Epsilon):
|
|
return 'root ::= ""'
|
|
frag = _node_to_gbnf(node)
|
|
return f'root ::= {frag}'
|
|
|
|
|
|
def to_gbnf_with_rules(node, name='root'):
|
|
"""Convert AST to GBNF with a named rule."""
|
|
if node is None or isinstance(node, Empty):
|
|
return f'{name} ::= ""'
|
|
if isinstance(node, Epsilon):
|
|
return f'{name} ::= ""'
|
|
frag = _node_to_gbnf(node)
|
|
return f'{name} ::= {frag}'
|
|
|
|
|
|
def grammar_structure_score(node):
|
|
"""Quantify how structured an AST is (0.0 = flat bag, 1.0 = fully structured)."""
|
|
if node is None or isinstance(node, Empty):
|
|
return 0.0
|
|
if isinstance(node, Symbol):
|
|
return 0.0
|
|
if isinstance(node, Epsilon):
|
|
return 0.0
|
|
if isinstance(node, (Plus, Optional, Star)):
|
|
return 0.5 + 0.5 * grammar_structure_score(node.child)
|
|
if isinstance(node, Alt):
|
|
child_scores = [grammar_structure_score(p) for p in node.parts]
|
|
return sum(child_scores) / max(len(child_scores), 1)
|
|
if isinstance(node, Concat):
|
|
n = len(node.parts)
|
|
if n <= 1:
|
|
return 1.0
|
|
child_scores = [grammar_structure_score(p) for p in node.parts]
|
|
return min(1.0, 0.3 + 0.2 * n + sum(child_scores) / max(len(child_scores), 1))
|
|
return 0.0
|