You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
24 lines
1.1 KiB
24 lines
1.1 KiB
import os
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = "3"
|
|
from vllm import LLM, SamplingParams
|
|
|
|
# Sample prompts.
|
|
prompts = [
|
|
"You are a helpful assistant.\n\nUser:张亮的爸爸叫张明，张明的爸爸有三个孩子，大儿子叫张大，二儿子叫张昊，三儿子叫什么？\nAssistant:",
|
|
"You are a helpful assistant.\n\nUser:你好\nAssistant:",
|
|
"You are a helpful assistant.\n\nUser:1+1等于几\nAssistant:",
|
|
"You are a helpful assistant.\n\nUser:你是谁\nAssistant:",
|
|
]
|
|
# Create a sampling params object.
|
|
sampling_params = SamplingParams(temperature=0, top_p=1, presence_penalty=0.9, max_tokens=1024)
|
|
|
|
# Create an LLM.
|
|
llm = LLM(model="/home/majiahui/project/models-llm/openbuddy-mistral-7b-v13.1", trust_remote_code=True)
|
|
# Generate texts from the prompts. The output is a list of RequestOutput objects
|
|
# that contain the prompt, generated text, and other information.
|
|
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}")
|
|
|