Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2
3tokenizer = AutoTokenizer.from_pretrained("khazarai/ClinicalReasoning-0.6B")
4model = AutoModelForCausalLM.from_pretrained(
5 "khazarai/ClinicalReasoning-0.6B",
6 device_map={"": 0}
7)
8
9system = "Please answer with one of the option in the bracket. Write reasoning in between <analysis></analysis>. Write answer in between <answer></answer>."
10
11question = """
12A 44-year-old female is admitted to the neurological service. You examine her chart and note that after admission she was started on nimodipine. Which of the following pathologies would benefit from this pharmacologic therapy??
13{'A': 'Pseudotumor cerebri', 'B': 'Thromboembolic stroke', 'C': 'Epidural hematoma', 'D': 'Subdural hematoma', 'E': 'Subarachnoid hemorrhage'}
14"""
15
16messages = [
17 {"role" : "system", "content" : system},
18 {"role" : "user", "content" : question}
19]
20text = tokenizer.apply_chat_template(
21 messages,
22 tokenize = False,
23 add_generation_prompt = True,
24 enable_thinking = False,
25)
26
27from transformers import TextStreamer
28_ = model.generate(
29 **tokenizer(text, return_tensors = "pt").to("cuda"),
30 max_new_tokens = 512,
31 temperature = 0.6,
32 top_p = 0.95,
33 top_k = 20,
34 streamer = TextStreamer(tokenizer, skip_prompt = True),
35)