baalajimaestro/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add $ref and $defs support to json schema converter

Browse source
This commit is contained in:
Andrei Betlen 2023-11-10 02:50:46 -05:00
parent faeae181b1
commit 770df34436
1 changed files with 14 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
llama_cpp/llama_grammar.py
View file

@ -1399,6 +1399,7 @@ class SchemaConverter:
def __init__(self, prop_order):
self._prop_order = prop_order
self._rules = {"space": SPACE_RULE}
self._defs = {}
def _format_literal(self, literal):
escaped = GRAMMAR_LITERAL_ESCAPE_RE.sub(
@ -1422,6 +1423,11 @@ class SchemaConverter:
schema_type = schema.get("type")
rule_name = name or "root"
if "$defs" in schema:
# add defs to self._defs for later inlining
for def_name, def_schema in schema["$defs"].items():
self._defs[def_name] = def_schema
if "oneOf" in schema or "anyOf" in schema:
rule = " | ".join(
(
@ -1471,6 +1477,14 @@ class SchemaConverter:
)
return self._add_rule(rule_name, rule)
elif "$ref" in schema:
ref = schema["$ref"]
assert ref.startswith("#/$defs/"), f"Unrecognized schema: {schema}"
# inline $defs
def_name = ref[len("#/$defs/") :]
def_schema = self._defs[def_name]
return self.visit(def_schema, f'{name}{"-" if name else ""}{def_name}')
else:
assert schema_type in PRIMITIVE_RULES, f"Unrecognized schema: {schema}"
return self._add_rule(