Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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