| import argparse |
| import uvicorn |
| import sys |
|
|
| from fastapi import FastAPI |
| from pydantic import BaseModel, Field |
| from sse_starlette.sse import EventSourceResponse |
| from utils.logger import logger |
| from networks.message_streamer import MessageStreamer |
| from messagers.message_composer import MessageComposer |
|
|
|
|
| class ChatAPIApp: |
| def __init__(self): |
| self.app = FastAPI( |
| docs_url="/", |
| title="HuggingFace LLM API", |
| swagger_ui_parameters={"defaultModelsExpandDepth": -1}, |
| version="1.0", |
| ) |
| self.setup_routes() |
|
|
| def get_available_models(self): |
| self.available_models = [ |
| { |
| "id": "mixtral-8x7b", |
| "description": "Mixtral-8x7B: https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1", |
| }, |
| ] |
| return self.available_models |
|
|
| class ChatCompletionsPostItem(BaseModel): |
| model: str = Field( |
| default="mixtral-8x7b", |
| description="(str) `mixtral-8x7b`", |
| ) |
| messages: list = Field( |
| default=[{"role": "user", "content": "Hello, who are you?"}], |
| description="(list) Messages", |
| ) |
| temperature: float = Field( |
| default=0.01, |
| description="(float) Temperature", |
| ) |
| max_tokens: int = Field( |
| default=32000, |
| description="(int) Max tokens", |
| ) |
| stream: bool = Field( |
| default=True, |
| description="(bool) Stream", |
| ) |
|
|
| def chat_completions(self, item: ChatCompletionsPostItem): |
| streamer = MessageStreamer(model=item.model) |
| composer = MessageComposer(model=item.model) |
| composer.merge(messages=item.messages) |
| return EventSourceResponse( |
| streamer.chat( |
| prompt=composer.merged_str, |
| temperature=item.temperature, |
| max_new_tokens=item.max_tokens, |
| stream=item.stream, |
| yield_output=True, |
| ), |
| media_type="text/event-stream", |
| ) |
|
|
| def setup_routes(self): |
| for prefix in ["", "/v1"]: |
| self.app.get( |
| prefix + "/models", |
| summary="Get available models", |
| )(self.get_available_models) |
|
|
| self.app.post( |
| prefix + "/chat/completions", |
| summary="Chat completions in conversation session", |
| )(self.chat_completions) |
|
|
|
|
| class ArgParser(argparse.ArgumentParser): |
| def __init__(self, *args, **kwargs): |
| super(ArgParser, self).__init__(*args, **kwargs) |
|
|
| self.add_argument( |
| "-s", |
| "--server", |
| type=str, |
| default="0.0.0.0", |
| help="Server IP for HF LLM Chat API", |
| ) |
| self.add_argument( |
| "-p", |
| "--port", |
| type=int, |
| default=23333, |
| help="Server Port for HF LLM Chat API", |
| ) |
|
|
| self.add_argument( |
| "-d", |
| "--dev", |
| default=False, |
| action="store_true", |
| help="Run in dev mode", |
| ) |
|
|
| self.args = self.parse_args(sys.argv[1:]) |
|
|
|
|
| app = ChatAPIApp().app |
|
|
| if __name__ == "__main__": |
| args = ArgParser().args |
| if args.dev: |
| uvicorn.run("__main__:app", host=args.server, port=args.port, reload=True) |
| else: |
| uvicorn.run("__main__:app", host=args.server, port=args.port, reload=False) |
|
|
| |
| |
|
|
