Skip to content

Commit 9313a08

Browse files
shihaobaishihaobai
andauthored
add RETURN_LIST for tgi_api (#742)
Co-authored-by: shihaobai <baishihao@sensetime.com>
1 parent 57bb6e1 commit 9313a08

File tree

1 file changed

+7
-1
lines changed

1 file changed

+7
-1
lines changed

lightllm/server/api_tgi.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import os
12
import collections
23
from typing import AsyncGenerator
34
from fastapi import BackgroundTasks, Request
@@ -8,6 +9,8 @@
89
from .httpserver.manager import HttpServerManager
910
import json
1011

12+
RETURN_LIST = os.getenv("RETURN_LIST", "FALSE").upper() in ["ON", "TRUE", "1"]
13+
1114

1215
def format_tgi_params(params, num_beam: int = 1):
1316
"""
@@ -132,7 +135,10 @@ async def tgi_generate_impl(request: Request, httpserver_manager: HttpServerMana
132135
if return_details:
133136
ret["details"]["beam_sequences"] = beam_sequences
134137
# wrap generation inside a Vec to match api-inference
135-
json_compatible_item_data = jsonable_encoder([ret])
138+
if RETURN_LIST:
139+
json_compatible_item_data = jsonable_encoder([ret])
140+
else:
141+
json_compatible_item_data = jsonable_encoder(ret)
136142
return JSONResponse(content=json_compatible_item_data)
137143

138144

0 commit comments

Comments
 (0)