tiny optimize for bench serving (#12553)

This commit is contained in:
Yi Zhang
2025-11-03 14:13:18 -08:00
committed by GitHub
parent 7a21d8b276
commit 21a8fa16ea
+8 -1
View File
@@ -25,6 +25,7 @@ import warnings
from argparse import ArgumentParser from argparse import ArgumentParser
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import datetime from datetime import datetime
from functools import lru_cache
from json import JSONDecodeError from json import JSONDecodeError
from pathlib import Path from pathlib import Path
from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Union
@@ -1484,9 +1485,15 @@ def sample_image_requests(
return dataset return dataset
@lru_cache(maxsize=1)
def get_available_tokens(tokenizer):
"""Get all available token ids from the tokenizer vocabulary."""
return list(tokenizer.get_vocab().values())
def gen_prompt(tokenizer, token_num): def gen_prompt(tokenizer, token_num):
"""Generate a random prompt of specified token length using tokenizer vocabulary.""" """Generate a random prompt of specified token length using tokenizer vocabulary."""
all_available_tokens = list(tokenizer.get_vocab().values()) all_available_tokens = get_available_tokens(tokenizer)
selected_tokens = random.choices(all_available_tokens, k=token_num) selected_tokens = random.choices(all_available_tokens, k=token_num)
return tokenizer.decode(selected_tokens) return tokenizer.decode(selected_tokens)