Skip to content

Commit

Permalink
fix input_ids
Browse files Browse the repository at this point in the history
  • Loading branch information
lvhan028 committed Nov 24, 2023
1 parent 20aecd8 commit 171c301
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions benchmark/profile_generation.py
Original file line number Diff line number Diff line change
Expand Up @@ -113,8 +113,10 @@ def profile_throughput(model_path: str,
# make up a prompt that can be tokenized into {input_seqlen} tokens
assert input_seqlen > 0, 'input_seqlen should > 0'
prompt = 'hi'
input_ids = tokenizer.encode(prompt)
input_ids = tokenizer.encode(prompt, add_bos=False)
input_ids = input_ids * input_seqlen
assert len(input_ids) == input_seqlen, \
'#input_token {input_seqlen} but #dummy_input_token {len(input_ids)}'

warmup(tm_model, concurrency, input_ids, output_seqlen)

Expand Down Expand Up @@ -294,7 +296,7 @@ def parse_args():
type=int,
help='how many tokens to be generated. One-to-one'
'correspondence with prompt-tokens',
default=[128, 2048, 128, 2048])
default=[128, 128, 2048, 128, 2048])
parser.add_argument('--tp', type=int, help='Tensor parallel', default=1)
parser.add_argument('--top_k',
type=int,
Expand Down

0 comments on commit 171c301

Please sign in to comment.