Update benchmarks to use HF token from environment. (#15421)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Frank
2025-12-18 13:47:27 -08:00
committed by GitHub
parent 2b0ddf89f5
commit 9749d3e346
2 changed files with 34 additions and 8 deletions

View File

@@ -918,7 +918,8 @@ class BenchmarkMetrics:
max_concurrent_requests: int = 0
SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json"
SHAREGPT_REPO_ID = "anon8231489123/ShareGPT_Vicuna_unfiltered"
SHAREGPT_FILENAME = "ShareGPT_V3_unfiltered_cleaned_split.json"
MOONCAKE_DATASET_URL = {
"mooncake": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/arxiv-trace/mooncake_trace.jsonl",
"conversation": "https://raw.githubusercontent.com/kvcache-ai/Mooncake/main/FAST25-release/traces/conversation_trace.jsonl",
@@ -927,6 +928,19 @@ MOONCAKE_DATASET_URL = {
}
def download_and_cache_hf_file(
repo_id: str,
filename: str,
repo_type: str = "dataset",
):
"""Download a file from Hugging Face and cache it locally."""
from huggingface_hub import hf_hub_download
return hf_hub_download(
repo_id=repo_id, filename=filename, repo_type=repo_type
)
def download_and_cache_file(url: str, filename: Optional[str] = None):
"""Read and cache a file from a url."""
if filename is None:
@@ -1171,7 +1185,10 @@ def sample_sharegpt_requests(
# Download sharegpt if necessary
if not is_file_valid_json(dataset_path) and dataset_path == "":
dataset_path = download_and_cache_file(SHAREGPT_URL)
dataset_path = download_and_cache_hf_file(
repo_id=SHAREGPT_REPO_ID,
filename=SHAREGPT_FILENAME,
)
# Load the dataset.
with open(dataset_path) as f:
@@ -1283,7 +1300,10 @@ def sample_random_requests(
# Download sharegpt if necessary
if not is_file_valid_json(dataset_path):
dataset_path = download_and_cache_file(SHAREGPT_URL)
dataset_path = download_and_cache_hf_file(
repo_id=SHAREGPT_REPO_ID,
filename=SHAREGPT_FILENAME,
)
# Load the dataset.
with open(dataset_path) as f: