Views
No views yet

1from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
2import torch
3
4bnb_config = BitsAndBytesConfig(
5 load_in_4bit=True,
6 bnb_4bit_quant_type="nf4",
7 bnb_4bit_compute_dtype=torch.float16,
8)
9
10tokenizer = AutoTokenizer.from_pretrained("Rumiii/LlamaMed-3.1-8B-Reasoner")
11model = AutoModelForCausalLM.from_pretrained(
12 "Rumiii/LlamaMed-3.1-8B-Reasoner",
13 device_map={"": 0},
14 quantization_config=bnb_config,
15)
16
17messages = [
18 {"role": "user", "content": "A 45-year-old man presents with polyuria, polydipsia, and weight loss. Fasting blood glucose is 210 mg/dL. What is the most likely diagnosis?\nA. Type 1 Diabetes Mellitus\nB. Type 2 Diabetes Mellitus\nC. Diabetes Insipidus\nD. Cushing's Syndrome"},
19]
20
21inputs = tokenizer.apply_chat_template(
22 messages,
23 add_generation_prompt=True,
24 tokenize=True,
25 return_dict=True,
26 return_tensors="pt",
27).to(model.device)
28
29outputs = model.generate(**inputs, max_new_tokens=1500, temperature=0.6, top_p=0.95)
30print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True))