#!/usr/bin/env -S uv run --script

# /// script

# requires-python = ">=3.8"

# dependencies = [

# "fastapi",

# "jinja2",

# "llama-cpp-python",

# "numpy",

# "openai",

# "pydantic",

# "uvicorn",

# "websockets",

# ]

# ///

from __future__ import annotations

import abc

import os

import re

import json

import math

import time

import uuid

import queue

import ctypes

import fnmatch

import asyncio

import argparse

import threading

import multiprocessing

import copy

import shutil

from pathlib import Path

from datetime import datetime

from dataclasses import dataclass, field

from collections import OrderedDict, deque

from openai.types.completion import Completion as OpenAICompletion

from openai.types.completion_choice import (

CompletionChoice,

Logprobs as CompletionLogprobs,

)

from openai.types.completion_usage import CompletionUsage

from openai.types.chat.chat_completion import (

ChatCompletion,

Choice as ChatCompletionChoice,

ChoiceLogprobs as ChatCompletionChoiceLogprobs,

)

from openai.types.chat.chat_completion_chunk import (

ChatCompletionChunk,

Choice as ChatCompletionChunkChoice,

ChoiceDelta,

ChoiceDeltaFunctionCall,

ChoiceDeltaToolCallFunction,

ChoiceDeltaToolCall,

ChoiceLogprobs as ChatCompletionChunkChoiceLogprobs,

)

from openai.types.chat.chat_completion_message import (

ChatCompletionMessage,

FunctionCall as ChatCompletionMessageFunctionCall,

)

from openai.types.chat.chat_completion_message_tool_call import (

ChatCompletionMessageToolCall,

Function as ChatCompletionMessageToolCallFunction,

)

from openai.types.chat.chat_completion_token_logprob import (

ChatCompletionTokenLogprob,

TopLogprob,

)

from typing import (

Any,

Callable,

Dict,

Generator,

Iterable,

List,

Optional,

Sequence,

Tuple,

Union,

Deque,

Literal,

Iterator,

TypedDict,

cast,

AsyncIterator,

)

import jinja2

import uvicorn

import numpy as np

from fastapi import FastAPI, HTTPException, Request, WebSocket, WebSocketDisconnect

from fastapi.middleware.cors import CORSMiddleware

from fastapi.responses import JSONResponse, Response, StreamingResponse

from jinja2.sandbox import ImmutableSandboxedEnvironment

from pydantic_core import from_json

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

from llama_cpp import llama_cpp # noqa: E402

JSON_GBNF = r"""

root ::= object

value ::= object | array | string | number | ("true" | "false" | "null") ws

object ::=

"{" ws (

string ":" ws value

("," ws string ":" ws value)*

)? "}" ws

array ::=

"[" ws (

value

("," ws value)*

)? "]" ws

string ::=

"\"" (

[^"\\\x7F\x00-\x1F] |

"\\" (["\\bfnrt] | "u" [0-9a-fA-F]{4})

)* "\"" ws

number ::= ("-"? ([0-9] | [1-9] [0-9]{0,15})) ("." [0-9]+)? ([eE] [-+]? [0-9] [1-9]{0,15})? ws

ws ::= | " " | "\n" [ \t]{0,20}

"""

class JsonSchemaConverter:

class BuiltinRule:

def __init__(self, content: str, deps: Optional[List[str]] = None):

self.content = content

self.deps = deps or []

SPACE_RULE = '" "?'

INVALID_RULE_CHARS_RE = re.compile(r"[^a-zA-Z0-9-]+")

GRAMMAR_LITERAL_ESCAPE_RE = re.compile(r'[\r\n"]')

GRAMMAR_LITERAL_ESCAPES = {"\r": "\\r", "\n": "\\n", '"': '\\"'}

DOTALL = "[\\U00000000-\\U0010FFFF]"

DOT = "[^\\x0A\\x0D]"

NON_LITERAL_SET = set("|.()[]{}*+?")

ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS = set("[]()|{}*+?")

PRIMITIVE_RULES: Optional[Dict[str, "JsonSchemaConverter.BuiltinRule"]] = None

STRING_FORMAT_RULES: Optional[Dict[str, "JsonSchemaConverter.BuiltinRule"]] = None

RESERVED_NAMES: Optional[set[str]] = None

@staticmethod

def _build_repetition(

item_rule: str,

min_items: int,

max_items: Optional[int],

separator_rule: Optional[str] = None,

item_rule_is_literal: bool = False,

) -> str:

if not separator_rule:

if min_items == 0 and max_items == 1:

return f"{item_rule}?"

if min_items == 1 and max_items is None:

return f"{item_rule}+"

result = ""

if min_items > 0:

if item_rule_is_literal and separator_rule is None:

result = '"' + (item_rule[1:-1] * min_items) + '"'

else:

result = (f" {separator_rule} " if separator_rule else " ").join(

[item_rule] * min_items

)

def opt_repetitions(up_to_n: int, prefix_with_sep: bool = False) -> str:

content = (

f"{separator_rule} {item_rule}"

if prefix_with_sep and separator_rule

else item_rule

)

if up_to_n == 0:

return ""

if up_to_n == 1:

return f"({content})?"

if separator_rule and not prefix_with_sep:

return f"({content} {opt_repetitions(up_to_n - 1, prefix_with_sep=True)})?"

return (f"({content} " * up_to_n).rstrip() + (")?" * up_to_n)

if min_items > 0 and max_items != min_items:

result += " "

if max_items is not None:

result += opt_repetitions(max_items - min_items, prefix_with_sep=min_items > 0)

else:

item_operator = f"({separator_rule + ' ' if separator_rule else ''}{item_rule})"

if min_items == 0 and separator_rule:

result = f"({item_rule} {item_operator}*)?"

else:

result += f"{item_operator}*"

return result

@classmethod

def _primitive_rules(cls) -> Dict[str, "JsonSchemaConverter.BuiltinRule"]:

if cls.PRIMITIVE_RULES is None:

up_to_15_digits = cls._build_repetition("[0-9]", 0, 15)

cls.PRIMITIVE_RULES = {

"boolean": cls.BuiltinRule('("true" | "false") space', []),

"decimal-part": cls.BuiltinRule("[0-9] " + up_to_15_digits, []),

"integral-part": cls.BuiltinRule(

"[0-9] | [1-9] " + up_to_15_digits,

[],

),

"number": cls.BuiltinRule(

'("-"? integral-part) ("." decimal-part)? ([eE] [-+]? integral-part)? space',

["integral-part", "decimal-part"],

),

"integer": cls.BuiltinRule('("-"? integral-part) space', ["integral-part"]),

"value": cls.BuiltinRule(

"object | array | string | number | boolean | null",

["object", "array", "string", "number", "boolean", "null"],

),

"object": cls.BuiltinRule(

'"{" space ( string ":" space value ("," space string ":" space value)* )? "}" space',

["string", "value"],

),

"array": cls.BuiltinRule(

'"[" space ( value ("," space value)* )? "]" space',

["value"],

),

"uuid": cls.BuiltinRule(

r'"\"" '

+ ' "-" '.join("[0-9a-fA-F]" * n for n in [8, 4, 4, 4, 12])

+ r' "\"" space',

[],

),

"char": cls.BuiltinRule(

r'[^"\\] | "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F])',

[],

),

"string": cls.BuiltinRule(r'"\"" char* "\"" space', ["char"]),

"null": cls.BuiltinRule('"null" space', []),

}

return cls.PRIMITIVE_RULES

@classmethod

def _string_format_rules(cls) -> Dict[str, "JsonSchemaConverter.BuiltinRule"]:

if cls.STRING_FORMAT_RULES is None:

cls.STRING_FORMAT_RULES = {

"date": cls.BuiltinRule(

'[0-9] [0-9] [0-9] [0-9] "-" ( "0" [1-9] | "1" [0-2] ) "-" ( "0" [1-9] | [1-2] [0-9] | "3" [0-1] )',

[],

),

"time": cls.BuiltinRule(

'([01] [0-9] | "2" [0-3]) ":" [0-5] [0-9] ":" [0-5] [0-9] ( "." [0-9] [0-9] [0-9] )? ( "Z" | ( "+" | "-" ) ( [01] [0-9] | "2" [0-3] ) ":" [0-5] [0-9] )',

[],

),

"date-time": cls.BuiltinRule('date "T" time', ["date", "time"]),

"date-string": cls.BuiltinRule('"\\"" date "\\"" space', ["date"]),

"time-string": cls.BuiltinRule('"\\"" time "\\"" space', ["time"]),

"date-time-string": cls.BuiltinRule(

'"\\"" date-time "\\"" space',

["date-time"],

),

}

return cls.STRING_FORMAT_RULES

@classmethod

def _reserved_names(cls) -> set[str]:

if cls.RESERVED_NAMES is None:

cls.RESERVED_NAMES = set(

["root", "dot", *cls._primitive_rules().keys(), *cls._string_format_rules().keys()]

)

return cls.RESERVED_NAMES

def __init__(

self,

*,

prop_order: Dict[str, int],

allow_fetch: bool,

dotall: bool,

raw_pattern: bool,

):

self._prop_order = prop_order

self._allow_fetch = allow_fetch

self._dotall = dotall

self._raw_pattern = raw_pattern

self._rules: Dict[str, str] = {"space": self.SPACE_RULE}

self._refs: Dict[str, Any] = {}

self._refs_being_resolved: set[str] = set()

def _format_literal(self, literal: str) -> str:

escaped = self.GRAMMAR_LITERAL_ESCAPE_RE.sub(

lambda match: self.GRAMMAR_LITERAL_ESCAPES[match.group(0)],

literal,

)

return f'"{escaped}"'

def _add_rule(self, name: str, rule: str) -> str:

escaped_name = self.INVALID_RULE_CHARS_RE.sub("-", name)

if escaped_name not in self._rules or self._rules[escaped_name] == rule:

key = escaped_name

else:

suffix = 0

while (

f"{escaped_name}{suffix}" in self._rules

and self._rules[f"{escaped_name}{suffix}"] != rule

):

suffix += 1

key = f"{escaped_name}{suffix}"

self._rules[key] = rule

return key

def resolve_refs(self, schema: Dict[str, Any], url: str) -> Dict[str, Any]:

def visit(node: Any) -> Any:

if isinstance(node, list):

return [visit(child) for child in node]

if isinstance(node, dict):

ref = node.get("$ref")

if ref is not None and ref not in self._refs:

if ref.startswith("https://"):

raise ValueError("remote schema fetch is not allowed")

elif ref.startswith("#/"):

target = schema

ref = f"{url}{ref}"

node["$ref"] = ref

else:

raise ValueError(f"Unsupported ref {ref}")

for selector in ref.split("#")[-1].split("/")[1:]:

assert target is not None and selector in target, (

f"Error resolving ref {ref}: {selector} not in {target}"

)

target = target[selector]

self._refs[ref] = target

else:

for value in node.values():

visit(value)

return node

return cast(Dict[str, Any], visit(schema))

def _generate_union_rule(self, name: str, alt_schemas: List[Dict[str, Any]]) -> str:

return " | ".join(

self.visit(alt_schema, f"{name}{'-' if name else 'alternative-'}{index}")

for index, alt_schema in enumerate(alt_schemas)

)

def _visit_pattern(self, pattern: str, name: str) -> str:

assert pattern.startswith("^") and pattern.endswith("$"), (

'Pattern must start with "^" and end with "$"'

)

pattern = pattern[1:-1]

sub_rule_ids: Dict[str, str] = {}

index = 0

length = len(pattern)

def to_rule(item: Tuple[str, bool]) -> str:

text, is_literal = item

return f'"{text}"' if is_literal else text

def transform() -> Tuple[str, bool]:

nonlocal index

start = index

sequence: List[Tuple[str, bool]] = []

def get_dot() -> str:

rule = self.DOTALL if self._dotall else self.DOT

return self._add_rule("dot", rule)

def join_sequence() -> Tuple[str, bool]:

if len(sequence) == 1:

return sequence[0]

return (" ".join(to_rule(item) for item in sequence), False)

while index < length:

char = pattern[index]

if char == ".":

sequence.append((get_dot(), False))

index += 1

elif char == "(":

index += 1

if index < length:

assert pattern[index] != "?", (

f'Unsupported pattern syntax "{pattern[index]}" at index {index} of /{pattern}/'

)

sequence.append((f"({to_rule(transform())})", False))

elif char == ")":

index += 1

assert start > 0 and pattern[start - 1] == "(", (

f"Unbalanced parentheses; start = {start}, index = {index}, pattern = {pattern}"

)

return join_sequence()

elif char == "[":

square_brackets = char

index += 1

while index < length and pattern[index] != "]":

if pattern[index] == "\\":

square_brackets += pattern[index : index + 2]

index += 2

else:

square_brackets += pattern[index]

index += 1

assert index < length, (

f"Unbalanced square brackets; start = {start}, index = {index}, pattern = {pattern}"

)

square_brackets += "]"

index += 1

sequence.append((square_brackets, False))

elif char == "|":

sequence.append(("|", False))

index += 1

elif char in ("*", "+", "?"):

sequence[-1] = (to_rule(sequence[-1]) + char, False)

index += 1

elif char == "{":

curly_brackets = char

index += 1

while index < length and pattern[index] != "}":

curly_brackets += pattern[index]

index += 1

assert index < length, (

f"Unbalanced curly brackets; start = {start}, index = {index}, pattern = {pattern}"

)

curly_brackets += "}"

index += 1

numbers = [part.strip() for part in curly_brackets[1:-1].split(",")]

min_times = 0

max_times: Optional[int] = None

try:

if len(numbers) == 1:

min_times = int(numbers[0])

max_times = min_times

else:

assert len(numbers) == 2

min_times = int(numbers[0]) if numbers[0] else 0

max_times = int(numbers[1]) if numbers[1] else None

except ValueError as exc:

raise ValueError(

f"Invalid quantifier {curly_brackets} in /{pattern}/"

) from exc

sub, sub_is_literal = sequence[-1]

if not sub_is_literal:

rule_id = sub_rule_ids.get(sub)

if rule_id is None:

rule_id = self._add_rule(f"{name}-{len(sub_rule_ids) + 1}", sub)

sub_rule_ids[sub] = rule_id

sub = rule_id

sequence[-1] = (

self._build_repetition(

f'"{sub}"' if sub_is_literal else sub,

min_times,

max_times,

item_rule_is_literal=sub_is_literal,

),

False,

)

else:

literal = ""

while index < length:

if pattern[index] == "\\" and index < length - 1:

next_char = pattern[index + 1]

if next_char in self.ESCAPED_IN_REGEXPS_BUT_NOT_IN_LITERALS:

index += 1

literal += pattern[index]

index += 1

else:

literal += pattern[index : index + 2]

index += 2

elif pattern[index] == '"' and not self._raw_pattern:

literal += '\\"'

index += 1

elif pattern[index] not in self.NON_LITERAL_SET and (

index == length - 1

or literal == ""

or pattern[index + 1] == "."

or pattern[index + 1] not in self.NON_LITERAL_SET

):

literal += pattern[index]

index += 1

else:

break

if literal:

sequence.append((literal, True))

return join_sequence()

return self._add_rule(

name,

(

to_rule(transform())

if self._raw_pattern

else '"\\"" ' + to_rule(transform()) + ' "\\"" space'

),

)

def _resolve_ref(self, ref: str) -> str:

ref_name = ref.split("/")[-1]

if ref_name not in self._rules and ref not in self._refs_being_resolved:

self._refs_being_resolved.add(ref)

resolved = self._refs[ref]

ref_name = self.visit(resolved, ref_name)

self._refs_being_resolved.remove(ref)

return ref_name

def _generate_constant_rule(self, value: Any) -> str:

return self._format_literal(json.dumps(value))

def visit(self, schema: Dict[str, Any], name: str) -> str:

schema_type = schema.get("type")

schema_format = schema.get("format")

rule_name = name + "-" if name in self._reserved_names() else name or "root"

ref = schema.get("$ref")

if ref is not None:

return self._add_rule(rule_name, self._resolve_ref(ref))

if "oneOf" in schema or "anyOf" in schema:

return self._add_rule(

rule_name,

self._generate_union_rule(name, cast(List[Dict[str, Any]], schema.get("oneOf") or schema["anyOf"])),

)

if isinstance(schema_type, list):

return self._add_rule(

rule_name,

self._generate_union_rule(name, [{"type": entry} for entry in schema_type]),

)

if "const" in schema:

return self._add_rule(rule_name, self._generate_constant_rule(schema["const"]))

if "enum" in schema:

rule = " | ".join(self._generate_constant_rule(value) for value in schema["enum"])

return self._add_rule(rule_name, rule)

if schema_type in (None, "object") and (

"properties" in schema

or ("additionalProperties" in schema and schema["additionalProperties"] is not True)

):

required_props = set(schema.get("required", []))

property_items = list(cast(Dict[str, Any], schema.get("properties", {})).items())

return self._add_rule(

rule_name,

self._build_object_rule(

property_items, required_props, name, schema.get("additionalProperties")

),

)

if schema_type in (None, "object") and "allOf" in schema:

allof_required_props: set[str] = set()

allof_property_items: List[Tuple[str, Any]] = []

def add_component(component_schema: Dict[str, Any], is_required: bool) -> None:

component_ref = component_schema.get("$ref")

if component_ref is not None:

component_schema = cast(Dict[str, Any], self._refs[component_ref])

if "properties" in component_schema:

for prop_name, prop_schema in cast(Dict[str, Any], component_schema["properties"]).items():

allof_property_items.append((prop_name, prop_schema))

if is_required:

allof_required_props.add(prop_name)

for entry in cast(List[Dict[str, Any]], schema["allOf"]):

if "anyOf" in entry:

for alt in cast(List[Dict[str, Any]], entry["anyOf"]):

add_component(alt, is_required=False)

else:

add_component(entry, is_required=True)

return self._add_rule(

rule_name,

self._build_object_rule(

allof_property_items,

allof_required_props,

name,

additional_properties=[],

),

)

if schema_type in (None, "array") and ("items" in schema or "prefixItems" in schema):

items = schema.get("items") or schema["prefixItems"]

if isinstance(items, list):

return self._add_rule(

rule_name,

'"[" space '

+ ' "," space '.join(

self.visit(item, f"{name}{'-' if name else ''}tuple-{index}")

for index, item in enumerate(items)

)

+ ' "]" space',

)

item_rule_name = self.visit(cast(Dict[str, Any], items), f"{name}{'-' if name else ''}item")

min_items = int(schema.get("minItems", 0))

max_items = cast(Optional[int], schema.get("maxItems"))

return self._add_rule(

rule_name,

'"[" space '

+ self._build_repetition(

item_rule_name,

min_items,

max_items,

separator_rule='"," space',

)

+ ' "]" space',

)

if schema_type in (None, "string") and "pattern" in schema:

return self._visit_pattern(cast(str, schema["pattern"]), rule_name)

if schema_type in (None, "string") and re.match(r"^uuid[1-5]?$", schema_format or ""):

return self._add_primitive(

"root" if rule_name == "root" else cast(str, schema_format),

self._primitive_rules()["uuid"],

)

if schema_type in (None, "string") and f"{schema_format}-string" in self._string_format_rules():

primitive_name = f"{schema_format}-string"

return self._add_rule(

rule_name,

self._add_primitive(primitive_name, self._string_format_rules()[primitive_name]),

)

if schema_type == "string" and ("minLength" in schema or "maxLength" in schema):

char_rule = self._add_primitive("char", self._primitive_rules()["char"])

min_len = int(schema.get("minLength", 0))

max_len = cast(Optional[int], schema.get("maxLength"))

return self._add_rule(

rule_name,

r'"\"" '

+ self._build_repetition(char_rule, min_len, max_len)

+ r' "\"" space',

)

if schema_type == "object" or len(schema) == 0:

return self._add_rule(

rule_name,

self._add_primitive("object", self._primitive_rules()["object"]),

)

primitive_rules = self._primitive_rules()

assert schema_type in primitive_rules, f"Unrecognized schema: {schema}"

return self._add_primitive(

"root" if rule_name == "root" else cast(str, schema_type),

primitive_rules[cast(str, schema_type)],

)

def _add_primitive(self, name: str, rule: "JsonSchemaConverter.BuiltinRule") -> str:

rule_name = self._add_rule(name, rule.content)

primitive_rules = self._primitive_rules()

string_format_rules = self._string_format_rules()

for dependency in rule.deps:

dependency_rule = primitive_rules.get(dependency) or string_format_rules.get(

dependency

)

assert dependency_rule is not None, f"Rule {dependency} not known"

if dependency not in self._rules:

self._add_primitive(dependency, dependency_rule)

return rule_name

def _build_object_rule(

self,

properties: List[Tuple[str, Any]],

required: set[str],

name: str,

additional_properties: Union[bool, Any],

) -> str:

prop_order = self._prop_order

sorted_props = [

key

for _, (key, _) in sorted(

enumerate(properties),

key=lambda indexed_key: (

prop_order.get(indexed_key[1][0], len(prop_order)),

indexed_key[0],

),

)

]

property_kv_rule_names: Dict[str, str] = {}

for prop_name, prop_schema in properties:

prop_rule_name = self.visit(

cast(Dict[str, Any], prop_schema),

f"{name}{'-' if name else ''}{prop_name}",

)

property_kv_rule_names[prop_name] = self._add_rule(

f"{name}{'-' if name else ''}{prop_name}-kv",

rf'{self._format_literal(json.dumps(prop_name))} space ":" space {prop_rule_name}',

)

required_props = [key for key in sorted_props if key in required]

optional_props = [key for key in sorted_props if key not in required]

if additional_properties is True or isinstance(additional_properties, dict):

sub_name = f"{name}{'-' if name else ''}additional"

value_rule = self.visit(

{} if additional_properties is True else cast(Dict[str, Any], additional_properties),

f"{sub_name}-value",

)

property_kv_rule_names["*"] = self._add_rule(

f"{sub_name}-kv",

self._add_primitive("string", self._primitive_rules()["string"])

+ f' ":" space {value_rule}',

)

optional_props.append("*")

rule = '"{" space '

rule += ' "," space '.join(property_kv_rule_names[key] for key in required_props)

if optional_props:

if required_props:

rule += ' ( "," space ( '

else:

rule += "( "

def get_recursive_refs(keys: List[str], first_is_optional: bool) -> str:

head, *rest = keys

kv_rule_name = property_kv_rule_names[head]

result = ""

if head == "*":

if first_is_optional:

result = f"({kv_rule_name})?"

else:

result = kv_rule_name

elif first_is_optional:

result = f'( "," space {kv_rule_name} )?'

else:

result = kv_rule_name

if rest:

result += " " + self._add_rule(

f"{name}{'-' if name else ''}{head}-rest",

get_recursive_refs(rest, first_is_optional=True),

)

return result

rule += " | ".join(

get_recursive_refs(optional_props[index:], first_is_optional=False)

for index in range(len(optional_props))

)

if required_props:

rule += " )"

rule += " )?"

rule += ' "}" space'

return rule

def format_grammar(self) -> str:

return "\n".join(

f"{name} ::= {rule}"

for name, rule in sorted(self._rules.items(), key=lambda item: item[0])

)

@classmethod

def to_gbnf(cls, schema: str, prop_order: Optional[List[str]] = None) -> str:

property_order = prop_order or []

loaded_schema = json.loads(schema)

order_index = {name: index for index, name in enumerate(property_order)}

converter = cls(

prop_order=order_index,

allow_fetch=False,

dotall=False,

raw_pattern=False,

)

resolved_schema = converter.resolve_refs(loaded_schema, "stdin")

converter.visit(resolved_schema, "")

return converter.format_grammar()

class RadixTrie:

__slots__ = ("root", "sequences", "sequence_lengths")

@dataclass

class Node:

label: Tuple[int, ...] = ()

parent: Optional["RadixTrie.Node"] = None

children: Dict[int, "RadixTrie.Node"] = field(default_factory=dict)

sequences: set[int] = field(default_factory=set)

tails: set[int] = field(default_factory=set)

def __init__(self) -> None:

self.root = RadixTrie.Node()

self.sequences: Dict[int, RadixTrie.Node] = {}

self.sequence_lengths: Dict[int, int] = {}

@staticmethod

def _pick_sequence(candidates: set[int], preferred_sequences: Optional[Any]) -> int:

if preferred_sequences is None:

return next(iter(candidates))

if isinstance(preferred_sequences, OrderedDict):

for sequence_id in reversed(preferred_sequences):

if sequence_id in candidates:

return sequence_id

return next(iter(candidates))

if isinstance(preferred_sequences, (list, tuple)):

for sequence_id in reversed(preferred_sequences):

if sequence_id in candidates:

return sequence_id

return next(iter(candidates))

preferred = candidates & preferred_sequences

if preferred:

return next(iter(preferred))

return next(iter(candidates))

@staticmethod

def _common_prefix_len(

label: Sequence[int],

tokens: Sequence[int],

offset: int,

) -> int:

limit = min(len(label), len(tokens) - offset)

match_len = 0

while match_len < limit and label[match_len] == tokens[offset + match_len]:

match_len += 1

return match_len

def _split_child(

self,

parent: "RadixTrie.Node",

child: "RadixTrie.Node",

prefix_len: int,

) -> "RadixTrie.Node":

assert 0 < prefix_len < len(child.label)

prefix = child.label[:prefix_len]

suffix = child.label[prefix_len:]

middle = RadixTrie.Node(

label=prefix,

parent=parent,

sequences=set(child.sequences),

)

parent.children[prefix[0]] = middle

child.label = suffix

child.parent = middle

middle.children[suffix[0]] = child

return middle

def _locate_prefix_node(

self,

sequence_id: int,

keep_len: int,

) -> "RadixTrie.Node":

total_len = self.sequence_lengths[sequence_id]

assert 0 <= keep_len <= total_len

if keep_len == 0:

return self.root

node = self.sequences[sequence_id]

drop_len = total_len - keep_len

while node is not self.root:

label_len = len(node.label)

if drop_len == 0:

return node

if drop_len < label_len:

parent = node.parent

assert parent is not None

return self._split_child(parent, node, label_len - drop_len)

drop_len -= label_len

parent = node.parent

assert parent is not None

node = parent

return self.root

def extend(

self,

sequence_id: int,

tokens: Sequence[int],

) -> None:

assert sequence_id >= 0

node = self.sequences.get(sequence_id, self.root)

if tokens:

node.tails.discard(sequence_id)

length = self.sequence_lengths.get(sequence_id, 0)

index = 0

while index < len(tokens):

token = tokens[index]

child = node.children.get(token)

if child is None:

child = RadixTrie.Node(

label=tuple(tokens[index:]),

parent=node,

sequences={sequence_id},

)

node.children[token] = child

node = child

length += len(tokens) - index

index = len(tokens)

break

match_len = self._common_prefix_len(child.label, tokens, index)

if match_len == len(child.label):

child.sequences.add(sequence_id)

node = child

length += match_len

index += match_len

continue

node = self._split_child(node, child, match_len)

node.sequences.add(sequence_id)

length += match_len

index += match_len

if index == len(tokens):

break

suffix = RadixTrie.Node(

label=tuple(tokens[index:]),

parent=node,

sequences={sequence_id},

)

node.children[suffix.label[0]] = suffix

node = suffix

length += len(tokens) - index

index = len(tokens)

break

if node is self.root:

self.sequences.pop(sequence_id, None)

self.sequence_lengths.pop(sequence_id, None)

self.root.tails.discard(sequence_id)

else:

self.sequences[sequence_id] = node

self.sequence_lengths[sequence_id] = length

node.tails.add(sequence_id)

def length(self, sequence_id: int) -> int:

return self.sequence_lengths.get(sequence_id, 0)

def truncate(self, sequence_id: int, keep_len: int) -> None:

assert sequence_id >= 0

assert sequence_id in self.sequence_lengths

assert 0 <= keep_len <= self.sequence_lengths[sequence_id]

current_len = self.sequence_lengths[sequence_id]

if keep_len == current_len:

return

boundary = self._locate_prefix_node(sequence_id, keep_len)

node = self.sequences.get(sequence_id, self.root)

if node is not self.root:

node.tails.discard(sequence_id)

while node is not boundary:

node.sequences.remove(sequence_id)

parent = node.parent

assert parent is not None

if not node.sequences:

del parent.children[node.label[0]]

node = parent

if boundary is self.root:

self.sequences.pop(sequence_id, None)

self.sequence_lengths.pop(sequence_id, None)

else:

self.sequences[sequence_id] = boundary

self.sequence_lengths[sequence_id] = keep_len

boundary.tails.add(sequence_id)

def copy(self, source_sequence_id: int, dest_sequence_id: int, keep_len: int) -> None:

assert source_sequence_id >= 0

assert dest_sequence_id >= 0

assert source_sequence_id in self.sequence_lengths

assert dest_sequence_id not in self.sequence_lengths

assert 0 <= keep_len <= self.sequence_lengths[source_sequence_id]

if keep_len == 0:

self.sequences[dest_sequence_id] = self.root

self.sequence_lengths[dest_sequence_id] = 0

self.root.tails.add(dest_sequence_id)

return

node = self._locate_prefix_node(source_sequence_id, keep_len)

self.sequences[dest_sequence_id] = node

self.sequence_lengths[dest_sequence_id] = keep_len

node.tails.add(dest_sequence_id)

while node is not self.root:

node.sequences.add(dest_sequence_id)

parent = node.parent

assert parent is not None

node = parent

def tokens(self, sequence_id: int, keep_len: Optional[int] = None) -> List[int]:

length = self.sequence_lengths[sequence_id]

target_len = length if keep_len is None else keep_len

assert 0 <= target_len <= length

node = self.sequences[sequence_id]

labels: List[Tuple[int, ...]] = []

while node is not self.root:

labels.append(node.label)

parent = node.parent

assert parent is not None

node = parent

values: List[int] = []

for label in reversed(labels):

values.extend(label)

return values[:target_len]

def longest_prefix(

self,

tokens: Sequence[int],

preferred_sequences: Optional[Any] = None,

*,

exact_only: bool = False,

) -> Tuple[int, int]:

node = self.root

longest_sequence_id = -1

longest_length = 0

index = 0

while index < len(tokens):

child = node.children.get(tokens[index])