Skip to content

feature: allow custom tokenizer in benchmark runs #96

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.

Sign up for GitHub

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feature: allow custom tokenizer in benchmark runs #96

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
20 changes: 16 additions & 4 deletions token_benchmark_ray.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,11 @@
)
from tqdm import tqdm

from transformers import LlamaTokenizerFast
from transformers import AutoTokenizer, LlamaTokenizerFast

def get_token_throughput_latencies(
model: str,
tokenizer_name_or_path: str,
mean_input_tokens: int,
stddev_input_tokens: int,
mean_output_tokens: int,
Expand All @@ -42,6 +43,7 @@ def get_token_throughput_latencies(

Args:
model: The name of the model to query.
tokenizer_name_or_path: The Hugging Face name or local path of the tokenizer to use for tokenizing prompts and completions.
mean_input_tokens: The mean number of tokens to send in the prompt for the request.
stddev_input_tokens: The standard deviation of the number of tokens to send in the prompt for the request.
mean_output_tokens: The mean number of tokens to generate per request.
Expand All @@ -60,9 +62,12 @@ def get_token_throughput_latencies(
"""
random.seed(11111)

tokenizer = LlamaTokenizerFast.from_pretrained(
"hf-internal-testing/llama-tokenizer"
)
if tokenizer_name_or_path:
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
else:
tokenizer = LlamaTokenizerFast.from_pretrained(
"hf-internal-testing/llama-tokenizer"
)
get_token_length = lambda text: len(tokenizer.encode(text))

if not additional_sampling_params:
Expand Down Expand Up @@ -282,6 +287,7 @@ def flatten(item):
def run_token_benchmark(
llm_api: str,
model: str,
tokenizer_name_or_path: str,
test_timeout_s: int,
max_num_completed_requests: int,
num_concurrent_requests: int,
Expand All @@ -297,6 +303,7 @@ def run_token_benchmark(
Args:
llm_api: The name of the llm api to use.
model: The name of the model to query.
tokenizer_name_or_path: The Hugging Face name or local path of the tokenizer to use for tokenizing prompts and completions.
max_num_completed_requests: The number of requests to complete before finishing the test.
test_timeout_s: The amount of time to run the test for before reporting results.
num_concurrent_requests: The number of concurrent requests to make. Increase
Expand All @@ -318,6 +325,7 @@ def run_token_benchmark(

summary, individual_responses = get_token_throughput_latencies(
model=model,
tokenizer_name_or_path=tokenizer_name_or_path,
llm_api=llm_api,
test_timeout_s=test_timeout_s,
max_num_completed_requests=max_num_completed_requests,
Expand Down Expand Up @@ -368,6 +376,9 @@ def run_token_benchmark(
args.add_argument(
"--model", type=str, required=True, help="The model to use for this load test."
)
args.add_argument(
"--tokenizer", type=str, required=False, default="", help="The tokenizer to use for this load test."
)
args.add_argument(
"--mean-input-tokens",
type=int,
Expand Down Expand Up @@ -478,6 +489,7 @@ def run_token_benchmark(
run_token_benchmark(
llm_api=args.llm_api,
model=args.model,
tokenizer_name_or_path=args.tokenizer,
test_timeout_s=args.timeout,
max_num_completed_requests=args.max_num_completed_requests,
mean_input_tokens=args.mean_input_tokens,
Expand Down