wehub-resource-sync
50 lines
1.2 KiB
Python
50 lines
1.2 KiB
Python
"""Batch inference with SGLang using Ray Data.
|
|
|
|
Usage:
|
|
RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES=0 python batch_sglang_example.py
|
|
"""
|
|
|
|
import ray
|
|
from ray.data.llm import SGLangEngineProcessorConfig, build_processor
|
|
|
|
config = SGLangEngineProcessorConfig(
|
|
model_source="unsloth/Llama-3.1-8B-Instruct",
|
|
engine_kwargs=dict(
|
|
dtype="half",
|
|
mem_fraction_static=0.8,
|
|
),
|
|
batch_size=32,
|
|
concurrency=1,
|
|
)
|
|
|
|
processor = build_processor(
|
|
config,
|
|
preprocess=lambda row: dict(
|
|
messages=[
|
|
{"role": "system", "content": "You are a helpful assistant."},
|
|
{"role": "user", "content": row["prompt"]},
|
|
],
|
|
sampling_params=dict(
|
|
temperature=0.7,
|
|
max_new_tokens=256,
|
|
),
|
|
),
|
|
postprocess=lambda row: dict(
|
|
prompt=row["prompt"],
|
|
response=row["generated_text"],
|
|
),
|
|
)
|
|
|
|
ds = ray.data.from_items(
|
|
[
|
|
{"prompt": "What is the capital of France?"},
|
|
{"prompt": "Explain photosynthesis in one sentence."},
|
|
{"prompt": "Write a haiku about programming."},
|
|
]
|
|
)
|
|
ds = processor(ds)
|
|
|
|
for row in ds.take_all():
|
|
print(f"Prompt: {row['prompt']}")
|
|
print(f"Response: {row['response']}\n")
|