chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,635 @@
|
||||
#!/usr/bin/env python
|
||||
"""
|
||||
Benchmark Ray Data LLM offline batch inference throughput.
|
||||
|
||||
Sample usage:
|
||||
python ray.llm._internal.batch.benchmark.benchmark_processor --mode vllm_engine --batch-size 64 --concurrency 1 --num-prompts 10000 --model facebook/opt-1.3b
|
||||
--tensor-parallel-size 2 --pipeline-parallel-size 2 --distributed-executor-backend ray
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from time import perf_counter, sleep
|
||||
|
||||
import ray
|
||||
from .dataset import ShareGPTDataset
|
||||
from ray import data, serve
|
||||
from ray.data.llm import (
|
||||
ChatTemplateStageConfig,
|
||||
DetokenizeStageConfig,
|
||||
ServeDeploymentProcessorConfig,
|
||||
TokenizerStageConfig,
|
||||
build_processor,
|
||||
vLLMEngineProcessorConfig,
|
||||
)
|
||||
from ray.serve.llm import (
|
||||
LLMConfig,
|
||||
ModelLoadingConfig,
|
||||
build_llm_deployment,
|
||||
)
|
||||
from ray.serve.llm.openai_api_models import CompletionRequest
|
||||
|
||||
|
||||
class Mode(Enum):
|
||||
"""Processor to benchmark."""
|
||||
|
||||
VLLM_ENGINE = "vllm_engine"
|
||||
SHARED_VLLM_ENGINE = "shared_vllm_engine"
|
||||
SERVE_DEPLOYMENT = "serve_deployment"
|
||||
SHARED_SERVE_DEPLOYMENT = "shared_serve_deployment"
|
||||
CLASSIFY = "classify"
|
||||
|
||||
|
||||
# Default sampling parameters -- ensure a fair comparison by omitting sampling-induced variance
|
||||
VLLM_SAMPLING_PARAMS = {
|
||||
"top_p": 1.0,
|
||||
"temperature": 1.0,
|
||||
"max_tokens": 100,
|
||||
"ignore_eos": True,
|
||||
}
|
||||
|
||||
# Default vLLM engine kwargs
|
||||
VLLM_ENGINE_KWARGS = {
|
||||
"max_num_batched_tokens": 4096,
|
||||
}
|
||||
|
||||
# Default tokenization kwargs for classification -- truncate to max_model_len.
|
||||
CLASSIFY_TOKENIZATION_KWARGS_DEFAULT = {"truncation": True, "max_length": 512}
|
||||
|
||||
|
||||
def build_vllm_engine_kwargs(**kwargs) -> dict:
|
||||
"""Build vLLM engine kwargs from command line arguments."""
|
||||
engine_kwargs = VLLM_ENGINE_KWARGS.copy()
|
||||
engine_kwargs.update({k: v for k, v in kwargs.items() if v is not None})
|
||||
return engine_kwargs
|
||||
|
||||
|
||||
def _build_vllm_engine_config(
|
||||
model: str,
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
pipeline_parallel_size: int = None,
|
||||
tensor_parallel_size: int = None,
|
||||
distributed_executor_backend: str = None,
|
||||
task_type: str = None,
|
||||
max_model_len: int = None,
|
||||
) -> vLLMEngineProcessorConfig:
|
||||
"""Helper to create vLLMEngineProcessorConfig."""
|
||||
engine_kwargs = build_vllm_engine_kwargs(
|
||||
pipeline_parallel_size=pipeline_parallel_size,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
)
|
||||
if max_model_len is not None:
|
||||
engine_kwargs["max_model_len"] = max_model_len
|
||||
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source=model,
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
chat_template_stage=False,
|
||||
tokenize_stage=False,
|
||||
detokenize_stage=False,
|
||||
engine_kwargs=engine_kwargs,
|
||||
)
|
||||
|
||||
if task_type is not None:
|
||||
config.task_type = task_type
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _build_serve_deployment_config(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
deployment_name: str = None,
|
||||
app_name: str = None,
|
||||
) -> ServeDeploymentProcessorConfig:
|
||||
"""Helper to create ServeDeploymentProcessorConfig."""
|
||||
return ServeDeploymentProcessorConfig(
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
dtype_mapping={
|
||||
"CompletionRequest": CompletionRequest,
|
||||
},
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class BenchmarkResult:
|
||||
mode: Mode
|
||||
batch_size: int
|
||||
concurrency: int
|
||||
samples: int
|
||||
elapsed_s: float
|
||||
|
||||
@property
|
||||
def throughput(self) -> float:
|
||||
return self.samples / self.elapsed_s if self.elapsed_s else 0.0
|
||||
|
||||
def show(self) -> None:
|
||||
print("\n" + "=" * 60)
|
||||
print(f"BENCHMARK - {self.mode}")
|
||||
print("=" * 60)
|
||||
print(f"Samples : {self.samples}")
|
||||
print(f"Batch size : {self.batch_size}")
|
||||
print(f"Concurrency : {self.concurrency}")
|
||||
print(f"Time (s) : {self.elapsed_s:.2f}")
|
||||
print(f"Throughput : {self.throughput:.2f} req/s")
|
||||
print("=" * 60)
|
||||
|
||||
|
||||
def build_single_vllm_engine_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
pipeline_parallel_size: int = None,
|
||||
tensor_parallel_size: int = None,
|
||||
distributed_executor_backend: str = None,
|
||||
):
|
||||
"""Build vLLM engine processor for single-turn benchmark."""
|
||||
config = _build_vllm_engine_config(
|
||||
model,
|
||||
batch_size,
|
||||
concurrency,
|
||||
pipeline_parallel_size,
|
||||
tensor_parallel_size,
|
||||
distributed_executor_backend,
|
||||
)
|
||||
return build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
prompt=row["prompt"],
|
||||
sampling_params=sampling_params,
|
||||
),
|
||||
postprocess=lambda row: row,
|
||||
)
|
||||
|
||||
|
||||
def build_shared_vllm_engine_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
pipeline_parallel_size: int = None,
|
||||
tensor_parallel_size: int = None,
|
||||
distributed_executor_backend: str = None,
|
||||
):
|
||||
"""Build vLLM engine processor for multi-turn benchmark."""
|
||||
config = _build_vllm_engine_config(
|
||||
model,
|
||||
batch_size,
|
||||
concurrency,
|
||||
pipeline_parallel_size,
|
||||
tensor_parallel_size,
|
||||
distributed_executor_backend,
|
||||
)
|
||||
|
||||
processor1 = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
prompt=row["prompt"],
|
||||
sampling_params=sampling_params,
|
||||
),
|
||||
postprocess=lambda row: {
|
||||
"prompt": row["generated_text"]
|
||||
if str(row.get("generated_text", "")).strip()
|
||||
else row["prompt"]
|
||||
},
|
||||
)
|
||||
|
||||
processor2 = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
prompt=row["prompt"],
|
||||
sampling_params=sampling_params,
|
||||
),
|
||||
postprocess=lambda row: row,
|
||||
)
|
||||
|
||||
def multi_turn_processor(dataset):
|
||||
return processor2(processor1(dataset))
|
||||
|
||||
return multi_turn_processor
|
||||
|
||||
|
||||
def build_classify_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
tokenization_kwargs: dict = CLASSIFY_TOKENIZATION_KWARGS_DEFAULT,
|
||||
max_model_len: int = 512,
|
||||
distributed_executor_backend: str = None,
|
||||
):
|
||||
"""Build vLLM engine processor for classification benchmark."""
|
||||
|
||||
engine_kwargs = VLLM_ENGINE_KWARGS.copy()
|
||||
if distributed_executor_backend is not None:
|
||||
engine_kwargs["distributed_executor_backend"] = distributed_executor_backend
|
||||
|
||||
# Truncate prompts to max_model_len to avoid errors on long inputs.
|
||||
tokenization_kwargs = {**tokenization_kwargs, "max_length": max_model_len}
|
||||
|
||||
config = vLLMEngineProcessorConfig(
|
||||
model_source=model,
|
||||
task_type="classify",
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
chat_template_stage=ChatTemplateStageConfig(enabled=False),
|
||||
tokenize_stage=TokenizerStageConfig(enabled=True),
|
||||
detokenize_stage=DetokenizeStageConfig(enabled=False),
|
||||
engine_kwargs=engine_kwargs,
|
||||
)
|
||||
return build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
prompt=row["prompt"],
|
||||
tokenization_kwargs=tokenization_kwargs,
|
||||
),
|
||||
postprocess=lambda row: {
|
||||
"probs": float(row["embeddings"][0])
|
||||
if row.get("embeddings") is not None and len(row["embeddings"]) > 0
|
||||
else None,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def setup_serve_deployment(model: str, concurrency: int) -> tuple[str, str]:
|
||||
"""Set up Ray Serve deployment for hosting the LLM model."""
|
||||
deployment_name = "benchmark_deployment"
|
||||
app_name = "benchmark_app"
|
||||
|
||||
llm_config = LLMConfig(
|
||||
model_loading_config=ModelLoadingConfig(
|
||||
model_id=model,
|
||||
model_source=model,
|
||||
),
|
||||
deployment_config=dict(
|
||||
name=deployment_name,
|
||||
# To fairly compare with vLLM engine processor, fix the number of replicas to the concurrency level
|
||||
autoscaling_config=dict(
|
||||
min_replicas=concurrency,
|
||||
max_replicas=concurrency,
|
||||
),
|
||||
),
|
||||
engine_kwargs=dict(
|
||||
enable_prefix_caching=True,
|
||||
enable_chunked_prefill=True,
|
||||
max_num_batched_tokens=4096,
|
||||
),
|
||||
)
|
||||
|
||||
override_serve_options = dict(name=deployment_name)
|
||||
llm_app = build_llm_deployment(
|
||||
llm_config, override_serve_options=override_serve_options
|
||||
)
|
||||
serve.run(llm_app, name=app_name)
|
||||
|
||||
print("Waiting for Serve deployment to be ready...")
|
||||
max_wait_time = 120 # seconds
|
||||
wait_time = 0
|
||||
while not _is_app_ready(app_name) and wait_time < max_wait_time:
|
||||
sleep(5)
|
||||
wait_time += 5
|
||||
|
||||
if wait_time >= max_wait_time:
|
||||
raise TimeoutError("Deployment failed to become ready within timeout")
|
||||
|
||||
print("Deployment is ready!")
|
||||
return deployment_name, app_name
|
||||
|
||||
|
||||
def _is_app_ready(app_name: str) -> bool:
|
||||
try:
|
||||
serve_status = serve.status()
|
||||
|
||||
if app_name in serve_status.applications:
|
||||
app_status = serve_status.applications[app_name]
|
||||
if app_status.status == "RUNNING":
|
||||
print(f"Application '{app_name}' is RUNNING.")
|
||||
return True
|
||||
else:
|
||||
print(f"Application '{app_name}' status: {app_status.status}")
|
||||
return False
|
||||
else:
|
||||
print(f"Application '{app_name}' not found in Serve status.")
|
||||
return False
|
||||
except Exception as e:
|
||||
print(f"Error checking app status: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def build_single_serve_deployment_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
deployment_name: str = None,
|
||||
app_name: str = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Build Serve deployment processor for single-turn benchmark."""
|
||||
config = _build_serve_deployment_config(
|
||||
batch_size,
|
||||
concurrency,
|
||||
deployment_name,
|
||||
app_name,
|
||||
)
|
||||
return build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
method="completions",
|
||||
dtype="CompletionRequest",
|
||||
request_kwargs=dict(
|
||||
model=model,
|
||||
prompt=row["prompt"],
|
||||
**sampling_params,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: row,
|
||||
)
|
||||
|
||||
|
||||
def build_shared_serve_deployment_processor(
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
deployment_name: str = None,
|
||||
app_name: str = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Build Serve deployment processor for multi-turn benchmark."""
|
||||
config = _build_serve_deployment_config(
|
||||
batch_size,
|
||||
concurrency,
|
||||
deployment_name,
|
||||
app_name,
|
||||
)
|
||||
|
||||
processor1 = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
method="completions",
|
||||
dtype="CompletionRequest",
|
||||
request_kwargs=dict(
|
||||
model=model,
|
||||
prompt=row["prompt"],
|
||||
stream=False,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: {
|
||||
# Fall back to original prompt if generated text is empty
|
||||
"prompt": (
|
||||
row["choices"][0]["text"]
|
||||
if row.get("choices") and str(row["choices"][0].get("text", "")).strip()
|
||||
else row["prompt"]
|
||||
)
|
||||
},
|
||||
)
|
||||
|
||||
processor2 = build_processor(
|
||||
config,
|
||||
preprocess=lambda row: dict(
|
||||
method="completions",
|
||||
dtype="CompletionRequest",
|
||||
request_kwargs=dict(
|
||||
model=model,
|
||||
prompt=row["prompt"],
|
||||
stream=False,
|
||||
),
|
||||
),
|
||||
postprocess=lambda row: row,
|
||||
)
|
||||
|
||||
def multi_turn_processor(dataset):
|
||||
return processor2(processor1(dataset))
|
||||
|
||||
return multi_turn_processor
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Benchmark execution
|
||||
# -----------------------------------------------------------------------------
|
||||
def run_processor(
|
||||
mode: Mode,
|
||||
dataset: data.Dataset,
|
||||
builder,
|
||||
**kwargs,
|
||||
) -> BenchmarkResult:
|
||||
processor = builder(**kwargs)
|
||||
|
||||
total_samples = dataset.count()
|
||||
|
||||
start = perf_counter()
|
||||
processor(dataset).materialize()
|
||||
elapsed = perf_counter() - start
|
||||
|
||||
return BenchmarkResult(
|
||||
mode=mode,
|
||||
batch_size=kwargs.get("batch_size"),
|
||||
concurrency=kwargs.get("concurrency"),
|
||||
samples=total_samples,
|
||||
elapsed_s=elapsed,
|
||||
)
|
||||
|
||||
|
||||
def benchmark(
|
||||
mode: Mode,
|
||||
dataset: data.Dataset,
|
||||
*,
|
||||
batch_size: int,
|
||||
concurrency: int,
|
||||
model: str,
|
||||
sampling_params: dict = VLLM_SAMPLING_PARAMS,
|
||||
pipeline_parallel_size: int = None,
|
||||
tensor_parallel_size: int = None,
|
||||
distributed_executor_backend: str = None,
|
||||
) -> BenchmarkResult:
|
||||
mode_to_builder = {
|
||||
Mode.VLLM_ENGINE: build_single_vllm_engine_processor,
|
||||
Mode.SHARED_VLLM_ENGINE: build_shared_vllm_engine_processor,
|
||||
Mode.SERVE_DEPLOYMENT: build_single_serve_deployment_processor,
|
||||
Mode.SHARED_SERVE_DEPLOYMENT: build_shared_serve_deployment_processor,
|
||||
Mode.CLASSIFY: build_classify_processor,
|
||||
}
|
||||
|
||||
if mode not in mode_to_builder:
|
||||
raise ValueError(f"Unknown benchmark mode: {mode}")
|
||||
|
||||
builder = mode_to_builder[mode]
|
||||
|
||||
if mode in [Mode.SERVE_DEPLOYMENT, Mode.SHARED_SERVE_DEPLOYMENT]:
|
||||
deployment_name, app_name = setup_serve_deployment(model, concurrency)
|
||||
try:
|
||||
return run_processor(
|
||||
mode,
|
||||
dataset,
|
||||
builder,
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
model=model,
|
||||
sampling_params=sampling_params,
|
||||
deployment_name=deployment_name,
|
||||
app_name=app_name,
|
||||
)
|
||||
finally:
|
||||
serve.delete(app_name)
|
||||
elif mode == Mode.CLASSIFY:
|
||||
return run_processor(
|
||||
mode,
|
||||
dataset,
|
||||
builder,
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
model=model,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
)
|
||||
else:
|
||||
return run_processor(
|
||||
mode,
|
||||
dataset,
|
||||
builder,
|
||||
batch_size=batch_size,
|
||||
concurrency=concurrency,
|
||||
model=model,
|
||||
sampling_params=sampling_params,
|
||||
pipeline_parallel_size=pipeline_parallel_size,
|
||||
tensor_parallel_size=tensor_parallel_size,
|
||||
distributed_executor_backend=distributed_executor_backend,
|
||||
)
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# CLI
|
||||
# -----------------------------------------------------------------------------
|
||||
def parse_args(argv: list[str]) -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="vLLM throughput benchmark")
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=[mode.value for mode in Mode],
|
||||
default=Mode.VLLM_ENGINE.value,
|
||||
help="Ray Data LLM processor to run benchmarks for",
|
||||
)
|
||||
# Dataset configuration
|
||||
parser.add_argument(
|
||||
"--dataset-path",
|
||||
type=str,
|
||||
default="/home/ubuntu/datasets/Code-feedback-sharegpt-renamed",
|
||||
help="Path to dataset on disk",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-prompts", type=int, default=1000, help="Number of prompts to process"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-dataset-id",
|
||||
type=str,
|
||||
default="Crystalcareai/Code-feedback-sharegpt-renamed",
|
||||
help="Hugging Face dataset ID to download",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--hf-split",
|
||||
type=str,
|
||||
default="train",
|
||||
help="Hugging Face dataset split to load",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=0,
|
||||
help="Random seed for dataset sampling",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--truncate-prompt",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Maximum prompt length",
|
||||
)
|
||||
# Engine configuration
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
required=True,
|
||||
help="LLM model to use",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pipeline-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Pipeline parallel size for vLLM engine",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Tensor parallel size for vLLM engine",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--distributed-executor-backend",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["ray", "mp", "uni"],
|
||||
help="Distributed executor backend for vLLM engine",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Maximum number of tokens to generate per request (default: 100)",
|
||||
)
|
||||
# Ray Data worker configuration
|
||||
parser.add_argument(
|
||||
"--batch-size",
|
||||
type=int,
|
||||
required=True,
|
||||
help="Ray Data batch size for processing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--concurrency", type=int, required=True, help="Ray Data concurrency level"
|
||||
)
|
||||
return parser.parse_args(argv)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args(sys.argv[1:])
|
||||
|
||||
ray.init()
|
||||
try:
|
||||
dataset = ShareGPTDataset(
|
||||
dataset_path=args.dataset_path,
|
||||
seed=args.seed,
|
||||
hf_dataset_id=args.hf_dataset_id,
|
||||
hf_split=args.hf_split,
|
||||
truncate_prompt=args.truncate_prompt,
|
||||
)
|
||||
prompts = dataset.sample(args.num_prompts)
|
||||
|
||||
dataset = data.from_items(prompts)
|
||||
|
||||
sampling_params = VLLM_SAMPLING_PARAMS.copy()
|
||||
if args.max_tokens is not None:
|
||||
sampling_params["max_tokens"] = args.max_tokens
|
||||
|
||||
result = benchmark(
|
||||
Mode(args.mode),
|
||||
dataset,
|
||||
batch_size=args.batch_size,
|
||||
concurrency=args.concurrency,
|
||||
model=args.model,
|
||||
sampling_params=sampling_params,
|
||||
pipeline_parallel_size=args.pipeline_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
distributed_executor_backend=args.distributed_executor_backend,
|
||||
)
|
||||
result.show()
|
||||
finally:
|
||||
ray.shutdown()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,175 @@
|
||||
"""
|
||||
This module defines a dataset framework for sampling benchmark requests.
|
||||
"""
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from datasets import load_dataset, load_from_disk
|
||||
|
||||
|
||||
class BenchmarkDataset(ABC):
|
||||
DEFAULT_RANDOM_SEED = 0
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_path: Optional[str] = None,
|
||||
random_seed: int = DEFAULT_RANDOM_SEED,
|
||||
) -> None:
|
||||
"""
|
||||
Abstract base class for benchmark datasets.
|
||||
|
||||
All benchmark datasets should inherit from this class and implement
|
||||
the required abstract methods.
|
||||
|
||||
Args:
|
||||
dataset_path: The path to the dataset on disk.
|
||||
random_seed: The seed for the random number generator.
|
||||
"""
|
||||
self._dataset_path = dataset_path
|
||||
self._random_seed = random_seed
|
||||
|
||||
@abstractmethod
|
||||
def load_data(self) -> None:
|
||||
"""
|
||||
Load data from the dataset source into memory.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the method is not implemented in subclasses.
|
||||
"""
|
||||
raise NotImplementedError("load_data must be implemented in subclasses.")
|
||||
|
||||
@abstractmethod
|
||||
def sample(self, num_requests: int) -> List[Dict]:
|
||||
"""
|
||||
Sample prompts from the loaded dataset.
|
||||
|
||||
Args:
|
||||
num_requests: The number of prompts to sample from the dataset.
|
||||
|
||||
Returns:
|
||||
A list of sampled request dictionaries.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: If the method is not implemented in subclasses.
|
||||
"""
|
||||
raise NotImplementedError("sample must be implemented in subclasses.")
|
||||
|
||||
|
||||
class ShareGPTDataset(BenchmarkDataset):
|
||||
"""Implements the ShareGPT dataset. The first human message of each conversation is used to build a prompt."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dataset_path: str,
|
||||
seed: int,
|
||||
hf_dataset_id: str = "Crystalcareai/Code-feedback-sharegpt-renamed",
|
||||
hf_split: str = "train",
|
||||
truncate_prompt: Optional[int] = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initializes the ShareGPTDataset.
|
||||
|
||||
Args:
|
||||
dataset_path: The path to the dataset on disk.
|
||||
seed: The seed for the random number generator.
|
||||
hf_dataset_id: The Hugging Face dataset ID to download if the dataset is not found on disk.
|
||||
hf_split: The Hugging Face split to load from the dataset.
|
||||
truncate_prompt: Maximum prompt length so that the prompt fits in the model's context window.
|
||||
"""
|
||||
super().__init__(dataset_path, seed)
|
||||
self._seed = seed
|
||||
|
||||
self._hf_dataset_id = hf_dataset_id
|
||||
self._hf_split = hf_split
|
||||
self._truncate_prompt = truncate_prompt
|
||||
|
||||
self._data: list[Dict] | None = None
|
||||
|
||||
def load_data(self) -> None:
|
||||
"""Load data from the dataset path into memory."""
|
||||
if self._data is None:
|
||||
self._data = self._load_dataset_data()
|
||||
|
||||
def sample(self, num_requests: int) -> List[Dict]:
|
||||
"""Sample prompts from the loaded dataset."""
|
||||
if self._data is None:
|
||||
self.load_data()
|
||||
|
||||
# Extract all valid prompts from the dataset
|
||||
all_prompts = []
|
||||
for item in self._data:
|
||||
prompt_data = self._extract_prompt(item)
|
||||
if prompt_data is not None:
|
||||
all_prompts.append(prompt_data)
|
||||
|
||||
if not all_prompts:
|
||||
raise ValueError("ShareGPT dataset yielded no usable prompts")
|
||||
|
||||
# Replicate samples if num_requests exceeds available samples
|
||||
if num_requests <= len(all_prompts):
|
||||
return all_prompts[:num_requests]
|
||||
|
||||
full_copies = num_requests // len(all_prompts)
|
||||
remainder = num_requests % len(all_prompts)
|
||||
prompts = all_prompts * full_copies + all_prompts[:remainder]
|
||||
return prompts
|
||||
|
||||
def _load_dataset(self):
|
||||
"""Load dataset from disk or Hugging Face."""
|
||||
path = Path(self._dataset_path)
|
||||
print(f"Attempting to load dataset from {path}")
|
||||
print(f"Dataset exists on disk: {path.exists()}")
|
||||
|
||||
try:
|
||||
if path.exists():
|
||||
dataset = load_from_disk(str(path))
|
||||
else:
|
||||
print(
|
||||
f"Dataset not found on disk, downloading from Hugging Face: {self._hf_dataset_id}"
|
||||
)
|
||||
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
dataset = load_dataset(self._hf_dataset_id, split=self._hf_split)
|
||||
dataset.save_to_disk(str(path))
|
||||
return dataset
|
||||
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Error loading ShareGPT dataset: {e}")
|
||||
|
||||
def _load_dataset_data(self) -> List[Dict]:
|
||||
"""Load and process dataset data into a list of dictionaries."""
|
||||
ds = self._load_dataset().shuffle(seed=self._seed)
|
||||
data = []
|
||||
|
||||
for i, row in enumerate(ds):
|
||||
data.append(row)
|
||||
|
||||
print(f"Loaded {len(data)} samples from dataset")
|
||||
return data
|
||||
|
||||
def _extract_prompt(self, item: Dict) -> Dict | None:
|
||||
"""
|
||||
Extracts the first human message of a conversation or None.
|
||||
|
||||
The ShareGPT schema uses {"role": "human", "value": ...} for user
|
||||
turns.
|
||||
"""
|
||||
messages = item.get("messages") or item.get("conversations") or []
|
||||
prompt = next(
|
||||
(
|
||||
str(msg.get("value", "")).strip()
|
||||
for msg in messages
|
||||
if msg.get("role") in {"human", "user"}
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
# Only return a valid prompt if it's not empty
|
||||
if prompt and prompt.strip():
|
||||
if self._truncate_prompt:
|
||||
prompt = prompt[: self._truncate_prompt]
|
||||
return {"prompt": prompt}
|
||||
|
||||
return None
|
||||
Reference in New Issue
Block a user