Skip to content

Commit

Permalink
standardize prompt formatting for vllm
Browse files Browse the repository at this point in the history
  • Loading branch information
wongjingping committed May 23, 2024
1 parent 16875cf commit e82bd92
Showing 1 changed file with 12 additions and 1 deletion.
13 changes: 12 additions & 1 deletion eval/vllm_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -121,8 +121,19 @@ def chunk_dataframe(df, chunk_size):
for batch in (pbar := tqdm(df_chunks, total=len(df))):
prompts = batch["prompt"].tolist()
print(f"Generating completions for {len(prompts)} prompts")
prompt_tokens = []
prompt_token_sizes = []
for prompt in prompts:
token_ids = tokenizer.encode(prompt, add_special_tokens=False)
# add bos token if not already present in prompt
if token_ids[0] != tokenizer.bos_token_id:
token_ids = [tokenizer.bos_token_id] + token_ids
prompt_tokens.append(token_ids)
prompt_token_sizes.append(len(token_ids))
print(f"Average prompt size: {sum(prompt_token_sizes)/len(prompt_token_sizes):.0f}")
start_time = time.time()
outputs = llm.generate(prompts, sampling_params)
# outputs = llm.generate(prompts, sampling_params) # if you prefer to use prompts instead of token_ids
outputs = llm.generate(sampling_params=sampling_params, prompt_token_ids=prompt_tokens)
print(
f"Generated {len(outputs)} completions in {time.time() - start_time:.2f} seconds"
)
Expand Down

0 comments on commit e82bd92

Please sign in to comment.