llama.cpp/llama_cpp/server/app.py

891 lines
30 KiB
Python
Raw Normal View History

2023-09-30 23:13:36 +00:00
import sys
import json
2023-10-10 19:56:04 +00:00
import traceback
2023-05-07 07:03:57 +00:00
import multiprocessing
import time
2023-07-16 05:57:39 +00:00
from re import compile, Match, Pattern
from threading import Lock
2023-05-27 13:12:58 +00:00
from functools import partial
from typing import Callable, Coroutine, Iterator, List, Optional, Tuple, Union, Dict
2023-05-07 07:03:57 +00:00
from typing_extensions import TypedDict, Literal
import llama_cpp
2023-05-27 13:12:58 +00:00
import anyio
from anyio.streams.memory import MemoryObjectSendStream
from starlette.concurrency import run_in_threadpool, iterate_in_threadpool
2023-07-16 05:57:39 +00:00
from fastapi import Depends, FastAPI, APIRouter, Request, Response
from fastapi.middleware import Middleware
from fastapi.middleware.cors import CORSMiddleware
2023-07-16 05:57:39 +00:00
from fastapi.responses import JSONResponse
from fastapi.routing import APIRoute
from pydantic import BaseModel, Field
from pydantic_settings import BaseSettings
from sse_starlette.sse import EventSourceResponse
from starlette_context import plugins
from starlette_context.middleware import RawContextMiddleware
import numpy as np
import numpy.typing as npt
2023-09-29 23:59:12 +00:00
# Disable warning for model and model_alias settings
2023-09-14 03:01:34 +00:00
BaseSettings.model_config['protected_namespaces'] = ()
class Settings(BaseSettings):
2023-05-07 06:52:20 +00:00
model: str = Field(
description="The path to the model to use for generating completions."
)
model_alias: Optional[str] = Field(
default=None,
description="The alias of the model to use for generating completions.",
)
# Model Params
2023-05-14 04:04:22 +00:00
n_gpu_layers: int = Field(
default=0,
ge=-1,
description="The number of layers to put on the GPU. The rest will be on the CPU. Set -1 to move all to GPU.",
2023-05-14 04:04:22 +00:00
)
main_gpu: int = Field(
default=0,
ge=0,
description="Main GPU to use.",
)
tensor_split: Optional[List[float]] = Field(
2023-07-07 09:22:10 +00:00
default=None,
description="Split layers across multiple GPUs in proportion.",
)
vocab_only: bool = Field(
default=False, description="Whether to only return the vocabulary."
2023-05-07 06:52:20 +00:00
)
use_mmap: bool = Field(
default=llama_cpp.llama_mmap_supported(),
2023-05-07 06:52:20 +00:00
description="Use mmap.",
)
use_mlock: bool = Field(
default=llama_cpp.llama_mlock_supported(),
description="Use mlock.",
)
# Context Params
seed: int = Field(default=llama_cpp.LLAMA_DEFAULT_SEED, description="Random seed. -1 for random.")
n_ctx: int = Field(default=2048, ge=1, description="The context size.")
n_batch: int = Field(
default=512, ge=1, description="The batch size to use per eval."
)
n_threads: int = Field(
default=max(multiprocessing.cpu_count() // 2, 1),
ge=1,
description="The number of threads to use.",
2023-06-15 02:13:42 +00:00
)
n_threads_batch: int = Field(
default=max(multiprocessing.cpu_count() // 2, 1),
ge=0,
description="The number of threads to use when batch processing.",
)
rope_scaling_type: int = Field(
default=llama_cpp.LLAMA_ROPE_SCALING_UNSPECIFIED
)
rope_freq_base: float = Field(
default=0.0, description="RoPE base frequency"
)
rope_freq_scale: float = Field(
default=0.0, description="RoPE frequency scaling factor"
)
yarn_ext_factor: float = Field(
2023-11-03 15:34:50 +00:00
default=-1.0
)
yarn_attn_factor: float = Field(
default=1.0
)
yarn_beta_fast: float = Field(
default=32.0
)
yarn_beta_slow: float = Field(
default=1.0
)
yarn_orig_ctx: int = Field(
default=0
)
mul_mat_q: bool = Field(
default=True, description="if true, use experimental mul_mat_q kernels"
)
f16_kv: bool = Field(default=True, description="Whether to use f16 key/value.")
logits_all: bool = Field(default=True, description="Whether to return logits.")
embedding: bool = Field(default=True, description="Whether to use embeddings.")
# Sampling Params
2023-05-07 06:52:20 +00:00
last_n_tokens_size: int = Field(
default=64,
ge=0,
description="Last n tokens to keep for repeat penalty calculation.",
)
# LoRA Params
lora_base: Optional[str] = Field(
default=None,
description="Optional path to base model, useful if using a quantized base model and you want to apply LoRA to an f16 model."
)
lora_path: Optional[str] = Field(
default=None,
description="Path to a LoRA file to apply to the model.",
)
# Backend Params
numa: bool = Field(
default=False,
description="Enable NUMA support.",
)
# Chat Format Params
2023-09-29 23:59:12 +00:00
chat_format: str = Field(
default="llama-2",
description="Chat format to use.",
)
# Cache Params
2023-05-07 06:52:20 +00:00
cache: bool = Field(
default=False,
description="Use a cache to reduce processing times for evaluated prompts.",
)
cache_type: Literal["ram", "disk"] = Field(
default="ram",
description="The type of cache to use. Only used if cache is True.",
)
2023-05-07 23:33:17 +00:00
cache_size: int = Field(
default=2 << 30,
description="The size of the cache in bytes. Only used if cache is True.",
)
# Misc
2023-05-07 09:09:10 +00:00
verbose: bool = Field(
default=True, description="Whether to print debug information."
)
# Server Params
2023-07-14 03:25:12 +00:00
host: str = Field(default="localhost", description="Listen address")
port: int = Field(default=8000, description="Listen port")
interrupt_requests: bool = Field(
default=True,
description="Whether to interrupt requests when a new request is received.",
)
2023-07-16 05:57:39 +00:00
class ErrorResponse(TypedDict):
"""OpenAI style error response"""
message: str
type: str
param: Optional[str]
code: Optional[str]
class ErrorResponseFormatters:
"""Collection of formatters for error responses.
Args:
request (Union[CreateCompletionRequest, CreateChatCompletionRequest]):
Request body
match (Match[str]): Match object from regex pattern
Returns:
Tuple[int, ErrorResponse]: Status code and error response
2023-07-16 05:57:39 +00:00
"""
@staticmethod
def context_length_exceeded(
2023-09-14 01:23:23 +00:00
request: Union["CreateCompletionRequest", "CreateChatCompletionRequest"],
match, # type: Match[str] # type: ignore
) -> Tuple[int, ErrorResponse]:
2023-07-16 05:57:39 +00:00
"""Formatter for context length exceeded error"""
context_window = int(match.group(2))
prompt_tokens = int(match.group(1))
completion_tokens = request.max_tokens
if hasattr(request, "messages"):
# Chat completion
message = (
"This model's maximum context length is {} tokens. "
"However, you requested {} tokens "
"({} in the messages, {} in the completion). "
"Please reduce the length of the messages or completion."
)
else:
# Text completion
message = (
"This model's maximum context length is {} tokens, "
"however you requested {} tokens "
"({} in your prompt; {} for the completion). "
"Please reduce your prompt; or completion length."
)
return 400, ErrorResponse(
message=message.format(
context_window,
completion_tokens + prompt_tokens,
prompt_tokens,
completion_tokens,
),
type="invalid_request_error",
param="messages",
code="context_length_exceeded",
)
@staticmethod
def model_not_found(
2023-09-14 01:23:23 +00:00
request: Union["CreateCompletionRequest", "CreateChatCompletionRequest"],
match, # type: Match[str] # type: ignore
) -> Tuple[int, ErrorResponse]:
2023-07-16 05:57:39 +00:00
"""Formatter for model_not_found error"""
model_path = str(match.group(1))
message = f"The model `{model_path}` does not exist"
return 400, ErrorResponse(
message=message,
type="invalid_request_error",
param=None,
code="model_not_found",
)
class RouteErrorHandler(APIRoute):
"""Custom APIRoute that handles application errors and exceptions"""
# key: regex pattern for original error message from llama_cpp
# value: formatter function
pattern_and_formatters: Dict[
2023-07-16 05:57:39 +00:00
"Pattern",
Callable[
[
Union["CreateCompletionRequest", "CreateChatCompletionRequest"],
"Match[str]",
2023-07-16 05:57:39 +00:00
],
Tuple[int, ErrorResponse],
2023-07-16 05:57:39 +00:00
],
] = {
compile(
r"Requested tokens \((\d+)\) exceed context window of (\d+)"
): ErrorResponseFormatters.context_length_exceeded,
compile(
r"Model path does not exist: (.+)"
): ErrorResponseFormatters.model_not_found,
}
def error_message_wrapper(
self,
error: Exception,
body: Optional[
Union[
"CreateChatCompletionRequest",
"CreateCompletionRequest",
"CreateEmbeddingRequest",
]
] = None,
) -> Tuple[int, ErrorResponse]:
2023-07-16 05:57:39 +00:00
"""Wraps error message in OpenAI style error response"""
2023-09-30 23:13:36 +00:00
print(f"Exception: {str(error)}", file=sys.stderr)
2023-10-10 19:56:04 +00:00
traceback.print_exc(file=sys.stderr)
2023-07-16 05:57:39 +00:00
if body is not None and isinstance(
body,
(
CreateCompletionRequest,
CreateChatCompletionRequest,
),
):
# When text completion or chat completion
for pattern, callback in self.pattern_and_formatters.items():
match = pattern.search(str(error))
if match is not None:
return callback(body, match)
# Wrap other errors as internal server error
return 500, ErrorResponse(
message=str(error),
type="internal_server_error",
param=None,
code=None,
)
def get_route_handler(
self,
) -> Callable[[Request], Coroutine[None, None, Response]]:
"""Defines custom route handler that catches exceptions and formats
in OpenAI style error response"""
original_route_handler = super().get_route_handler()
async def custom_route_handler(request: Request) -> Response:
try:
start_sec = time.perf_counter()
response = await original_route_handler(request)
elapsed_time_ms = int((time.perf_counter() - start_sec) * 1000)
response.headers["openai-processing-ms"] = f"{elapsed_time_ms}"
return response
2023-07-16 05:57:39 +00:00
except Exception as exc:
json_body = await request.json()
try:
if "messages" in json_body:
# Chat completion
body: Optional[
Union[
CreateChatCompletionRequest,
CreateCompletionRequest,
CreateEmbeddingRequest,
]
] = CreateChatCompletionRequest(**json_body)
elif "prompt" in json_body:
# Text completion
body = CreateCompletionRequest(**json_body)
else:
# Embedding
body = CreateEmbeddingRequest(**json_body)
except Exception:
# Invalid request body
body = None
# Get proper error message from the exception
(
status_code,
error_message,
) = self.error_message_wrapper(error=exc, body=body)
return JSONResponse(
{"error": error_message},
status_code=status_code,
)
return custom_route_handler
router = APIRouter(route_class=RouteErrorHandler)
2023-05-02 02:38:46 +00:00
settings: Optional[Settings] = None
2023-05-02 02:38:46 +00:00
llama: Optional[llama_cpp.Llama] = None
2023-05-02 02:38:46 +00:00
def create_app(settings: Optional[Settings] = None):
if settings is None:
settings = Settings()
middleware = [
2023-09-14 01:23:23 +00:00
Middleware(RawContextMiddleware, plugins=(plugins.RequestIdPlugin(),))
]
2023-05-02 02:38:46 +00:00
app = FastAPI(
middleware=middleware,
2023-05-02 02:38:46 +00:00
title="🦙 llama.cpp Python API",
version="0.0.1",
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
app.include_router(router)
global llama
llama = llama_cpp.Llama(
2023-05-02 02:38:46 +00:00
model_path=settings.model,
# Model Params
2023-05-14 04:04:22 +00:00
n_gpu_layers=settings.n_gpu_layers,
main_gpu=settings.main_gpu,
2023-07-07 09:22:10 +00:00
tensor_split=settings.tensor_split,
vocab_only=settings.vocab_only,
use_mmap=settings.use_mmap,
use_mlock=settings.use_mlock,
# Context Params
seed=settings.seed,
n_ctx=settings.n_ctx,
n_batch=settings.n_batch,
n_threads=settings.n_threads,
n_threads_batch=settings.n_threads_batch,
rope_scaling_type=settings.rope_scaling_type,
2023-07-18 08:34:36 +00:00
rope_freq_base=settings.rope_freq_base,
rope_freq_scale=settings.rope_freq_scale,
yarn_ext_factor=settings.yarn_ext_factor,
yarn_attn_factor=settings.yarn_attn_factor,
yarn_beta_fast=settings.yarn_beta_fast,
yarn_beta_slow=settings.yarn_beta_slow,
yarn_orig_ctx=settings.yarn_orig_ctx,
mul_mat_q=settings.mul_mat_q,
f16_kv=settings.f16_kv,
logits_all=settings.logits_all,
embedding=settings.embedding,
# Sampling Params
last_n_tokens_size=settings.last_n_tokens_size,
# LoRA Params
lora_base=settings.lora_base,
lora_path=settings.lora_path,
# Backend Params
2023-09-29 23:59:12 +00:00
numa=settings.numa,
# Chat Format Params
2023-09-29 23:59:12 +00:00
chat_format=settings.chat_format,
# Misc
2023-05-07 09:09:10 +00:00
verbose=settings.verbose,
)
if settings.cache:
if settings.cache_type == "disk":
2023-06-15 01:46:48 +00:00
if settings.verbose:
print(f"Using disk cache with size {settings.cache_size}")
cache = llama_cpp.LlamaDiskCache(capacity_bytes=settings.cache_size)
else:
2023-06-15 01:46:48 +00:00
if settings.verbose:
print(f"Using ram cache with size {settings.cache_size}")
cache = llama_cpp.LlamaRAMCache(capacity_bytes=settings.cache_size)
2023-05-07 23:33:17 +00:00
cache = llama_cpp.LlamaCache(capacity_bytes=settings.cache_size)
llama.set_cache(cache)
def set_settings(_settings: Settings):
global settings
settings = _settings
set_settings(settings)
2023-05-02 02:38:46 +00:00
return app
llama_outer_lock = Lock()
llama_inner_lock = Lock()
2023-05-02 02:38:46 +00:00
def get_llama():
# NOTE: This double lock allows the currently streaming llama model to
# check if any other requests are pending in the same thread and cancel
# the stream if so.
llama_outer_lock.acquire()
release_outer_lock = True
try:
llama_inner_lock.acquire()
try:
llama_outer_lock.release()
release_outer_lock = False
yield llama
finally:
llama_inner_lock.release()
finally:
if release_outer_lock:
llama_outer_lock.release()
2023-05-07 06:52:20 +00:00
def get_settings():
yield settings
2023-07-16 05:57:39 +00:00
async def get_event_publisher(
request: Request,
inner_send_chan: MemoryObjectSendStream,
iterator: Iterator,
):
async with inner_send_chan:
try:
async for chunk in iterate_in_threadpool(iterator):
await inner_send_chan.send(dict(data=json.dumps(chunk)))
if await request.is_disconnected():
raise anyio.get_cancelled_exc_class()()
if settings.interrupt_requests and llama_outer_lock.locked():
await inner_send_chan.send(dict(data="[DONE]"))
raise anyio.get_cancelled_exc_class()()
await inner_send_chan.send(dict(data="[DONE]"))
except anyio.get_cancelled_exc_class() as e:
print("disconnected")
with anyio.move_on_after(1, shield=True):
2023-09-14 01:23:23 +00:00
print(f"Disconnected from client (via refresh/close) {request.client}")
2023-07-16 05:57:39 +00:00
raise e
2023-09-14 01:23:23 +00:00
model_field = Field(
description="The model to use for generating completions.", default=None
)
max_tokens_field = Field(
2023-07-16 05:57:39 +00:00
default=16, ge=1, description="The maximum number of tokens to generate."
)
temperature_field = Field(
default=0.8,
ge=0.0,
le=2.0,
2023-05-07 06:52:20 +00:00
description="Adjust the randomness of the generated text.\n\n"
+ "Temperature is a hyperparameter that controls the randomness of the generated text. It affects the probability distribution of the model's output tokens. A higher temperature (e.g., 1.5) makes the output more random and creative, while a lower temperature (e.g., 0.5) makes the output more focused, deterministic, and conservative. The default value is 0.8, which provides a balance between randomness and determinism. At the extreme, a temperature of 0 will always pick the most likely next token, leading to identical outputs in each run.",
)
top_p_field = Field(
default=0.95,
ge=0.0,
le=1.0,
2023-05-07 06:52:20 +00:00
description="Limit the next token selection to a subset of tokens with a cumulative probability above a threshold P.\n\n"
+ "Top-p sampling, also known as nucleus sampling, is another text generation method that selects the next token from a subset of tokens that together have a cumulative probability of at least p. This method provides a balance between diversity and quality by considering both the probabilities of tokens and the number of tokens to sample from. A higher value for top_p (e.g., 0.95) will lead to more diverse text, while a lower value (e.g., 0.5) will generate more focused and conservative text.",
)
stop_field = Field(
default=None,
2023-05-07 06:52:20 +00:00
description="A list of tokens at which to stop generation. If None, no stop tokens are used.",
)
stream_field = Field(
default=False,
2023-05-07 06:52:20 +00:00
description="Whether to stream the results as they are generated. Useful for chatbots.",
)
top_k_field = Field(
default=40,
ge=0,
2023-05-07 06:52:20 +00:00
description="Limit the next token selection to the K most probable tokens.\n\n"
+ "Top-k sampling is a text generation method that selects the next token only from the top k most likely tokens predicted by the model. It helps reduce the risk of generating low-probability or nonsensical tokens, but it may also limit the diversity of the output. A higher value for top_k (e.g., 100) will consider more tokens and lead to more diverse text, while a lower value (e.g., 10) will focus on the most probable tokens and generate more conservative text.",
)
repeat_penalty_field = Field(
2023-05-08 22:49:11 +00:00
default=1.1,
ge=0.0,
2023-05-07 06:52:20 +00:00
description="A penalty applied to each token that is already generated. This helps prevent the model from repeating itself.\n\n"
+ "Repeat penalty is a hyperparameter used to penalize the repetition of token sequences during text generation. It helps prevent the model from generating repetitive or monotonous text. A higher value (e.g., 1.5) will penalize repetitions more strongly, while a lower value (e.g., 0.9) will be more lenient.",
)
presence_penalty_field = Field(
default=0.0,
ge=-2.0,
le=2.0,
description="Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics.",
)
frequency_penalty_field = Field(
default=0.0,
ge=-2.0,
le=2.0,
description="Positive values penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim.",
)
2023-05-07 06:52:20 +00:00
mirostat_mode_field = Field(
default=0,
ge=0,
le=2,
2023-07-14 03:25:12 +00:00
description="Enable Mirostat constant-perplexity algorithm of the specified version (1 or 2; 0 = disabled)",
)
mirostat_tau_field = Field(
default=5.0,
ge=0.0,
le=10.0,
2023-07-14 03:25:12 +00:00
description="Mirostat target entropy, i.e. the target perplexity - lower values produce focused and coherent text, larger values produce more diverse and less coherent text",
)
mirostat_eta_field = Field(
2023-07-14 03:25:12 +00:00
default=0.1, ge=0.001, le=1.0, description="Mirostat learning rate"
)
grammar = Field(
default=None,
description="A CBNF grammar (as string) to be used for formatting the model's output."
)
class CreateCompletionRequest(BaseModel):
prompt: Union[str, List[str]] = Field(
default="", description="The prompt to generate completions for."
)
suffix: Optional[str] = Field(
default=None,
2023-05-07 06:52:20 +00:00
description="A suffix to append to the generated text. If None, no suffix is appended. Useful for chatbots.",
)
max_tokens: int = max_tokens_field
temperature: float = temperature_field
top_p: float = top_p_field
mirostat_mode: int = mirostat_mode_field
mirostat_tau: float = mirostat_tau_field
mirostat_eta: float = mirostat_eta_field
grammar: Optional[str] = None
echo: bool = Field(
default=False,
2023-05-07 06:52:20 +00:00
description="Whether to echo the prompt in the generated text. Useful for chatbots.",
)
2023-05-19 07:15:08 +00:00
stop: Optional[Union[str, List[str]]] = stop_field
stream: bool = stream_field
logprobs: Optional[int] = Field(
default=None,
ge=0,
2023-05-07 06:52:20 +00:00
description="The number of logprobs to generate. If None, no logprobs are generated.",
)
presence_penalty: Optional[float] = presence_penalty_field
frequency_penalty: Optional[float] = frequency_penalty_field
logit_bias: Optional[Dict[str, float]] = Field(None)
2023-06-15 02:08:28 +00:00
logprobs: Optional[int] = Field(None)
# ignored or currently unsupported
model: Optional[str] = model_field
n: Optional[int] = 1
best_of: Optional[int] = 1
2023-07-14 03:25:12 +00:00
user: Optional[str] = Field(default=None)
# llama.cpp specific parameters
top_k: int = top_k_field
repeat_penalty: float = repeat_penalty_field
2023-06-15 02:08:28 +00:00
logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None)
2023-07-14 03:25:12 +00:00
model_config = {
"json_schema_extra": {
"examples": [
{
"prompt": "\n\n### Instructions:\nWhat is the capital of France?\n\n### Response:\n",
"stop": ["\n", "###"],
}
]
}
2023-07-14 03:25:12 +00:00
}
def make_logit_bias_processor(
llama: llama_cpp.Llama,
logit_bias: Dict[str, float],
logit_bias_type: Optional[Literal["input_ids", "tokens"]],
):
if logit_bias_type is None:
logit_bias_type = "input_ids"
to_bias: Dict[int, float] = {}
if logit_bias_type == "input_ids":
for input_id, score in logit_bias.items():
input_id = int(input_id)
to_bias[input_id] = score
elif logit_bias_type == "tokens":
for token, score in logit_bias.items():
2023-07-14 03:25:12 +00:00
token = token.encode("utf-8")
for input_id in llama.tokenize(token, add_bos=False, special=True):
to_bias[input_id] = score
def logit_bias_processor(
input_ids: npt.NDArray[np.intc],
scores: npt.NDArray[np.single],
) -> npt.NDArray[np.single]:
new_scores = np.copy(scores) # Does it make sense to copy the whole array or can we just overwrite the original one?
for input_id, score in to_bias.items():
new_scores[input_id] = score + scores[input_id]
return new_scores
return logit_bias_processor
2023-05-02 02:38:46 +00:00
@router.post(
"/v1/completions",
)
@router.post("/v1/engines/copilot-codex/completions")
2023-05-27 13:12:58 +00:00
async def create_completion(
request: Request,
body: CreateCompletionRequest,
llama: llama_cpp.Llama = Depends(get_llama),
2023-07-14 03:25:12 +00:00
) -> llama_cpp.Completion:
2023-05-27 13:12:58 +00:00
if isinstance(body.prompt, list):
assert len(body.prompt) <= 1
body.prompt = body.prompt[0] if len(body.prompt) > 0 else ""
exclude = {
"n",
"best_of",
"logit_bias",
"logit_bias_type",
2023-05-27 13:12:58 +00:00
"user",
}
2023-07-14 03:25:12 +00:00
kwargs = body.model_dump(exclude=exclude)
if body.logit_bias is not None:
2023-07-19 07:48:27 +00:00
kwargs["logits_processor"] = llama_cpp.LogitsProcessorList(
[
make_logit_bias_processor(llama, body.logit_bias, body.logit_bias_type),
]
)
if body.grammar is not None:
kwargs["grammar"] = llama_cpp.LlamaGrammar.from_string(body.grammar)
2023-09-14 01:23:23 +00:00
iterator_or_completion: Union[
llama_cpp.Completion, Iterator[llama_cpp.CompletionChunk]
] = await run_in_threadpool(llama, **kwargs)
2023-05-27 13:12:58 +00:00
2023-07-16 05:57:39 +00:00
if isinstance(iterator_or_completion, Iterator):
# EAFP: It's easier to ask for forgiveness than permission
first_response = await run_in_threadpool(next, iterator_or_completion)
2023-07-16 05:57:39 +00:00
# If no exception was raised from first_response, we can assume that
# the iterator is valid and we can use it to stream the response.
def iterator() -> Iterator[llama_cpp.CompletionChunk]:
yield first_response
yield from iterator_or_completion
send_chan, recv_chan = anyio.create_memory_object_stream(10)
2023-05-27 13:12:58 +00:00
return EventSourceResponse(
2023-09-14 01:23:23 +00:00
recv_chan,
data_sender_callable=partial( # type: ignore
2023-07-16 05:57:39 +00:00
get_event_publisher,
request=request,
inner_send_chan=send_chan,
iterator=iterator(),
2023-09-14 01:23:23 +00:00
),
2023-07-16 05:57:39 +00:00
)
2023-05-27 13:12:58 +00:00
else:
2023-07-16 05:57:39 +00:00
return iterator_or_completion
class CreateEmbeddingRequest(BaseModel):
model: Optional[str] = model_field
input: Union[str, List[str]] = Field(description="The input to embed.")
2023-07-14 03:25:12 +00:00
user: Optional[str] = Field(default=None)
model_config = {
"json_schema_extra": {
"examples": [
{
"input": "The food was delicious and the waiter...",
}
]
}
2023-07-14 03:25:12 +00:00
}
2023-05-02 02:38:46 +00:00
@router.post(
"/v1/embeddings",
)
2023-05-27 13:12:58 +00:00
async def create_embedding(
request: CreateEmbeddingRequest, llama: llama_cpp.Llama = Depends(get_llama)
):
2023-05-27 13:12:58 +00:00
return await run_in_threadpool(
2023-07-14 03:25:12 +00:00
llama.create_embedding, **request.model_dump(exclude={"user"})
2023-05-27 13:12:58 +00:00
)
class ChatCompletionRequestMessage(BaseModel):
role: Literal["system", "user", "assistant", "function"] = Field(
default="user", description="The role of the message."
)
content: Optional[str] = Field(default="", description="The content of the message.")
from typing import Any
class CreateChatCompletionRequest(BaseModel):
messages: List[Any] = Field(
2023-05-07 06:52:20 +00:00
default=[], description="A list of messages to generate completions for."
)
2023-07-19 07:48:20 +00:00
functions: Optional[List[llama_cpp.ChatCompletionFunction]] = Field(
default=None,
description="A list of functions to apply to the generated completions.",
)
function_call: Optional[Union[Literal["auto", "none"], llama_cpp.ChatCompletionFunctionCallOption]] = Field(
2023-07-19 07:48:20 +00:00
default=None,
description="A function to apply to the generated completions.",
)
max_tokens: int = max_tokens_field
temperature: float = temperature_field
top_p: float = top_p_field
mirostat_mode: int = mirostat_mode_field
mirostat_tau: float = mirostat_tau_field
mirostat_eta: float = mirostat_eta_field
grammar: Optional[str] = None
stop: Optional[List[str]] = stop_field
stream: bool = stream_field
presence_penalty: Optional[float] = presence_penalty_field
frequency_penalty: Optional[float] = frequency_penalty_field
logit_bias: Optional[Dict[str, float]] = Field(None)
# ignored or currently unsupported
model: Optional[str] = model_field
n: Optional[int] = 1
user: Optional[str] = Field(None)
# llama.cpp specific parameters
top_k: int = top_k_field
repeat_penalty: float = repeat_penalty_field
2023-06-15 02:08:28 +00:00
logit_bias_type: Optional[Literal["input_ids", "tokens"]] = Field(None)
2023-07-14 03:25:12 +00:00
model_config = {
"json_schema_extra": {
"examples": [
{
"messages": [
ChatCompletionRequestMessage(
role="system", content="You are a helpful assistant."
).model_dump(),
ChatCompletionRequestMessage(
role="user", content="What is the capital of France?"
).model_dump(),
]
}
]
}
2023-07-14 03:25:12 +00:00
}
2023-05-02 02:38:46 +00:00
@router.post(
"/v1/chat/completions",
)
2023-05-27 13:12:58 +00:00
async def create_chat_completion(
request: Request,
body: CreateChatCompletionRequest,
llama: llama_cpp.Llama = Depends(get_llama),
settings: Settings = Depends(get_settings),
2023-07-14 03:25:12 +00:00
) -> llama_cpp.ChatCompletion:
2023-05-27 13:12:58 +00:00
exclude = {
"n",
"logit_bias",
"logit_bias_type",
2023-05-27 13:12:58 +00:00
"user",
}
2023-07-14 03:25:12 +00:00
kwargs = body.model_dump(exclude=exclude)
if body.logit_bias is not None:
2023-07-19 07:48:27 +00:00
kwargs["logits_processor"] = llama_cpp.LogitsProcessorList(
[
make_logit_bias_processor(llama, body.logit_bias, body.logit_bias_type),
]
)
if body.grammar is not None:
kwargs["grammar"] = llama_cpp.LlamaGrammar.from_string(body.grammar)
2023-09-14 01:23:23 +00:00
iterator_or_completion: Union[
llama_cpp.ChatCompletion, Iterator[llama_cpp.ChatCompletionChunk]
] = await run_in_threadpool(llama.create_chat_completion, **kwargs)
2023-05-27 13:12:58 +00:00
2023-07-16 05:57:39 +00:00
if isinstance(iterator_or_completion, Iterator):
# EAFP: It's easier to ask for forgiveness than permission
first_response = await run_in_threadpool(next, iterator_or_completion)
# If no exception was raised from first_response, we can assume that
# the iterator is valid and we can use it to stream the response.
def iterator() -> Iterator[llama_cpp.ChatCompletionChunk]:
yield first_response
yield from iterator_or_completion
2023-07-16 05:57:39 +00:00
send_chan, recv_chan = anyio.create_memory_object_stream(10)
return EventSourceResponse(
2023-09-14 01:23:23 +00:00
recv_chan,
data_sender_callable=partial( # type: ignore
2023-07-16 05:57:39 +00:00
get_event_publisher,
request=request,
inner_send_chan=send_chan,
iterator=iterator(),
2023-09-14 01:23:23 +00:00
),
)
2023-07-16 05:57:39 +00:00
else:
return iterator_or_completion
class ModelData(TypedDict):
id: str
object: Literal["model"]
owned_by: str
permissions: List[str]
class ModelList(TypedDict):
object: Literal["list"]
data: List[ModelData]
@router.get("/v1/models")
2023-05-27 13:12:58 +00:00
async def get_models(
settings: Settings = Depends(get_settings),
2023-05-08 00:17:52 +00:00
) -> ModelList:
assert llama is not None
return {
"object": "list",
"data": [
{
"id": settings.model_alias
if settings.model_alias is not None
else llama.model_path,
"object": "model",
"owned_by": "me",
"permissions": [],
}
],
}