Spaces:
Running
Running
import argparse | |
import uvicorn | |
import sys | |
import json | |
from fastapi import FastAPI | |
from fastapi.encoders import jsonable_encoder | |
from fastapi.responses import JSONResponse | |
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 | |
from googletrans import Translator | |
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): | |
f = open('apis/lang_name.json', "r") | |
self.available_models = json.loads(f.read()) | |
return self.available_models | |
class ChatCompletionsPostItem(BaseModel): | |
from_language: str = Field( | |
default="auto", | |
description="(str) `Detect`", | |
) | |
to_language: str = Field( | |
default="en", | |
description="(str) `en`", | |
) | |
text: str = Field( | |
default="Hello", | |
description="(str) `Text for translate`", | |
) | |
def chat_completions(self, item: ChatCompletionsPostItem): | |
translator = Translator() | |
f = open('apis/lang_name.json', "r") | |
available_langs = json.loads(f.read()) | |
from_lang = "en" | |
for lang_item in available_langs: | |
if item.to_language == lang_item['code']: | |
from_lang = item.to_language | |
break | |
item_response = { | |
"from_language": item.from_language, | |
"to_language": item.to_language, | |
"text": item.text, | |
"translate": translator.translate(item.text, dest=from_lang) | |
} | |
json_compatible_item_data = jsonable_encoder(item_response) | |
return JSONResponse(content=json_compatible_item_data) | |
def setup_routes(self): | |
for prefix in ["", "/v1"]: | |
self.app.get( | |
prefix + "/models", | |
summary="Get available languages", | |
)(self.get_available_models) | |
self.app.post( | |
prefix + "/translate", | |
summary="translate text", | |
)(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) | |
# python -m apis.chat_api # [Docker] on product mode | |
# python -m apis.chat_api -d # [Dev] on develop mode | |