ForcePilot/scripts/vllm/main.py
Wenjie Zhang 39ff8a07a9 sync
2024-10-08 22:16:17 +08:00

20 lines
550 B
Python

from vllm import LLM, SamplingParams
llm = LLM(model="/hdd/zwj/models/meta-llama/Meta-Llama-3-8B-Instruct")
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")