Unverified Commit 2342605e authored by fzyzcjy's avatar fzyzcjy Committed by GitHub
Browse files

Tiny cleanup send_single (#12056)

parent dbf17a83
...@@ -94,18 +94,19 @@ class BenchArgs: ...@@ -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( ...@@ -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): ...@@ -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)
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment