Skip to content

Commit

Permalink
add RETURN_LIST for tgi_api (#742)
Browse files Browse the repository at this point in the history
Co-authored-by: shihaobai <[email protected]>
  • Loading branch information
shihaobai and shihaobai authored Feb 20, 2025
1 parent 57bb6e1 commit 9313a08
Showing 1 changed file with 7 additions and 1 deletion.
8 changes: 7 additions & 1 deletion lightllm/server/api_tgi.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import os
import collections
from typing import AsyncGenerator
from fastapi import BackgroundTasks, Request
Expand All @@ -8,6 +9,8 @@
from .httpserver.manager import HttpServerManager
import json

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


def format_tgi_params(params, num_beam: int = 1):
"""
Expand Down Expand Up @@ -132,7 +135,10 @@ async def tgi_generate_impl(request: Request, httpserver_manager: HttpServerMana
if return_details:
ret["details"]["beam_sequences"] = beam_sequences
# wrap generation inside a Vec to match api-inference
json_compatible_item_data = jsonable_encoder([ret])
if RETURN_LIST:
json_compatible_item_data = jsonable_encoder([ret])
else:
json_compatible_item_data = jsonable_encoder(ret)
return JSONResponse(content=json_compatible_item_data)


Expand Down

0 comments on commit 9313a08

Please sign in to comment.