llama.cpp/examples/fastapi_server.py

100 lines
2.4 KiB
Python
Raw Normal View History

2023-03-24 23:10:31 +00:00
"""Example FastAPI server for llama.cpp.
"""
import json
from typing import List, Optional, Iterator
2023-03-24 05:41:24 +00:00
import llama_cpp
2023-03-24 05:41:24 +00:00
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, BaseSettings, Field, create_model_from_typeddict
from sse_starlette.sse import EventSourceResponse
2023-03-24 05:41:24 +00:00
2023-03-24 18:35:41 +00:00
2023-03-24 05:41:24 +00:00
class Settings(BaseSettings):
model: str
2023-03-24 18:35:41 +00:00
2023-03-24 05:41:24 +00:00
app = FastAPI(
title="🦙 llama.cpp Python API",
version="0.0.1",
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
2023-03-24 05:41:24 +00:00
settings = Settings()
llama = llama_cpp.Llama(
settings.model,
f16_kv=True,
use_mlock=True,
embedding=True,
n_threads=6,
n_batch=2048,
)
2023-03-24 05:41:24 +00:00
2023-03-24 18:35:41 +00:00
class CreateCompletionRequest(BaseModel):
2023-03-24 05:41:24 +00:00
prompt: str
suffix: Optional[str] = Field(None)
max_tokens: int = 16
temperature: float = 0.8
top_p: float = 0.95
logprobs: Optional[int] = Field(None)
echo: bool = False
stop: List[str] = []
repeat_penalty: float = 1.1
top_k: int = 40
stream: bool = False
2023-03-24 05:41:24 +00:00
class Config:
schema_extra = {
"example": {
"prompt": "\n\n### Instructions:\nWhat is the capital of France?\n\n### Response:\n",
2023-03-24 18:35:41 +00:00
"stop": ["\n", "###"],
2023-03-24 05:41:24 +00:00
}
}
CreateCompletionResponse = create_model_from_typeddict(llama_cpp.Completion)
@app.post(
"/v1/completions",
response_model=CreateCompletionResponse,
)
def create_completion(request: CreateCompletionRequest):
if request.stream:
chunks: Iterator[llama_cpp.CompletionChunk] = llama(**request.dict()) # type: ignore
return EventSourceResponse(dict(data=json.dumps(chunk)) for chunk in chunks)
2023-03-24 18:35:41 +00:00
return llama(**request.dict())
class CreateEmbeddingRequest(BaseModel):
model: Optional[str]
input: str
user: Optional[str]
class Config:
schema_extra = {
"example": {
"input": "The food was delicious and the waiter...",
}
}
CreateEmbeddingResponse = create_model_from_typeddict(llama_cpp.Embedding)
@app.post(
"/v1/embeddings",
response_model=CreateEmbeddingResponse,
)
def create_embedding(request: CreateEmbeddingRequest):
# print(request)
# return llama.create_embedding(**request.dict(exclude={"model", "user"}))
return llama.create_embedding(request.input)