diff --git a/benchmark/hicache/data_processing.py b/benchmark/hicache/data_processing.py index 1fb3650ce..dd0cbf669 100644 --- a/benchmark/hicache/data_processing.py +++ b/benchmark/hicache/data_processing.py @@ -11,10 +11,10 @@ from nextqa import NExTQALoader from tqdm.asyncio import tqdm from transformers import PreTrainedTokenizerBase -SHAREGPT_URL = "https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json" - from sglang.bench_serving import ( - download_and_cache_file, + SHAREGPT_FILENAME, + SHAREGPT_REPO_ID, + download_and_cache_hf_file, gen_prompt, get_gen_prefix_cache_path, ) @@ -104,7 +104,10 @@ def sample_sharegpt_requests( # Download sharegpt if necessary if not os.path.isfile(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: @@ -367,7 +370,10 @@ def sample_random_requests( # Download sharegpt if necessary if not os.path.isfile(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: diff --git a/python/sglang/bench_serving.py b/python/sglang/bench_serving.py index f0b3b4472..60ccb3b84 100644 --- a/python/sglang/bench_serving.py +++ b/python/sglang/bench_serving.py @@ -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: