from vllm import LLM, SamplingParams

model_path = "../model_weights"

model = LLM(model=model_path,tokenizer=model_path, dtype='bfloat16',seed=1234)
sampling_params = SamplingParams(repetition_penalty = 1.05,
                                 temperature = 0.7,
                                 top_k = 20,
                                 top_p = 0.8,
                                 max_tokens = 512)

messages = [
    {"role": "system", "content": "下午好！这里是曜影医疗预约中心，我是Lisa,请问有什么可以帮您？"}
]

#费用
prompts = ["能告诉我你们全科问诊收费价格吗？",'那首次购买感染及持续咳嗽诊断套餐的价格呢？',"能告诉我你们专科问诊收费价格吗？",
          '能告诉我你们急诊收费价格吗？','能告诉我你们心理科&精神科问诊收费价格吗？','能告诉我你们康复理疗收费价格吗？',
          '能告诉我你们整脊收费价格吗？']
# prompts = ['我现在需要和医生通电话！','我感到呼吸困难，胸痛……','今天我打篮球被篮球砸到了胸口，现在胸口很难受。没有其他症状。',
#           '好的，麻烦尽快帮我安排急诊。','我叫张梅，我的生日是1990/1/1,  电话是19937679835','好的']
# prompts = ['我咳嗽3天了，现在情况越来越严重，需要看医生。','没有了','我住在人民广场附近','太好了，那我约医院看诊。',
#           '下午3点吧','我的生日是1990/1/1,  电话是19937679835','好的']
# prompts = ['我今天早上鼻子出血了，很不舒服，需要看医生。','没有，就是早上刷牙时，鼻子开始流血。近半年都没有出现过。',
#           '我在东方体育中心附近。','太好了，那我约医生看诊。','下午3点吧','我的生日是1990/1/1,  电话是19937679835',
#           '好的']
#prompts = ['曜影医疗公司一共有哪些门诊部？','上海商城门诊部位于哪里？','天山门诊部位于哪里？','曜影医疗的服务模式是什么？']
#prompts = ['谢谢']
response = ''
count = 1
for question in prompts:
    messages.append({"role":"user", "content": question})
    response = model.chat(messages, add_generation_prompt=True, sampling_params=sampling_params)
    print(response)
    response = response[0].outputs[0].text
    messages.append({"role": "system", "content":response})
    print("##第",count,"轮次##")
    for message in messages:
        print(message)
    count +=1