TOKEN={eos_token}")
95 | try:
96 | fail = acceptor.advance_token(eos_token)
97 | except acceptor.TokenRejected:
98 | fail = True
99 | if fail:
100 | print("[FAIL]")
101 | result = 1
102 | else:
103 | print("[SUCCESS]")
104 | result = 0
105 | if debug:
106 | debug("\n".join(repr(c) for c in acceptor.cursors))
107 | if args.paths:
108 | print(json.dumps(values_by_path, indent=2))
109 | return result
110 |
111 |
112 | if __name__ == "__main__":
113 | sys.exit(main())
114 |
--------------------------------------------------------------------------------
/src/examples/llm_schema.py:
--------------------------------------------------------------------------------
1 | # pylint: disable=missing-class-docstring,missing-function-docstring
2 | """
3 | Example of JSON schema decoding with MLX.
4 | """
5 | import argparse
6 | import json
7 | import time
8 | from math import inf
9 | from operator import itemgetter
10 | from typing import Iterable, Optional, Union
11 |
12 | import mlx.core as mx
13 | import mlx.nn as nn
14 |
15 | from mlx_lm.utils import load
16 |
17 | from llm_structured_output import JsonSchemaAcceptorDriver
18 | from llm_structured_output.util.bitmap import (
19 | bias_logits,
20 | count_set_bits,
21 | enumerate_set_bits,
22 | )
23 | from llm_structured_output.util.output import info, bold, bolddim, debug
24 | from llm_structured_output.util.tokenization import HuggingfaceTokenizerHelper
25 |
26 | from .reusable_kv_cache import ReusableKVCache
27 |
28 |
29 | class RejectedCompletion(Exception):
30 | """
31 | It's rare, but sometimes we reach a state where it's not possible to
32 | advance the acceptor. For example, when closing a JSON string we get
33 | a higher probability for slanted quotes than straight ones and select
34 | the wrong token. At that point, the LLM will continue generating with
35 | the prior that the string is closed, but our acceptor will remain in
36 | the string-accepting state. This can indicate an issue with the
37 | tokenizer vocabulary passed to the acceptor, or a bug in the code
38 | used to decode tokens from the LLM. If none of these apply, check that
39 | the LLM is actually able to generate JSON, although most are.
40 | """
41 |
42 |
43 | class Model:
44 | def __init__(self):
45 | mx.random.seed(0)
46 | self.model = None
47 | self.tokenizer = None
48 | self.vocabulary = None
49 | self.eos_id = None
50 | self.json_schema_acceptor_driver_factory = None
51 | self._cached_prompt = None
52 | self._cached_cache = None
53 |
54 | def load(self, model_path: str):
55 | """
56 | Load locally or download from Huggingface hub.
57 | """
58 | self.model, tokenizer = load(model_path)
59 | self.tokenizer = HuggingfaceTokenizerHelper(tokenizer)
60 | self.vocabulary, self.eos_id = self.tokenizer.extract_vocabulary()
61 | self.json_schema_acceptor_driver_factory = (
62 | JsonSchemaAcceptorDriver.driver_factory_for_model(
63 | self.vocabulary, self.eos_id
64 | )
65 | )
66 |
67 | def get_driver_for_json_schema(self, schema, encapsulated: bool = False):
68 | return self.json_schema_acceptor_driver_factory(
69 | schema, is_encapsulated_json=encapsulated
70 | )
71 |
72 | def _evaluate_prompt(
73 | self, prompt: list[int], prior_prompt: list[int] = None, prior_cache=None
74 | ):
75 | if prior_prompt:
76 | i = 0
77 | for i, t in enumerate(prior_prompt):
78 | # We need to leave at least one token to evaluate because we don't
79 | # save the past logits.
80 | if i >= len(prompt) - 1 or prompt[i] != t:
81 | break
82 | cache = prior_cache
83 | for layer_cache in cache:
84 | layer_cache.reuse(len(prompt), i)
85 | tokens = prompt[i:]
86 | else:
87 | cache = ReusableKVCache.for_model(self.model)
88 | tokens = prompt
89 |
90 | logits = self.model(mx.array(tokens)[None], cache=cache)
91 | return logits, cache
92 |
93 | def _decode(self, tokens):
94 | return self.tokenizer.no_strip_decode(tokens)
95 |
96 | def _debug_top_tokens(self, logits, count=10):
97 | token_logits = sorted(
98 | enumerate(logits.tolist()), key=itemgetter(1), reverse=True
99 | )
100 | top_tokens = [
101 | (self._decode([t]), p) for t, p in token_logits[:count] if p != -inf
102 | ]
103 | debug("TOP TOKENS:", top_tokens)
104 |
105 | def _sample(self, logits, temp: float = 0):
106 | if temp == 0:
107 | result = mx.argmax(logits, axis=-1)
108 | else:
109 | result = mx.random.categorical(logits * (1 / temp))
110 | return result.item()
111 |
112 | def _sample_with_bias(
113 | self, logits, temp: float = 0, token_acceptor=None, lazy_bias: bool = True
114 | ):
115 | if token_acceptor is None:
116 | return self._sample(logits, temp)
117 |
118 | if lazy_bias:
119 | token = self._sample(logits, temp)
120 | try:
121 | token_acceptor.advance_token(token)
122 | return token
123 | except JsonSchemaAcceptorDriver.TokenRejected:
124 | pass
125 |
126 | accepted_token_bitmap = token_acceptor.select_valid_tokens()
127 | if not accepted_token_bitmap:
128 | debug(token_acceptor.cursors)
129 | self._debug_top_tokens(logits)
130 | raise RejectedCompletion()
131 | token = self._sample(bias_logits(mx, logits, accepted_token_bitmap), temp)
132 | token_acceptor.advance_token(token)
133 | return token
134 |
135 | def generate_without_schema(self, logits, cache, temp: Optional[float] = 0.0):
136 | """
137 | For testing / comparison purposes.
138 | """
139 | while True:
140 | tokens = [self._sample(logits[0, -1, :], temp)]
141 | yield tokens
142 | if tokens[-1] == self.eos_id:
143 | break
144 | logits = self.model(mx.array(tokens)[None], cache=cache)
145 |
146 | def generate_with_schema(
147 | self, logits, cache, token_acceptor, temp: Optional[float] = 0.0
148 | ):
149 | while True:
150 | tokens = [self._sample_with_bias(logits[0, -1, :], temp, token_acceptor)]
151 | yield tokens
152 | if tokens[-1] == self.eos_id:
153 | break
154 | logits = self.model(mx.array(tokens)[None], cache=cache)
155 |
156 | def generate_with_preemptive_decoding(
157 | self,
158 | logits,
159 | cache,
160 | token_acceptor,
161 | temp: Optional[float] = 0.0,
162 | max_batch_size=5,
163 | ):
164 | """
165 | Try to generate faster by precomputing two tokens at a time when possible.
166 | If we know that the acceptor will only accept a small set of tokens after
167 | the current one, we can evaluate a batch with one entry per possible
168 | future token. Each entry in the batch contains the current token sampled,
169 | which we have to evaluate anyway, and a second token corresponding to one
170 | of the possible tokens that could be sampled from the output to the first
171 | token. We get back logits for both tokens for each item in the batch: the
172 | logits for the first token will be the same (as long as the model applies
173 | a causal mask), and we can sample those logits to select from which of the
174 | items in the batch we can select the second token.
175 | In practice, this only seems to accelerate things for unquantized models.
176 | """
177 | # Sample token from prompt evaluation
178 | accepted_token_bitmap = token_acceptor.select_valid_tokens()
179 | first_token_logits = bias_logits(mx, logits[0, -1, :], accepted_token_bitmap)
180 | first_token = self._sample(first_token_logits, temp)
181 | tokens = [first_token]
182 | yield tokens
183 | token_acceptor.advance_token(first_token)
184 | accepted_token_bitmap = token_acceptor.select_valid_tokens()
185 |
186 | while True:
187 | last_token = tokens[-1]
188 | if count_set_bits(accepted_token_bitmap) in range(1, max_batch_size + 1):
189 | # If the number of possible follow-up tokens is small, submit for
190 | # evaluation a batch of 2-token continuations.
191 | batch = []
192 | for followup_token in enumerate_set_bits(accepted_token_bitmap):
193 | batch.append([last_token, followup_token])
194 | # Re-shape the cache to match the input.
195 | for layer_cache in cache:
196 | layer_cache.keys = mx.concatenate([layer_cache.keys] * len(batch))
197 | layer_cache.values = mx.concatenate(
198 | [layer_cache.values] * len(batch)
199 | )
200 | else: # Otherwise, submit the normal one-token continuation.
201 | batch = [[last_token]]
202 |
203 | logits = self.model(mx.array(batch), cache=cache)
204 | mx.eval(logits)
205 |
206 | first_token_logits = bias_logits(mx, logits[0, 0, :], accepted_token_bitmap)
207 | first_token = self._sample(first_token_logits, temp)
208 | tokens = [first_token]
209 |
210 | if first_token == self.eos_id:
211 | yield tokens
212 | break
213 |
214 | token_acceptor.advance_token(first_token)
215 | accepted_token_bitmap = token_acceptor.select_valid_tokens()
216 | if not accepted_token_bitmap:
217 | raise RejectedCompletion()
218 |
219 | # If we had submitted 2-token continuations, we can decode a second token
220 | if len(batch[0]) > 1:
221 | index = next( # Find which of the second tokens was selected
222 | i
223 | for i, batch_item in enumerate(batch)
224 | if batch_item[1] == first_token
225 | )
226 | second_token_logits = bias_logits(
227 | mx, logits[index, 1, :], accepted_token_bitmap
228 | )
229 | second_token = self._sample(second_token_logits, temp)
230 | tokens.append(second_token)
231 |
232 | token_acceptor.advance_token(second_token)
233 | accepted_token_bitmap = token_acceptor.select_valid_tokens()
234 |
235 | # Select the accepted generation in the cache, restoring it to batch dimension 1.
236 | for layer_cache in cache:
237 | layer_cache.keys = layer_cache.keys.split([index, index + 1])[1]
238 | layer_cache.values = layer_cache.values.split([index, index + 1])[1]
239 |
240 | yield tokens
241 |
242 | def _generate_tokens(
243 | self,
244 | generator: Iterable,
245 | max_tokens: int = 1000,
246 | ) -> Iterable:
247 | start_time = time.time_ns()
248 | token_count = 0
249 |
250 | for tokens in generator:
251 | token_count += len(tokens)
252 |
253 | try:
254 | eos_index = tokens.index(self.eos_id)
255 | tokens = tokens[0:eos_index]
256 | except ValueError:
257 | eos_index = -1
258 |
259 | if tokens:
260 | text = self._decode(tokens)
261 | yield {
262 | "op": "generatedTokens",
263 | "text": text,
264 | "token_count": len(tokens),
265 | "time_ms": (time.time_ns() - start_time) / 1e6,
266 | }
267 |
268 | if eos_index >= 0:
269 | yield {"op": "stop", "reason": "end"}
270 | return
271 |
272 | if token_count >= max_tokens:
273 | yield {"op": "stop", "reason": "max_tokens"}
274 | return
275 |
276 | start_time = time.time_ns()
277 |
278 | assert False
279 |
280 | def completion(
281 | self,
282 | prompt: Union[str, Iterable[dict[str, str]]],
283 | schema: dict,
284 | encapsulated: bool = False,
285 | max_tokens: int = 1000,
286 | temp: float = 0.0,
287 | seed: int = None,
288 | preemptive_batch_size: int = 0,
289 | cache_prompt: bool = False,
290 | ):
291 | if seed is not None:
292 | mx.random.seed(seed)
293 |
294 | start_time = time.time_ns()
295 | prompt_tokens = self.tokenizer.encode_prompt(prompt)
296 | logits, cache = self._evaluate_prompt(
297 | prompt_tokens, self._cached_prompt, self._cached_cache
298 | )
299 | if cache_prompt:
300 | self._cached_prompt = prompt_tokens
301 | self._cached_cache = cache
302 | # Eager eval to more accurately reflect the prompt evaluation time.
303 | mx.eval(logits)
304 | prompt_time = time.time_ns() - start_time
305 | yield {
306 | "op": "evaluatedPrompt",
307 | "prompt": prompt,
308 | "token_count": len(prompt_tokens),
309 | "time_ms": prompt_time / 1e6,
310 | "prompt_tps": len(prompt_tokens) / (prompt_time / 1e9),
311 | }
312 |
313 | if schema:
314 | token_acceptor = self.get_driver_for_json_schema(schema, encapsulated)
315 | if preemptive_batch_size > 0:
316 | generator = self.generate_with_preemptive_decoding(
317 | logits,
318 | cache,
319 | token_acceptor,
320 | temp,
321 | max_batch_size=preemptive_batch_size,
322 | )
323 | else:
324 | generator = self.generate_with_schema(
325 | logits, cache, token_acceptor, temp
326 | )
327 | else:
328 | generator = self.generate_without_schema(logits, cache, temp)
329 |
330 | token_count = 0
331 | generation_time = 0
332 | for generation_result in self._generate_tokens(generator, max_tokens):
333 | if generation_result["op"] == "generatedTokens":
334 | token_count += generation_result["token_count"]
335 | generation_time += generation_result["time_ms"]
336 | elif generation_result["op"] == "stop":
337 | generation_result["token_count"] = token_count
338 | generation_result["time_ms"] = generation_time
339 | # This is slightly incorrect, because the first token is generated
340 | # from the prompt evaluation.
341 | generation_result["generation_tps"] = token_count / (
342 | generation_time / 1e3
343 | )
344 | yield generation_result
345 |
346 |
347 | def main():
348 | parser = argparse.ArgumentParser(
349 | description="LLM inference script with schema-constrained sampling"
350 | )
351 | parser.add_argument(
352 | "--model-path",
353 | type=str,
354 | default="mlx_model",
355 | help="The path to the model weights and tokenizer",
356 | )
357 | parser.add_argument(
358 | "--prompt",
359 | default="Once upon a midnight dreary",
360 | help="The message to be processed by the model",
361 | )
362 | parser.add_argument(
363 | "--max-tokens",
364 | "-m",
365 | type=int,
366 | default=100,
367 | help="Maximum number of tokens to generate",
368 | )
369 | parser.add_argument(
370 | "--temp",
371 | help="The sampling temperature.",
372 | type=float,
373 | default=0.0,
374 | )
375 | parser.add_argument("--seed", type=int, default=None, help="The PRNG seed")
376 | parser.add_argument(
377 | "--repeat-prompt",
378 | action=argparse.BooleanOptionalAction,
379 | help="Print prompt before start of generation",
380 | )
381 | parser.add_argument(
382 | "--schema",
383 | help="A JSON schema to constrain the output.",
384 | type=str,
385 | default=None,
386 | )
387 | parser.add_argument(
388 | "--encapsulated",
389 | action=argparse.BooleanOptionalAction,
390 | help="Whether the LLM is expected to encapsulate the JSON within ```json and ```.",
391 | )
392 | parser.add_argument(
393 | "--preemptive",
394 | type=int,
395 | default=0,
396 | help="If greater than zero, the maximum size of the batch for pre-emptive decoding",
397 | )
398 |
399 | args = parser.parse_args()
400 |
401 | info("Loading model from disk.")
402 | model = Model()
403 | model.load(args.model_path)
404 |
405 | if args.schema is not None:
406 | schema = json.loads(args.schema)
407 | info("Using schema")
408 | else:
409 | schema = None
410 | info("Starting generation...")
411 |
412 | for result in model.completion(
413 | prompt=args.prompt,
414 | schema=schema,
415 | encapsulated=args.encapsulated,
416 | max_tokens=args.max_tokens,
417 | temp=args.temp,
418 | seed=args.seed,
419 | preemptive_batch_size=args.preemptive,
420 | ):
421 | if result["op"] == "evaluatedPrompt":
422 | prompt_token_count = result["token_count"]
423 | prompt_time = result["time_ms"]
424 | prompt_tps = result["prompt_tps"]
425 | if args.repeat_prompt:
426 | bolddim(result["prompt"], flush=True)
427 | elif result["op"] == "generatedTokens":
428 | bold(result["text"], end="", flush=True)
429 | elif result["op"] == "stop":
430 | end_reason = result["reason"]
431 | generated_token_count = result["token_count"]
432 | generation_time = result["time_ms"]
433 | generation_tps = result["generation_tps"]
434 | else:
435 | assert False
436 |
437 | print()
438 | info(f"End reason: {end_reason}")
439 | info(f"Tokens: prompt {prompt_token_count}, generation {generated_token_count}")
440 | info(f"Tokens per second: prompt {prompt_tps:.2f}, generation {generation_tps:.2f}")
441 | info(f"Total time: prompt {prompt_time:.2f}ms, generation {generation_time:.2f}ms")
442 |
443 |
444 | if __name__ == "__main__":
445 | main()
446 |
--------------------------------------------------------------------------------
/src/examples/reluctance.py:
--------------------------------------------------------------------------------
1 | # pylint: disable=missing-class-docstring,missing-function-docstring
2 | """
3 | Example of JSON schema decoding with MLX.
4 | """
5 | import argparse
6 | import json
7 |
8 | import mlx.core as mx
9 | import mlx.nn as nn
10 | from mlx_lm.utils import load
11 |
12 | from llm_structured_output import (
13 | JsonSchemaAcceptorDriver,
14 | HuggingfaceTokenizerHelper,
15 | bias_logits,
16 | )
17 | from llm_structured_output.util.output import info, setbg, setfg, clear
18 |
19 | from .reusable_kv_cache import ReusableKVCache
20 |
21 |
22 | def compute_reluctance(logits, accepted_token_bitmap) -> float:
23 | """
24 | Sum the probabilities of each token that has higher probability than
25 | the highest-probability token selected by the schema. This gives an
26 | idea of the model's preference for tokens that don't follow the schema.
27 | """
28 | p = nn.softmax(logits)
29 | indices = mx.argsort(p)[::-1]
30 | r = 0
31 | for i in indices.tolist():
32 | if (1 << i) & accepted_token_bitmap:
33 | break
34 | r += p[i].item()
35 | return r
36 |
37 |
38 | def main():
39 | parser = argparse.ArgumentParser(
40 | description="Visualize LLM reluctance to generate according to the schema."
41 | )
42 | parser.add_argument(
43 | "--model-path",
44 | type=str,
45 | default="mlx_model",
46 | help="The path to the model weights and tokenizer",
47 | )
48 | parser.add_argument(
49 | "--schema",
50 | help="A JSON schema to constrain the output.",
51 | type=str,
52 | )
53 | parser.add_argument(
54 | "--prompt",
55 | help="The message to be processed by the model",
56 | )
57 | parser.add_argument(
58 | "--max-tokens",
59 | "-m",
60 | type=int,
61 | default=1000,
62 | help="Maximum number of tokens to generate",
63 | )
64 |
65 | args = parser.parse_args()
66 |
67 | info("Loading model from disk...")
68 | model, tokenizer = load(args.model_path)
69 | schema = json.loads(args.schema)
70 |
71 | tokenizer_helper = HuggingfaceTokenizerHelper(tokenizer)
72 | vocabulary, eos_id = tokenizer_helper.extract_vocabulary()
73 | token_acceptor_factory = JsonSchemaAcceptorDriver.driver_factory_for_model(vocabulary, eos_id)
74 | token_acceptor = token_acceptor_factory(schema)
75 |
76 |
77 | info("Starting generation...")
78 | tokens = tokenizer_helper.encode_prompt(args.prompt)
79 | cache = ReusableKVCache.for_model(model)
80 | while tokens[-1] != eos_id:
81 | logits = model(mx.array(tokens)[None], cache)
82 | accepted_token_bitmap = token_acceptor.select_valid_tokens()
83 | reluctance = compute_reluctance(logits[0, -1, :], accepted_token_bitmap)
84 | biased_logits = bias_logits(mx, logits[0, -1, :], accepted_token_bitmap)
85 | token = mx.argmax(biased_logits, axis=-1).item()
86 | if token == eos_id:
87 | break
88 | tokens = [token]
89 | text = tokenizer_helper.no_strip_decode(tokens)
90 | setbg(reluctance, 0.8 * (1 - reluctance), 0)
91 | setfg(1, 1, 1)
92 | print(text, end="")
93 | token_acceptor.advance_token(token)
94 | clear()
95 | print()
96 |
97 |
98 | if __name__ == "__main__":
99 | main()
100 |
--------------------------------------------------------------------------------
/src/examples/requirements.txt:
--------------------------------------------------------------------------------
1 | mlx >= 0.19.1
2 | mlx-lm >= 0.19.2
3 | tokenizers >= 0.20.1
4 | sentencepiece
5 | fastapi
6 | pydantic
7 | uvicorn
8 |
--------------------------------------------------------------------------------
/src/examples/reusable_kv_cache.py:
--------------------------------------------------------------------------------
1 | """
2 | Helper with improvements over mlx-lm's KVCache.
3 | """
4 |
5 | import mlx.core as mx
6 | from mlx_lm.models.cache import KVCache
7 |
8 |
9 | class ReusableKVCache(KVCache):
10 | """
11 | Usability improvements over KVCache.
12 | """
13 |
14 | @classmethod
15 | def for_model(cls, model):
16 | return [cls() for _ in model.layers]
17 |
18 | def reuse(self, new_prompt_length, common_prefix_length):
19 | """
20 | Reuse (part of) this cache for a new prompt that shares a prefix with it.
21 | """
22 | if self.keys is None:
23 | return
24 | # Clip the cache to the common length.
25 | self.offset = common_prefix_length
26 | # Make sure the cache can fit the whole prompt. Because the offset is
27 | # (very likely) not a multiple of the step size, update_and_fetch()
28 | # won't resize the cache when evaluating the rest of the prompt as it
29 | # would if it were an empty cache.
30 | current_size = self.keys.shape[2]
31 | if current_size < new_prompt_length:
32 | _, n_kv_heads, _, k_head_dim = self.keys.shape
33 | v_head_dim = self.values.shape[3]
34 | n_steps = (self.step + new_prompt_length - 1) // self.step
35 | k_add_shape = (1, n_kv_heads, n_steps * self.step - current_size, k_head_dim)
36 | v_add_shape = (1, n_kv_heads, n_steps * self.step - current_size, v_head_dim)
37 | k_zeros = mx.zeros(k_add_shape, self.keys.dtype)
38 | v_zeros = mx.zeros(v_add_shape, self.values.dtype)
39 | self.keys = mx.concatenate([self.keys, k_zeros], axis=2)
40 | self.values = mx.concatenate([self.values, v_zeros], axis=2)
41 |
42 | def update_and_fetch(self, keys, values):
43 | """
44 | Override the base class method to allow the cache to be used with batches of
45 | size greater than 1.
46 | This is just a tiny change in the line that determines the shape.
47 | """
48 | prev = self.offset
49 | if self.keys is None or (prev + keys.shape[2]) > self.keys.shape[2]:
50 | B, n_kv_heads, _, k_head_dim = keys.shape
51 | v_head_dim = values.shape[3]
52 | n_steps = (self.step + keys.shape[2] - 1) // self.step
53 | k_shape = (B, n_kv_heads, n_steps * self.step, k_head_dim)
54 | v_shape = (B, n_kv_heads, n_steps * self.step, v_head_dim)
55 | new_k = mx.zeros(k_shape, keys.dtype)
56 | new_v = mx.zeros(v_shape, values.dtype)
57 | if self.keys is not None:
58 | if prev % self.step != 0:
59 | self.keys = self.keys[..., :prev, :]
60 | self.values = self.values[..., :prev, :]
61 | self.keys = mx.concatenate([self.keys, new_k], axis=2)
62 | self.values = mx.concatenate([self.values, new_v], axis=2)
63 | else:
64 | self.keys, self.values = new_k, new_v
65 |
66 | self.offset += keys.shape[2]
67 | self.keys[..., prev : self.offset, :] = keys
68 | self.values[..., prev : self.offset, :] = values
69 | return self.keys[..., : self.offset, :], self.values[..., : self.offset, :]
70 |
71 |
--------------------------------------------------------------------------------
/src/examples/server.py:
--------------------------------------------------------------------------------
1 | # pylint: disable=missing-function-docstring,missing-class-docstring
2 | """
3 | Example model server with OpenAI-like API, including function calls / tools.
4 | """
5 | import json
6 | import time
7 | import os
8 | from enum import Enum
9 | from traceback import format_exc
10 | from typing import Literal, List, Optional, Union
11 |
12 | from fastapi import FastAPI, Request, status
13 | from fastapi.responses import FileResponse, JSONResponse, StreamingResponse
14 | from fastapi.exceptions import RequestValidationError
15 | from pydantic import BaseModel
16 |
17 | from examples.llm_schema import Model
18 | from llm_structured_output.util.output import info, warning, debug
19 |
20 |
21 | app = FastAPI()
22 |
23 | model = Model()
24 | info("Loading model...")
25 | try:
26 | model_path = os.environ["MODEL_PATH"]
27 | model.load(model_path)
28 | except KeyError:
29 | warning("Need to specify MODEL_PATH environment variable")
30 |
31 |
32 | @app.exception_handler(RequestValidationError)
33 | # pylint: disable-next=unused-argument
34 | async def validation_exception_handler(request: Request, exc: RequestValidationError):
35 | exc_str = f"{exc}"
36 | warning(f"RequestValidationError: {exc_str}")
37 | content = {"error": exc_str}
38 | return JSONResponse(
39 | content=content, status_code=status.HTTP_422_UNPROCESSABLE_ENTITY
40 | )
41 |
42 |
43 | @app.get("/status")
44 | def get_status():
45 | return {"status": "OK"}
46 |
47 |
48 | @app.get("/")
49 | def get_root():
50 | return FileResponse(f"{os.path.dirname(os.path.realpath(__file__))}/static/ui.html")
51 |
52 |
53 | class V1ChatMessageRole(str, Enum):
54 | SYSTEM = "system"
55 | USER = "user"
56 | ASSISTANT = "assistant"
57 |
58 |
59 | class V1ChatMessage(BaseModel):
60 | role: V1ChatMessageRole
61 | content: str
62 |
63 |
64 | class V1Function(BaseModel):
65 | name: str
66 | description: str = ""
67 | parameters: dict = {}
68 |
69 |
70 | class V1ToolFunction(BaseModel):
71 | type: Literal["function"]
72 | function: V1Function
73 |
74 |
75 | class V1ToolChoiceKeyword(str, Enum):
76 | AUTO = "auto"
77 | NONE = "none"
78 |
79 |
80 | class V1ToolChoiceFunction(BaseModel):
81 | type: Optional[Literal["function"]] = None
82 | name: str
83 |
84 |
85 | class V1ToolOptions(BaseModel): # Non-standard, our addition.
86 | # We automatically add instructions with the JSON schema
87 | # for the tool calls to the prompt. This option disables
88 | # it and is useful when the user prompt already includes
89 | # the schema and relevant instructions.
90 | no_prompt_steering: bool = False
91 |
92 |
93 | class V1ResponseFormatType(str, Enum):
94 | JSON_OBJECT = "json_object"
95 |
96 |
97 | class V1ResponseFormat(BaseModel):
98 | type: V1ResponseFormatType
99 | # schema is our addition, not an OpenAI API parameter
100 | schema: str = None
101 |
102 |
103 | class V1StreamOptions(BaseModel):
104 | include_usage: bool = False
105 |
106 |
107 | class V1ChatCompletionsRequest(
108 | BaseModel
109 | ): # pylint: disable=too-many-instance-attributes
110 | model: str = "default"
111 | max_tokens: int = 1000
112 | temperature: float = 0.0
113 | messages: List[V1ChatMessage]
114 | # The 'functions' and 'function_call' fields have been dreprecated and
115 | # replaced with 'tools' and 'tool_choice', that work similarly but allow
116 | # for multiple functions to be invoked.
117 | functions: List[V1Function] = None
118 | function_call: Union[V1ToolChoiceKeyword, V1ToolChoiceFunction] = None
119 | tools: List[V1ToolFunction] = None
120 | tool_choice: Union[V1ToolChoiceKeyword, V1ToolChoiceFunction] = None
121 | tool_options: V1ToolOptions = None
122 | response_format: V1ResponseFormat = None
123 | stream: bool = False
124 | stream_options: V1StreamOptions = None
125 |
126 |
127 | @app.post("/v1/chat/completions")
128 | async def post_v1_chat_completions(request: V1ChatCompletionsRequest):
129 | debug("REQUEST", request)
130 | if request.stream:
131 | async def get_content():
132 | try:
133 | async for message in post_v1_chat_completions_impl(request):
134 | yield message
135 | # pylint: disable-next=broad-exception-caught
136 | except Exception as e:
137 | warning(format_exc())
138 | yield 'data: {"choices": [{"index": 0, "finish_reason": "error: ' + str(e) + '"}]}'
139 | return StreamingResponse(
140 | content=get_content(),
141 | media_type="text/event-stream",
142 | )
143 | else:
144 | # FUTURE: Python 3.10 can use `await anext(x))` instead of `await x.__anext__()`.
145 | try:
146 | response = await post_v1_chat_completions_impl(request).__anext__()
147 | # pylint: disable-next=broad-exception-caught
148 | except Exception as e:
149 | warning(format_exc())
150 | content = {"error": str(e)}
151 | response = JSONResponse(
152 | content=content, status_code=status.HTTP_500_INTERNAL_SERVER_ERROR
153 | )
154 | debug("RESPONSE", response)
155 | return response
156 |
157 |
158 | async def post_v1_chat_completions_impl(request: V1ChatCompletionsRequest):
159 | messages = request.messages[:]
160 |
161 | # Extract valid functions from the request.
162 | functions = []
163 | is_legacy_function_call = False
164 | if request.tool_choice == "none":
165 | pass
166 | elif request.tool_choice == "auto":
167 | functions = [tool.function for tool in request.tools if tool.type == "function"]
168 | elif request.tool_choice is not None:
169 | functions = [
170 | next(
171 | tool.function
172 | for tool in request.tools
173 | if tool.type == "function"
174 | and tool.function.name == request.function_call.name
175 | )
176 | ]
177 | elif request.function_call == "none":
178 | pass
179 | elif request.function_call == "auto":
180 | functions = request.functions
181 | is_legacy_function_call = True
182 | elif request.function_call is not None:
183 | functions = [
184 | next(
185 | fn for fn in request.functions if fn.name == request.function_call.name
186 | )
187 | ]
188 | is_legacy_function_call = True
189 |
190 | model_name = model_path
191 | schema = None
192 | if functions:
193 | # If the request includes functions, create a system prompt to instruct the LLM
194 | # to use tools, and assemble a JSON schema to steer the LLM output.
195 | if request.stream:
196 | responder = ToolCallStreamingResponder(
197 | model_name,
198 | functions,
199 | is_legacy_function_call,
200 | model,
201 | )
202 | else:
203 | responder = ToolCallResponder(
204 | model_name, functions, is_legacy_function_call
205 | )
206 | if not (request.tool_options and request.tool_options.no_prompt_steering):
207 | messages.insert(
208 | 0,
209 | V1ChatMessage(
210 | role="system",
211 | content=responder.tool_prompt,
212 | ),
213 | )
214 | schema = responder.schema
215 | else:
216 | if request.response_format:
217 | assert request.response_format.type == V1ResponseFormatType.JSON_OBJECT
218 | # The request may specify a JSON schema (this option is not in the OpenAI API)
219 | if request.response_format.schema:
220 | schema = json.loads(request.response_format.schema)
221 | else:
222 | schema = {"type": "object"}
223 | if request.stream:
224 | responder = ChatCompletionStreamingResponder(model_name, schema, model)
225 | else:
226 | responder = ChatCompletionResponder(model_name)
227 |
228 | if schema is not None:
229 | debug("Using schema:", schema)
230 |
231 | info("Starting generation...")
232 |
233 | prompt_tokens = None
234 |
235 | for result in model.completion(
236 | messages,
237 | schema=schema,
238 | max_tokens=request.max_tokens,
239 | temp=request.temperature,
240 | cache_prompt=True,
241 | ):
242 | if result["op"] == "evaluatedPrompt":
243 | prompt_tokens = result["token_count"]
244 | elif result["op"] == "generatedTokens":
245 | message = responder.generated_tokens(result["text"])
246 | if message:
247 | yield message
248 | elif result["op"] == "stop":
249 | completion_tokens = result["token_count"]
250 | yield responder.generation_stopped(
251 | result["reason"], prompt_tokens, completion_tokens
252 | )
253 | else:
254 | assert False
255 |
256 |
257 | class ChatCompletionResponder:
258 | def __init__(self, model_name: str):
259 | self.object_type = "chat.completion"
260 | self.model_name = model_name
261 | self.created = int(time.time())
262 | self.id = f"{id(self)}_{self.created}"
263 | self.content = ""
264 |
265 | def message_properties(self):
266 | return {
267 | "object": self.object_type,
268 | "id": f"chatcmpl-{self.id}",
269 | "created": self.created,
270 | "model": self.model_name,
271 | }
272 |
273 | def translate_reason(self, reason):
274 | """
275 | Translate our reason codes to OpenAI ones.
276 | """
277 | if reason == "end":
278 | return "stop"
279 | if reason == "max_tokens":
280 | return "length"
281 | return f"error: {reason}" # Not a standard OpenAI API reason
282 |
283 | def format_usage(self, prompt_tokens: int, completion_tokens: int):
284 | return {
285 | "usage": {
286 | "completion_tokens": completion_tokens,
287 | "prompt_tokens": prompt_tokens,
288 | "total_tokens": completion_tokens + prompt_tokens,
289 | },
290 | }
291 |
292 | def generated_tokens(
293 | self,
294 | text: str,
295 | ):
296 | self.content += text
297 | return None
298 |
299 | def generation_stopped(
300 | self,
301 | stop_reason: str,
302 | prompt_tokens: int,
303 | completion_tokens: int,
304 | ):
305 | finish_reason = self.translate_reason(stop_reason)
306 | message = {"role": "assistant", "content": self.content}
307 | return {
308 | "choices": [
309 | {"index": 0, "message": message, "finish_reason": finish_reason}
310 | ],
311 | **self.format_usage(prompt_tokens, completion_tokens),
312 | **self.message_properties(),
313 | }
314 |
315 |
316 | class ChatCompletionStreamingResponder(ChatCompletionResponder):
317 | def __init__(self, model_name: str, schema: dict = None, _model = None):
318 | super().__init__(model_name)
319 | self.object_type = "chat.completion.chunk"
320 | if schema:
321 | assert _model
322 | self.schema_parser = _model.get_driver_for_json_schema(schema)
323 | else:
324 | self.schema_parser = None
325 |
326 | def generated_tokens(
327 | self,
328 | text: str,
329 | ):
330 | delta = {"role": "assistant", "content": text}
331 | if self.schema_parser:
332 | values = {}
333 | for char in text:
334 | self.schema_parser.advance_char(char)
335 | for path in self.schema_parser.get_current_value_paths():
336 | values[path] = values.get(path, "") + char
337 | delta["values"] = values
338 | message = {
339 | "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
340 | **self.message_properties(),
341 | }
342 | return f"data: {json.dumps(message)}\n"
343 |
344 | def generation_stopped(
345 | self,
346 | stop_reason: str,
347 | prompt_tokens: int,
348 | completion_tokens: int,
349 | ):
350 | finish_reason = self.translate_reason(stop_reason)
351 | delta = {"role": "assistant", "content": ""}
352 | message = {
353 | "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
354 | # Usage field notes:
355 | # - OpenAI only sends usage in streaming if the option
356 | # stream_options.include_usage is true, but we send it always.
357 | **self.format_usage(prompt_tokens, completion_tokens),
358 | **self.message_properties(),
359 | }
360 | return f"data: {json.dumps(message)}\ndata: [DONE]\n"
361 |
362 |
363 | class ToolCallResponder(ChatCompletionResponder):
364 | def __init__(
365 | self, model_name: str, functions: list[dict], is_legacy_function_call: bool
366 | ):
367 | super().__init__(model_name)
368 |
369 | self.is_legacy_function_call = is_legacy_function_call
370 |
371 | function_schemas = [
372 | {
373 | "type": "object",
374 | "properties": {
375 | "name": {"type": "const", "const": fn.name},
376 | "arguments": fn.parameters,
377 | },
378 | "required": ["name", "arguments"],
379 | }
380 | for fn in functions
381 | ]
382 | if len(function_schemas) == 1:
383 | self.schema = function_schemas[0]
384 | self.tool_prompt = self._one_tool_prompt(functions[0], function_schemas[0])
385 | elif is_legacy_function_call: # Only allows one function to be called.
386 | self.schema = {"oneOf": function_schemas}
387 | self.tool_prompt = self._select_tool_prompt(functions, function_schemas)
388 | else:
389 | self.schema = {"type": "array", "items": {"anyOf": function_schemas}}
390 | self.tool_prompt = self._multiple_tool_prompt(functions, function_schemas)
391 |
392 | def translate_reason(self, reason):
393 | if reason == "end":
394 | if self.is_legacy_function_call:
395 | return "function_call"
396 | return "tool_calls"
397 | return super().translate_reason(reason)
398 |
399 | def generation_stopped(
400 | self,
401 | stop_reason: str,
402 | prompt_tokens: int,
403 | completion_tokens: int,
404 | ):
405 | finish_reason = self.translate_reason(stop_reason)
406 | if finish_reason == "tool_calls":
407 | tool_calls = json.loads(self.content)
408 | if not isinstance(tool_calls, list):
409 | # len(functions) == 1 was special cased
410 | tool_calls = [tool_calls]
411 | message = {
412 | "role": "assistant",
413 | "tool_calls": [
414 | {
415 | "id": f"call_{self.id}_{i}",
416 | "type": "function",
417 | "function": {
418 | "name": function_call["name"],
419 | "arguments": json.dumps(function_call["arguments"]),
420 | },
421 | }
422 | for i, function_call in enumerate(tool_calls)
423 | ],
424 | }
425 | elif finish_reason == "function_call":
426 | function_call = json.loads(self.content)
427 | message = {
428 | "role": "assistant",
429 | "function_call": {
430 | "name": function_call["name"],
431 | "arguments": json.dumps(function_call["arguments"]),
432 | },
433 | }
434 | else:
435 | message = None
436 | return {
437 | "choices": [
438 | {"index": 0, "message": message, "finish_reason": finish_reason}
439 | ],
440 | **self.format_usage(prompt_tokens, completion_tokens),
441 | **self.message_properties(),
442 | }
443 |
444 | def _one_tool_prompt(self, tool, tool_schema):
445 | return f"""
446 | You are a helpful assistant with access to a tool that you must invoke to answer the user's request.
447 | The tool is:
448 | Tool {tool.name}: {tool.description}
449 | Invocation schema: {json.dumps(tool_schema)}
450 | Your answer is a JSON object according to the invocation schema in order to answer the user request below.
451 | """
452 |
453 | def _multiple_tool_prompt(self, tools, tool_schemas, separator="\n"):
454 | return f"""
455 | You are a helpful assistant with access to tools that you must invoke to answer the user's request.
456 | The following tools are available:
457 | {separator.join([ f'''
458 | Tool {tool.name}: {tool.description}
459 | Invocation schema: {json.dumps(tool_schema)}
460 | ''' for tool, tool_schema in zip(tools, tool_schemas) ])}
461 | Your answer is a JSON array with one or more tool invocations according to the appropriate schema(s)
462 | in order to answer the user request below.
463 | """
464 |
465 | def _select_tool_prompt(self, tools, tool_schemas, separator="\n"):
466 | return f"""
467 | You are a helpful assistant with access to tools that you must invoke to answer the user's request.
468 | The following tools are available:
469 | {separator.join([ f'''
470 | Function {tool.name}: {tool.description}
471 | Tool schema: {json.dumps(tool_schema)}
472 | ''' for tool, tool_schema in zip(tools, tool_schemas) ])}
473 | Your answer is a JSON object according to the invocation schema of the most appropriate tool to use
474 | to answer the user request below.
475 | """
476 |
477 |
478 | class ToolCallStreamingResponder(ToolCallResponder):
479 | def __init__(
480 | self,
481 | model_name: str,
482 | functions: list[dict],
483 | is_legacy_function_call: bool,
484 | _model,
485 | ):
486 | super().__init__(model_name, functions, is_legacy_function_call)
487 | self.object_type = "chat.completion.chunk"
488 |
489 | # We need to parse the output as it's being generated in order to send
490 | # streaming messages that contain the name and arguments of the function
491 | # being called.
492 |
493 | self.current_function_index = -1
494 | self.current_function_name = None
495 | self.in_function_arguments = False
496 |
497 | def set_function_name(_prop_name: str, prop_value):
498 | self.current_function_index += 1
499 | self.current_function_name = prop_value
500 |
501 | def start_function_arguments(_prop_name: str):
502 | self.in_function_arguments = True
503 |
504 | def end_function_arguments(_prop_name: str, _prop_value: str):
505 | self.in_function_arguments = False
506 |
507 | hooked_function_schemas = [
508 | {
509 | "type": "object",
510 | "properties": {
511 | "name": {
512 | "type": "const",
513 | "const": fn.name,
514 | "__hooks": {
515 | "value_end": set_function_name,
516 | },
517 | },
518 | "arguments": {
519 | **fn.parameters,
520 | "__hooks": {
521 | "value_start": start_function_arguments,
522 | "value_end": end_function_arguments,
523 | },
524 | },
525 | },
526 | "required": ["name", "arguments"],
527 | }
528 | for fn in functions
529 | ]
530 | if len(hooked_function_schemas) == 1:
531 | hooked_schema = hooked_function_schemas[0]
532 | elif is_legacy_function_call:
533 | hooked_schema = {"oneOf": hooked_function_schemas}
534 | else:
535 | hooked_schema = {
536 | "type": "array",
537 | "items": {"anyOf": hooked_function_schemas},
538 | }
539 | self.tool_call_parser = _model.get_driver_for_json_schema(hooked_schema)
540 |
541 | def generated_tokens(
542 | self,
543 | text: str,
544 | ):
545 | argument_text = ""
546 | for char in text:
547 | if self.in_function_arguments:
548 | argument_text += char
549 | # Update state. This is certain to parse, no need to check for rejections.
550 | self.tool_call_parser.advance_char(char)
551 | if not argument_text:
552 | return None
553 | assert self.current_function_name
554 | if self.is_legacy_function_call:
555 | delta = {
556 | "function_call": {
557 | "name": self.current_function_name,
558 | "arguments": argument_text,
559 | }
560 | }
561 | else:
562 | delta = {
563 | "tool_calls": [
564 | {
565 | "index": self.current_function_index,
566 | "id": f"call_{self.id}_{self.current_function_index}",
567 | "type": "function",
568 | "function": {
569 | # We send the name on every update, but OpenAI only sends it on
570 | # the first one for each call, with empty arguments (""). Further
571 | # updates only have the arguments field. This is something we may
572 | # want to emulate if client code depends on this behavior.
573 | "name": self.current_function_name,
574 | "arguments": argument_text,
575 | },
576 | }
577 | ]
578 | }
579 | message = {
580 | "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
581 | **self.message_properties(),
582 | }
583 | return f"data: {json.dumps(message)}\n"
584 |
585 | def generation_stopped(
586 | self,
587 | stop_reason: str,
588 | prompt_tokens: int,
589 | completion_tokens: int,
590 | ):
591 | finish_reason = self.translate_reason(stop_reason)
592 | message = {
593 | "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}],
594 | # Usage field notes:
595 | # - OpenAI only sends usage in streaming if the option
596 | # stream_options.include_usage is true, but we send it always.
597 | # - OpenAI sends two separate messages: one with the finish_reason and no
598 | # usage field, and one with an empty choices array and the usage field.
599 | **self.format_usage(prompt_tokens, completion_tokens),
600 | **self.message_properties(),
601 | }
602 | return f"data: {json.dumps(message)}\ndata: [DONE]\n"
603 |
--------------------------------------------------------------------------------
/src/examples/static/attention.html:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 | Attention
6 |
176 |
386 |
387 |
388 |
389 |
390 |
Prompt
391 |
392 |
393 |
394 |
395 |
396 |
397 |
399 |
400 |
401 |
402 |
403 |
404 |
406 |
Prompt perplexity
407 |
408 |
409 |
410 |
Attention
411 |
425 |
426 |
436 |
437 |
453 |
454 |
455 |
--------------------------------------------------------------------------------
/src/examples/static/ui.html:
--------------------------------------------------------------------------------
1 |
2 |
3 |
4 |
5 |
6 |
7 | LLM
8 |
163 |
295 |
296 |
297 |
298 |
299 |
Prompt
300 |
301 |
302 |
303 |
304 |
305 |
306 |
308 |
309 |
310 |
321 |
322 |
323 |
324 |
325 |
327 |
Output
328 |
329 |
330 |
334 |