Tiny cleanup send_single (#12056)
This commit is contained in:
@@ -94,18 +94,19 @@ class BenchArgs:
|
|||||||
|
|
||||||
def send_single(
|
def send_single(
|
||||||
args,
|
args,
|
||||||
batch_size: int = 1,
|
|
||||||
profile: bool = False,
|
profile: bool = False,
|
||||||
profile_steps: int = 3,
|
profile_steps: int = 3,
|
||||||
profile_by_stage: bool = False,
|
profile_by_stage: bool = False,
|
||||||
return_full_response: bool = False,
|
return_full_response: bool = False,
|
||||||
input_ids: List[int] = None,
|
input_ids: List[int] = None,
|
||||||
|
prompt: List[str] = None,
|
||||||
max_new_tokens: int = None,
|
max_new_tokens: int = None,
|
||||||
):
|
):
|
||||||
base_url = f"http://{args.host}:{args.port}"
|
base_url = f"http://{args.host}:{args.port}"
|
||||||
|
|
||||||
# Use input_ids if provided, otherwise use text prompts
|
# Use input_ids if provided, otherwise use text prompts
|
||||||
if input_ids is not None:
|
if input_ids is not None:
|
||||||
|
assert prompt is None
|
||||||
json_data = {
|
json_data = {
|
||||||
"input_ids": input_ids,
|
"input_ids": input_ids,
|
||||||
"sampling_params": {
|
"sampling_params": {
|
||||||
@@ -122,7 +123,7 @@ def send_single(
|
|||||||
"stream": args.stream,
|
"stream": args.stream,
|
||||||
}
|
}
|
||||||
else:
|
else:
|
||||||
prompt = [PROMPT_1] * batch_size
|
assert input_ids is None
|
||||||
json_data = {
|
json_data = {
|
||||||
"text": prompt,
|
"text": prompt,
|
||||||
"sampling_params": {
|
"sampling_params": {
|
||||||
@@ -225,7 +226,7 @@ def test_deterministic(args):
|
|||||||
texts = []
|
texts = []
|
||||||
for i in range(1, args.n_trials + 1):
|
for i in range(1, args.n_trials + 1):
|
||||||
batch_size = i
|
batch_size = i
|
||||||
text = send_single(args, batch_size, args.profile)
|
text = send_single(args, args.profile, prompt=[PROMPT_1] * batch_size)
|
||||||
text = text.replace("\n", " ")
|
text = text.replace("\n", " ")
|
||||||
print(f"Trial {i} with batch size {batch_size}: {text}")
|
print(f"Trial {i} with batch size {batch_size}: {text}")
|
||||||
texts.append(text)
|
texts.append(text)
|
||||||
|
|||||||
Reference in New Issue
Block a user