BDav24 · BDav24 · Feb 6, 2024 · Feb 8, 2024 · Feb 8, 2024 · Feb 8, 2024
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -2,7 +2,6 @@
 
 import os
 import sys
-import abc
 import uuid
 import time
 import multiprocessing
@@ -15,7 +14,6 @@
     Iterator,
     Deque,
     Callable,
-    Any,
 )
 from collections import deque
 
@@ -31,6 +29,10 @@
     LlamaDiskCache,  # type: ignore
     LlamaRAMCache,  # type: ignore
 )
+from .llama_tokenizer import (
+    BaseLlamaTokenizer,
+    LlamaTokenizer
+)
 import llama_cpp.llama_cpp as llama_cpp
 import llama_cpp.llama_chat_format as llama_chat_format
 
@@ -948,8 +950,7 @@ def logit_bias_processor(
 
             if stream:
                 remaining_tokens = completion_tokens[returned_tokens:]
-                prev_tokens = completion_tokens[:returned_tokens]
-                remaining_text = self.detokenize(completion_tokens, prev_tokens)
+                remaining_text = self.detokenize(remaining_tokens)
                 remaining_length = len(remaining_text)
 
                 # We want to avoid yielding any characters from
@@ -971,13 +972,13 @@ def logit_bias_processor(
                     for token in remaining_tokens:
                         if token == self.token_bos():
                             continue
-                        token_end_position += len(remaining_text)
+                        token_end_position += len(self.detokenize([token]))
                         # Check if stop sequence is in the token
                         if token_end_position > (
                             remaining_length - first_stop_position
                         ):
                             break
-                        token_str = remaining_text.decode(
+                        token_str = self.detokenize([token]).decode(
                             "utf-8", errors="ignore"
                         )
                         text_offset = len(prompt) + len(
@@ -1002,7 +1003,11 @@ def logit_bias_processor(
                         }
                         top_logprob.update({token_str: current_logprobs[int(token)]})
                         logprobs_or_none = {
-                            "tokens": [token_str],
+                            "tokens": [
+                                self.detokenize([token]).decode(
+                                    "utf-8", errors="ignore"
+                                )
+                            ],
                             "text_offset": [text_offset],
                             "token_logprobs": [current_logprobs[int(token)]],
                             "top_logprobs": [top_logprob],
@@ -1015,7 +1020,9 @@ def logit_bias_processor(
                             "model": model_name,
                             "choices": [
                                 {
-                                    "text": token_str,
+                                    "text": self.detokenize([token]).decode(
+                                        "utf-8", errors="ignore"
+                                    ),
                                     "index": 0,
                                     "logprobs": logprobs_or_none,
                                     "finish_reason": None,
@@ -1027,7 +1034,7 @@ def logit_bias_processor(
                         decode_success = False
                         for i in range(1, len(remaining_tokens) + 1):
                             try:
-                                bs = remaining_text
+                                bs = self.detokenize(remaining_tokens[:i])
                                 ts = bs.decode("utf-8")
                                 decode_success = True
                                 break
@@ -1063,7 +1070,6 @@ def logit_bias_processor(
 
             if len(completion_tokens) >= max_tokens:
                 text = self.detokenize(completion_tokens)
-
                 finish_reason = "length"
                 break
 
@@ -1747,69 +1753,6 @@ def longest_token_prefix(a: Sequence[int], b: Sequence[int]):
         return longest_prefix
 
 
-class BaseLlamaTokenizer(abc.ABC):
-    @abc.abstractmethod
-    def tokenize(self, text: bytes, add_bos: bool = True, special: bool = True) -> List[int]:
-        raise NotImplementedError
-
-    @abc.abstractmethod
-    def detokenize(self, tokens: List[int], prev_tokens: Optional[List[int]] = None) -> bytes:
-        raise NotImplementedError
-
-
-class LlamaTokenizer(BaseLlamaTokenizer):
-    def __init__(self, llama: Llama):
-        self.llama = llama
-        self._model = llama._model # type: ignore
-
-    def tokenize(self, text: bytes, add_bos: bool = True, special: bool = True) -> List[int]:
-        return self._model.tokenize(text, add_bos=add_bos, special=special)
-
-    def detokenize(self, tokens: List[int], prev_tokens: Optional[List[int]] = None) -> bytes:
-        if prev_tokens is not None:
-            return self._model.detokenize(tokens[len(prev_tokens):])
-        else:
-            return self._model.detokenize(tokens)
-
-    def encode(self, text: str, add_bos: bool = True, special: bool = True) -> List[int]:
-        return self.tokenize(
-            text.encode("utf-8", errors="ignore"), add_bos=add_bos, special=special
-        )
-
-    def decode(self, tokens: List[int]) -> str:
-        return self.detokenize(tokens).decode("utf-8", errors="ignore")
-
-    @classmethod
-    def from_ggml_file(cls, path: str) -> "LlamaTokenizer":
-        return cls(Llama(model_path=path, vocab_only=True))
-
-
-class LlamaHFTokenizer(BaseLlamaTokenizer):
-    def __init__(self, hf_tokenizer: Any):
-        self.hf_tokenizer = hf_tokenizer
-
-    def tokenize(self, text: bytes, add_bos: bool = True, special: bool = True) -> List[int]:
-        return self.hf_tokenizer.encode(text.decode("utf-8", errors="ignore"), add_special_tokens=special)
-
-    def detokenize(self, tokens: List[int], prev_tokens: Optional[List[int]] = None) -> bytes:
-        if prev_tokens is not None:
-            text = self.hf_tokenizer.decode(tokens).encode("utf-8", errors="ignore")
-            prev_text = self.hf_tokenizer.decode(prev_tokens).encode("utf-8", errors="ignore")
-            return text[len(prev_text):]
-        else:
-            return self.hf_tokenizer.decode(tokens).encode("utf-8", errors="ignore")
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: str) -> "LlamaHFTokenizer":
-        try:
-            from transformers import AutoTokenizer
-        except ImportError:
-            raise ImportError(
-                "The `transformers` library is required to use the `HFTokenizer`."
-                "You can install it with `pip install transformers`."
-            )
-        hf_tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path=pretrained_model_name_or_path)
-        return cls(hf_tokenizer)
 
 
 class LlamaState:

diff --git a/llama_cpp/llama_grammar.py b/llama_cpp/llama_grammar.py
@@ -81,9 +81,15 @@ def from_json_schema(
         cls,
         json_schema: str,
         verbose: bool = True,
+        treat_optional_as_nullable: bool = False,
     ) -> "LlamaGrammar":
         """Convert a JSON schema to a Llama grammar."""
-        return cls.from_string(json_schema_to_gbnf(json_schema), verbose=verbose)
+        return cls.from_string(
+            json_schema_to_gbnf(
+                json_schema, treat_optional_as_nullable=treat_optional_as_nullable
+            ),
+            verbose=verbose,
+        )
 
     @classmethod
     def from_file(cls, file: Union[str, Path], verbose: bool = True) -> "LlamaGrammar":
@@ -1392,14 +1398,14 @@ def print_grammar(file: TextIO, state: parse_state) -> None:
 SPACE_RULE = '" "?'
 
 PRIMITIVE_RULES = {
-    "boolean": '("true" | "false") space',
-    "number": '("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? space',
-    "integer": '("-"? ([0-9] | [1-9] [0-9]*)) space',
+    "boolean": '("true" | "false")',
+    "number": '("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)?',
+    "integer": '("-"? ([0-9] | [1-9] [0-9]*))',
     "string": r""" "\"" (
         [^"\\] |
         "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])
-      )* "\"" space """,
-    "null": '"null" space',
+      )* "\"" """,
+    "null": '"null"',
 }
 
 INVALID_RULE_CHARS_RE = re.compile(r"[^a-zA-Z0-9-]+")
@@ -1408,30 +1414,43 @@ def print_grammar(file: TextIO, state: parse_state) -> None:
 
 
 class SchemaConverter:
-    def __init__(self, prop_order):
+    def __init__(self, prop_order, treat_optional_as_nullable: bool = False):
         self._prop_order = prop_order
         self._rules = {"space": SPACE_RULE}
         self._defs: Dict[str, Any] = {}
+        self._treat_optional_as_nullable = treat_optional_as_nullable
 
     def _format_literal(self, literal: str):
         escaped: str = GRAMMAR_LITERAL_ESCAPE_RE.sub(
             lambda m: GRAMMAR_LITERAL_ESCAPES.get(m.group(0)), json.dumps(literal)
         )
         return f'"{escaped}"'
 
-    def _add_rule(self, name: str, rule: str):
+    def _add_rule(
+        self, name: str, rule: str, is_required: bool = True, with_space: bool = False
+    ):
         esc_name = INVALID_RULE_CHARS_RE.sub("-", name)
-        if esc_name not in self._rules or self._rules[esc_name] == rule:
+        complete_rule = rule
+
+        if self._treat_optional_as_nullable and not is_required:
+            esc_name += "-or-null"
+            complete_rule = f"({complete_rule} | {PRIMITIVE_RULES['null']})"
+
+        if with_space:
+            complete_rule += " space"
+
+        if esc_name not in self._rules or self._rules[esc_name] == complete_rule:
             key = esc_name
         else:
             i = 0
             while f"{esc_name}{i}" in self._rules:
                 i += 1
             key = f"{esc_name}{i}"
-        self._rules[key] = rule
+
+        self._rules[key] = complete_rule
         return key
 
-    def visit(self, schema: Dict[str, Any], name: str) -> str:
+    def visit(self, schema: Dict[str, Any], name: str, is_required: bool = True) -> str:
         rule_name = name or "root"
 
         if "$defs" in schema:
@@ -1448,14 +1467,16 @@ def visit(self, schema: Dict[str, Any], name: str) -> str:
                     )
                 )
             )
-            return self._add_rule(rule_name, rule)
+            return self._add_rule(rule_name, rule, is_required, False)
 
         elif "const" in schema:
-            return self._add_rule(rule_name, self._format_literal(schema["const"]))
+            return self._add_rule(
+                rule_name, self._format_literal(schema["const"]), is_required, False
+            )
 
         elif "enum" in schema:
             rule = " | ".join((self._format_literal(v) for v in schema["enum"]))
-            return self._add_rule(rule_name, rule)
+            return self._add_rule(rule_name, rule, is_required, False)
 
         elif "$ref" in schema:
             ref = schema["$ref"]
@@ -1465,56 +1486,78 @@ def visit(self, schema: Dict[str, Any], name: str) -> str:
             def_schema = self._defs[def_name]
             return self.visit(def_schema, f'{name}{"-" if name else ""}{def_name}')
 
-
-        schema_type: Optional[str] = schema.get("type") # type: ignore
+        schema_type: Optional[str] = schema.get("type")  # type: ignore
         assert isinstance(schema_type, str), f"Unrecognized schema: {schema}"
 
         if schema_type == "object" and "properties" in schema:
-            # TODO: `required` keyword
             prop_order = self._prop_order
             prop_pairs = sorted(
                 schema["properties"].items(),
                 # sort by position in prop_order (if specified) then by key
                 key=lambda kv: (prop_order.get(kv[0], len(prop_order)), kv[0]),
             )
 
-            rule = '"{" space'
+            rule = ""
+            previous_is_prop_required = False
             for i, (prop_name, prop_schema) in enumerate(prop_pairs):
+                is_prop_required = (
+                    "required" not in schema or prop_name in schema["required"]
+                )
                 prop_rule_name = self.visit(
-                    prop_schema, f'{name}{"-" if name else ""}{prop_name}'
+                    prop_schema,
+                    f'{name}{"-" if name else ""}{prop_name}',
+                    is_prop_required,
                 )
-                if i > 0:
-                    rule += ' "," space'
-                rule += rf' {self._format_literal(prop_name)} space ":" space {prop_rule_name}'
-            rule += ' "}" space'
+                prop_rule = rf'{self._format_literal(prop_name)} space ":" space {prop_rule_name}'
+                if i == 0:
+                    rule += prop_rule
+                else:
+                    if self._treat_optional_as_nullable or (
+                        previous_is_prop_required and is_prop_required
+                    ):
+                        rule = f'{rule} "," space {prop_rule}'
+                    elif previous_is_prop_required and not is_prop_required:
+                        rule = f'{rule} ("," space {prop_rule})?'
+                    elif not previous_is_prop_required and is_prop_required:
+                        rule = f'({rule} "," space)? {prop_rule}'
+                    elif not previous_is_prop_required and not is_prop_required:
+                        rule = f'({rule} | {prop_rule} | {rule} "," space {prop_rule})'
 
-            return self._add_rule(rule_name, rule)
+                previous_is_prop_required |= is_prop_required
+
+            rule = '"{" space ' + rule + ' space "}"'
+
+            return self._add_rule(rule_name, rule, is_required, True)
 
         elif schema_type == "array" and "items" in schema:
             # TODO `prefixItems` keyword
             item_rule_name = self.visit(
                 schema["items"], f'{name}{"-" if name else ""}item'
             )
-            rule = (
-                f'"[" space ({item_rule_name} ("," space {item_rule_name})*)? "]" space'
-            )
-            return self._add_rule(rule_name, rule)
+            rule = f'"[" space ({item_rule_name} ("," space {item_rule_name})*)? "]"'
+            return self._add_rule(rule_name, rule, is_required, True)
 
         else:
             assert schema_type in PRIMITIVE_RULES, f"Unrecognized schema: {schema}"
             return self._add_rule(
                 "root" if rule_name == "root" else schema_type,
                 PRIMITIVE_RULES[schema_type],
+                is_required,
+                True,
             )
 
     def format_grammar(self):
         return "\n".join((f"{name} ::= {rule}" for name, rule in self._rules.items()))
 
 
-def json_schema_to_gbnf(schema: str, prop_order: Optional[List[str]] = None):
+def json_schema_to_gbnf(
+    schema: str,
+    prop_order: Optional[List[str]] = None,
+    treat_optional_as_nullable: bool = False,
+):
     prop_order = prop_order or []
     schema = json.loads(schema)
     prop_order = {name: idx for idx, name in enumerate(prop_order)}
-    converter = SchemaConverter(prop_order)
+    converter = SchemaConverter(prop_order, treat_optional_as_nullable)
     converter.visit(schema, "")
     return converter.format_grammar()