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}")