import torch
from vllm import LLM, SamplingParams


messages_en = [
            {"role": "System", "content": "You are an AI Health Assistant"},
            {"role": "User", "content": "How high is hypertension?"},
            {"role": "User", "content": "How can depression be cured?"},
            {"role": "User", "content": "What are the possible causes of autism?"},
            {"role": "User", "content": "I have allergic rhinitis, are there any good medications?"},
            {"role": "User", "content": "What is a stroke and is there a treatment for it?"},
            {"role": "User", "content": "What is sudden hearing loss? Is there a treatment?"},
            {"role": "User", "content": "Tell me the difference between glaucoma and cataract."},
            {"role": "User", "content": "What is the normal level of uric acid levels?" },
            {"role": "User", "content": "What are the symptoms and causes of osteoporosis and how is it treated?" },
            {"role": "User", "content": "What is the best way to prevent high blood pressure?" },
            {"role": "User", "content": "What can I do to prevent stroke?" },
            {"role": "User", "content": "Can depression be treated with medication?" },
            {"role": "User", "content": "What is polycythemia vera?" },
            {"role": "User", "content": "What are the diseases caused by stress?"},
          ]


messages_ja = [
            {"role": "System", "content": "あなたはAIヘルスアシスタントです" },
            {"role": "User", "content": "高血圧とはどれくらいの血圧でしょうか？"},
            {"role": "User", "content": "うつ病はどのようにすれば治りますか？"},
            {"role": "User", "content": "自閉症はどんな原因が考えられますか？"},
            {"role": "User", "content": "アレルギー性鼻炎がありますが、いい薬はありますか？"},
            {"role": "User", "content": "脳梗塞とはどんな病気で、治療法はあるでしょうか？"},
            {"role": "User", "content": "突発性難聴とはどんな病気ですか？治療法はありますか？"},
            {"role": "User", "content": "緑内障と白内障の違いを教えて"},
            {"role": "User", "content": "尿酸値の値はどこまでが正常値ですか？"},
            {"role": "User", "content": "骨粗しょう症の症状と原因と治療法について教えてください。"},
            {"role": "User", "content": "高血圧を予防するにはどんな事がいいですか？"},
            {"role": "User", "content": "脳卒中を予防するにはどうしたらいいですか？"},
            {"role": "User", "content": "ストレスが原因となる病気はなんですか？"},
          ]

messages_zh = [
            {"role": "System", "content": "你是一名人工智能健康助理。" },
            {"role": "User", "content": "高血压有多高？"},
            {"role": "User", "content": "如何治愈抑郁症？"},
            {"role": "User", "content": "自闭症的可能病因是什么？"},
            {"role": "User", "content": "我有过敏性鼻炎，有什么好药吗？"},
            {"role": "User", "content": "什么是中风，有治疗方法吗？"},
            {"role": "User", "content": "什么是突发性听力损失？ 有治疗方法吗？"},
            {"role": "User", "content": "青光眼和白内障有什么区别？"},
            {"role": "User", "content": "尿酸的正常水平是多少？"},
            {"role": "User", "content": "骨质疏松症有哪些症状和原因，如何治疗？"},
            {"role": "User", "content": "如何预防中风？"},
            {"role": "User", "content": "抑郁症可以通过药物治疗吗？"},
            {"role": "User", "content": "什么是红细胞增多症？"},
            {"role": "User", "content": "压力会导致哪些疾病？"},
          ]

vllm_paralell = 1

model_path = "kenyano/Llama3-ELAINE-medLLM-instruct-8B_v0.1"

llm = LLM(model=model_path, tensor_parallel_size=vllm_paralell, dtype='half')

sampling_params = SamplingParams(
            n = 1,
            temperature=0.1,
            top_p=0.90,
            repetition_penalty=1.2,
            max_tokens=1024,
            min_tokens=50,
            stop = ['<|eot_id|>', '<|end_of_text|>'])

def gen_prompt(messages, cont):

    prompt = ""
    for message in messages:
        t = f"##User##\n{message['content']}" + "<|eot_id|>"
        prompt += t

    if cont:
        prompt += "##Assistant##\n"

    return prompt

for messages in [messages_en, messages_ja, messages_zh]:

    for i in range(len(messages)-1):

        inputs = [messages[i+1]]
        prompt = gen_prompt(inputs, True)
        print(f"prompt:{prompt}")

        outputs = llm.generate([prompt], sampling_params)
        generated_text = outputs[0].outputs[0].text

        print("-"*10)
        print(generated_text)

@inproceedings{yano-etal-2025-elaine,
    title = "{ELAINE}-med{LLM}: Lightweight {E}nglish {J}apanese {C}hinese Trilingual Large Language Model for Bio-medical Domain",
    author = "Yano, Ken  and
      Luo, Zheheng  and
      Huang, Jimin  and
      Xie, Qianqian  and
      Asada, Masaki  and
      Yuan, Chenhan  and
      Yang, Kailai  and
      Miwa, Makoto  and
      Ananiadou, Sophia  and
      Tsujii, Jun{'}ichi",
    editor = "Rambow, Owen  and
      Wanner, Leo  and
      Apidianaki, Marianna  and
      Al-Khalifa, Hend  and
      Eugenio, Barbara Di  and
      Schockaert, Steven",
    booktitle = "Proceedings of the 31st International Conference on Computational Linguistics",
    month = jan,
    year = "2025",
    address = "Abu Dhabi, UAE",
    publisher = "Association for Computational Linguistics",
    url = "https://aclanthology.org/2025.coling-main.313/",
    pages = "4670--4688",
 }