54 lines
2.1 KiB
Python
54 lines
2.1 KiB
Python
|
import os
|
|||
|
|
from vllm import LLM
|
|||
|
|
from vllm import SamplingParams
|
|||
|
|
from transformers import AutoTokenizer
|
|||
|
|
|
|||
|
|
os.environ['VLLM_USE_MODELSCOPE'] = 'True'
|
|||
|
|
|
|||
|
|
|
|||
|
|
def get_completion(prompts, model, tokenizer=None, temperature = 1.0, top_p = 0.95, top_k=20, min_p=0,
|
|||
|
|
max_tokens = 2048, max_model_len = 4096):
|
|||
|
|
stop_token_ids = [151645, 151643]
|
|||
|
|
# 创建采样参数。temperature 控制生成文本的多样性，
|
|||
|
|
# top_p 控制核心采样的概率，
|
|||
|
|
# top_k 通过限制候选词的数量来控制生成文本的质量和多样性,
|
|||
|
|
# min_p 通过设置概率阈值来筛选候选词，从而在保证文本质量的同时增加多样性
|
|||
|
|
sampling_params = SamplingParams(temperature=temperature, top_p=top_p,
|
|||
|
|
top_k=top_k, min_p=min_p, max_tokens=max_tokens, stop_token_ids=stop_token_ids)
|
|||
|
|
#初始化vllm推理引擎
|
|||
|
|
llm = LLM(
|
|||
|
|
model=model,
|
|||
|
|
tokenizer=tokenizer,
|
|||
|
|
max_model_len=max_model_len,
|
|||
|
|
gpu_memory_utilization=0.85,
|
|||
|
|
trust_remote_code=True,
|
|||
|
|
enforce_eager=True,
|
|||
|
|
swap_space=2 # 使用2GB交换空间
|
|||
|
|
)
|
|||
|
|
outputs = llm.generate(prompts, sampling_params)
|
|||
|
|
return outputs
|
|||
|
|
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
model = '/home/tong/AIProject/Qwen/Qwen/Qwen3-0.6B'
|
|||
|
|
tokenizer = AutoTokenizer.from_pretrained(model, use_fast=False) #加载分词器
|
|||
|
|
prompt = "给我一个关于大模型的简短介绍"
|
|||
|
|
messages = [
|
|||
|
|
{"role": "user", "content": prompt}
|
|||
|
|
]
|
|||
|
|
text = tokenizer.apply_chat_template(
|
|||
|
|
messages,
|
|||
|
|
tokenize=False,
|
|||
|
|
add_generation_prompt=True,
|
|||
|
|
enable_thinking=False)
|
|||
|
|
|
|||
|
|
outputs = get_completion(text, model, tokenizer=None, temperature=0.6, top_p = 0.95, top_k=20, min_p=0) # 对于思考模式，官方建议使用以下参数：temperature = 0.6，TopP = 0.95，TopK = 20，MinP = 0。
|
|||
|
|
|
|||
|
|
# 输出是一个包含 prompt、生成文本和其他信息的 RequestOutput 对象列表。
|
|||
|
|
# 打印输出。
|
|||
|
|
for output in outputs:
|
|||
|
|
prompt = output.prompt
|
|||
|
|
generated_text = output.outputs[0].text
|
|||
|
|
print(f"Prompt: {prompt!r}, \nResponse: {generated_text!r}")
|