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

Fix type warnings for json schema grammar converter

Browse source
This commit is contained in:
Andrei Betlen 2023-11-21 13:32:00 -05:00
parent 230fc8b535
commit 8b6ca22846
1 changed files with 9 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

15
llama_cpp/llama_grammar.py
View file

@ -7,7 +7,9 @@ from ctypes import * # type: ignore
from enum import Enum from enum import Enum
from itertools import islice from itertools import islice
from typing import ( from typing import (
Any,
Callable, Callable,
Dict,
Generic, Generic,
List, List,
Optional, Optional,
@ -1399,15 +1401,15 @@ class SchemaConverter:
def __init__(self, prop_order): def __init__(self, prop_order):
self._prop_order = prop_order self._prop_order = prop_order
self._rules = {"space": SPACE_RULE} self._rules = {"space": SPACE_RULE}
self._defs = {} self._defs: Dict[str, Any] = {}
def _format_literal(self, literal): def _format_literal(self, literal: str):
escaped = GRAMMAR_LITERAL_ESCAPE_RE.sub( escaped: str = GRAMMAR_LITERAL_ESCAPE_RE.sub(
lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)), json.dumps(literal) lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)), json.dumps(literal)
) )
return f'"{escaped}"' return f'"{escaped}"'
def _add_rule(self, name, rule): def _add_rule(self, name: str, rule: str):
esc_name = INVALID_RULE_CHARS_RE.sub("-", name) esc_name = INVALID_RULE_CHARS_RE.sub("-", name)
if esc_name not in self._rules or self._rules[esc_name] == rule: if esc_name not in self._rules or self._rules[esc_name] == rule:
key = esc_name key = esc_name
@ -1419,8 +1421,9 @@ class SchemaConverter:
self._rules[key] = rule self._rules[key] = rule
return key return key
def visit(self, schema, name): def visit(self, schema: Dict[str, Any], name: str) -> str:
schema_type = schema.get("type") schema_type: Optional[str] = schema.get("type") # type: ignore
assert isinstance(schema_type, str), f"Unrecognized schema: {schema}"
rule_name = name or "root" rule_name = name or "root"
if "$defs" in schema: if "$defs" in schema: