forked from LucienShui/huggingface-vscode-endpoint-server
-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
38 lines (30 loc) · 1.04 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
import uvicorn
from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from generators import GeneratorBase, StarCoder
import json
from util import logger, get_parser
app = FastAPI()
app.add_middleware(
CORSMiddleware
)
generator: GeneratorBase = ...
@app.post("/api/generate/")
async def api(request: Request):
json_request: dict = await request.json()
inputs: str = json_request['inputs']
parameters: dict = json_request['parameters']
logger.info(f'{request.client.host}:{request.client.port} inputs = {json.dumps(inputs)}')
generated_text: str = generator.generate(inputs, parameters)
logger.info(f'{request.client.host}:{request.client.port} generated_text = {json.dumps(generated_text)}')
return {
"generated_text": generated_text,
"status": 200
}
def main():
global generator
args = get_parser().parse_args()
generator = StarCoder(args.pretrained, device_map='auto')
uvicorn.run(app, host=args.host, port=args.port)
if __name__ == '__main__':
main()