wehub-resource-sync
63 lines
1.6 KiB
Python
63 lines
1.6 KiB
Python
"""
|
|
Quickstart: vLLM + Ray Data batch inference.
|
|
|
|
1. Installation
|
|
2. Dataset creation
|
|
3. Processor configuration
|
|
4. Running inference
|
|
5. Getting results
|
|
"""
|
|
|
|
# __minimal_vllm_quickstart_start__
|
|
import ray
|
|
from ray.data.llm import vLLMEngineProcessorConfig, build_processor
|
|
|
|
# Initialize Ray
|
|
ray.init()
|
|
|
|
# simple dataset
|
|
ds = ray.data.from_items([
|
|
{"prompt": "What is machine learning?"},
|
|
{"prompt": "Explain neural networks in one sentence."},
|
|
])
|
|
|
|
# Minimal vLLM configuration
|
|
config = vLLMEngineProcessorConfig(
|
|
model_source="unsloth/Llama-3.1-8B-Instruct",
|
|
concurrency=1, # 1 vLLM engine replica
|
|
batch_size=32, # 32 samples per batch
|
|
engine_kwargs={
|
|
"max_model_len": 4096, # Fit into test GPU memory
|
|
}
|
|
)
|
|
|
|
# Build processor
|
|
# preprocess: converts input row to format expected by vLLM (OpenAI chat format)
|
|
# postprocess: extracts generated text from vLLM output
|
|
processor = build_processor(
|
|
config,
|
|
preprocess=lambda row: {
|
|
"messages": [{"role": "user", "content": row["prompt"]}],
|
|
"sampling_params": {"temperature": 0.7, "max_tokens": 100},
|
|
},
|
|
postprocess=lambda row: {
|
|
"prompt": row["prompt"],
|
|
"response": row["generated_text"],
|
|
},
|
|
)
|
|
|
|
# inference
|
|
ds = processor(ds)
|
|
|
|
# iterate through the results
|
|
for result in ds.iter_rows():
|
|
print(f"Q: {result['prompt']}")
|
|
print(f"A: {result['response']}\n")
|
|
|
|
# Alternative ways to get results:
|
|
# results = ds.take(10) # Get first 10 results
|
|
# ds.show(limit=5) # Print first 5 results
|
|
# ds.write_parquet("output.parquet") # Save to file
|
|
# __minimal_vllm_quickstart_end__
|
|
|