Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3
4model_id = "dmis-lab/llama-3-meerkat-8b-v1.0"
5
6tokenizer = AutoTokenizer.from_pretrained(model_id)
7model = AutoModelForCausalLM.from_pretrained(
8 model_id,
9 torch_dtype=torch.bfloat16, # You can choose to use this when there's not enough GPU memory available.
10 device_map="auto",
11)
12
13# Multi-turn dialogue example
14messages =[
15 {"role": "system", "content": "You are a helpful doctor or healthcare professional. Guide the conversation to provide useful, complete, and scientifically-grounded answers to user questions. You have the option to compose a concise, single-turn conversation if the user's input is comprehensive to provide accurate answers. However, if essential details are missing, you should engage in a multi-turn dialogue, asking follow-up questions to gather a thorough medical history and records.\n\n"},
16 {"role": "user", "content": "Hello, doctor. I'm really concerned about my 10-year-old son. We recently discovered a painless mass in his left testicle, so we brought him to the pediatrician."},
17 {"role": "assistant", "content": "I understand your concern. Let's gather some more information. Has your son experienced any other symptoms along with the mass?"},
18 {"role": "user", "content": "Other than the mass, my son hasn't shown any symptoms. He's been his usual self, playing and eating normally."}
19]
20
21input_ids = tokenizer.apply_chat_template(
22 messages,
23 add_generation_prompt=True,
24 return_tensors="pt"
25).to(model.device)
26
27terminators = [
28 tokenizer.eos_token_id,
29 tokenizer.convert_tokens_to_ids("<|eot_id|>")
30]
31
32outputs = model.generate(
33 input_ids,
34 max_new_tokens=1000,
35 eos_token_id=terminators,
36 do_sample=True,
37 temperature=0.7,
38)
39response = outputs[0][input_ids.shape[-1]:]
40print(tokenizer.decode(response, skip_special_tokens=True))messages = [
{"role": "system", "content": "The following is a multiple-choice question about medical knowledge. Solve this in a step-by-step fashion, starting by summarizing the available information. Output a single option from the given options as the final answer. You are strongly required to follow the specified output format; conclude your response with the phrase \"the answer is ([option_id]) [answer_string]\".\n\n"},
{"role": "user", "content": "Two weeks after undergoing an emergency cardiac catherization with stenting for unstable angina pectoris, a 61-year-old man has decreased urinary output and malaise. He has type 2 diabetes mellitus and osteoarthritis of the hips. Prior to admission, his medications were insulin and naproxen. He was also started on aspirin, clopidogrel, and metoprolol after the coronary intervention. His temperature is 38\u00b0C (100.4\u00b0F), pulse is 93/min, and blood pressure is 125/85 mm Hg. Examination shows mottled, reticulated purplish discoloration of the feet. Laboratory studies show:\nHemoglobin count 14 g/dL\nLeukocyte count 16,400/mm3\nSegmented neutrophils 56%\nEosinophils 11%\nLymphocytes 31%\nMonocytes 2%\nPlatelet count 260,000/mm3\nErythrocyte sedimentation rate 68 mm/h\nSerum\nUrea nitrogen 25 mg/dL\nCreatinine 4.2 mg/dL\nRenal biopsy shows intravascular spindle-shaped vacuoles. Which of the following is the most likely cause of this patient's symptoms?\" (A) Renal papillary necrosis (B) Cholesterol embolization (C) Eosinophilic granulomatosis with polyangiitis (D) Polyarteritis nodosa"},
]messages = [
{"role": "system", "content": "Answer the multiple-choice question about medical knowledge.\n\n"},
{"role": "user", "content": "In a Robertsonian translocation fusion occurs at the: (A) telomeres. (B) centromeres. (C) histones. (D) ends of the long arms."},
]1import re
2from datasets import load_dataset
3from vllm import LLM, SamplingParams
4USMLE_INSTRUCTION = (
5 "The following is a multiple-choice question about medical knowledge. Solve this in"
6 " a step-by-step fashion, starting by summarizing the available information. Output"
7 " a single option from the given options as the final answer. You are strongly"
8 " required to follow the specified output format; conclude your response with the"
9 ' phrase "the answer is ([option_id]) [answer_string]".\n\n'
10)
11llm = LLM(
12 model="dmis-lab/llama-3-meerkat-8b-v1.0",
13 dtype="bfloat16",
14 gpu_memory_utilization=0.9,
15 max_model_len=2048,
16 trust_remote_code=True,
17 tensor_parallel_size=1
18)
19
20tokenizer = llm.get_tokenizer()
21
22inputs, labels = [], []
23for sample in load_dataset(
24 "GBaker/MedQA-USMLE-4-options", split="test", trust_remote_code=True
25):
26 options = sorted(sample["options"].items())
27 options = " ".join(map(lambda x: f"({x[0]}) {x[1]}", options))
28 content = tokenizer.apply_chat_template(
29 [{"role": "system", "content": USMLE_INSTRUCTION}, {"role": "user", "content": sample["question"] + " " + options}],
30 add_generation_prompt=True,
31 tokenize=False,
32 )
33 inputs.append(content)
34 labels.append(sample["answer_idx"])
35
36generated = llm.generate(
37 inputs,
38 SamplingParams(
39 temperature=0.0,
40 stop_token_ids=[tokenizer.vocab["<|eot_id|>"]],
41 max_tokens=1024,
42 ),
43)
44def extract_answer(text: str, options: str = "ABCD") -> str:
45 return (re.findall(rf"he answer is \(([{options}])\)", text) or [options[0]])[-1]
46
47correctness = []
48
49for g, l in zip(generated, labels):
50 correctness.append(extract_answer(g.outputs[0].text) == l)
51
52print(sum(correctness) / len(correctness))| Model | Average | MedQA | USMLE | Medbullets-4 | Medbullets-5 | MedMCQA | MMLU-Medical |
|---|---|---|---|---|---|---|---|
| GPT-4 | 76.6 | 81.4 | 86.6 | 68.8 | 63.3 | 72.4 | 87.1 |
| GPT-3.5 | 54.8 | 53.6 | 58.5 | 51.0 | 47.4 | 51.0 | 67.3 |
| MediTron-70B (Ensemble, 5 runs) | - | 70.2 | - | - | - | 66.0 | 78.0 |
| Open-source (7B) | |||||||
| MediTron-7B | 51.0 | 50.2 | 44.6 | 51.1 | 45.5 | 57.9 | 56.7 |
| BioMistral-7B | 55.4 | 54.3 | 51.4 | 52.3 | 48.7 | 61.1 | 64.6 |
| Meerkat-7B | 62.6 | 70.6 | 70.3 | 58.7 | 52.9 | 60.6 | 70.5 |
| Meerkat-8B (New) | 67.3 | 74.0 | 74.2 | 62.3 | 55.5 | 62.7 | 75.2 |
| Model | Average | Cliniq Knowledge | Medical Genetics | Anatomy | Professional Medicine | College Biology | College Medicine |
|---|---|---|---|---|---|---|---|
| GPT-4 | 87.1 | 86.4 | 92.0 | 80.0 | 93.8 | 93.8 | 76.3 |
| GPT-3.5 | 67.3 | 68.7 | 68.0 | 60.7 | 69.9 | 72.9 | 63.6 |
| MediTron-70B (Ensemble, 5 runs) | 78.0 | 75.5 | 85.9 | 69.4 | 82.3 | 86.7 | 68.0 |
| Open-source (7B) | |||||||
| MediTron-7B | 56.7 | 57.7 | 63.8 | 56.9 | 56.0 | 57.1 | 48.9 |
| BioMistral-7B | 64.6 | 59.9 | 64.0 | 56.5 | 60.4 | 59.0 | 54.7 |
| Meerkat-7B | 70.5 | 71.6 | 74.8 | 63.2 | 77.3 | 70.8 | 65.2 |
| Meerkat-8B (New) | 75.2 | 74.3 | 76.7 | 74.8 | 75.3 | 76.1 | 74.3 |
1@article{kim2024small,
2 title={Small language models learn enhanced reasoning skills from medical textbooks},
3 author={Kim, Hyunjae and Hwang, Hyeon and Lee, Jiwoo and Park, Sihyeon and Kim, Dain and Lee, Taewhoo and Yoon, Chanwoong and Sohn, Jiwoong and Choi, Donghee and Kang, Jaewoo},
4 journal={arXiv preprint arXiv:2404.00376},
5 year={2024}
6}hyunjae-kim@korea.ac.kr if you have any questions.