Skip to content

add RETURN_LIST for tgi_api #742

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 20, 2025
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