from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("khazarai/Qwen3.5-2B-Qwen3.6-plus-Distilled")
model = AutoModelForCausalLM.from_pretrained(
    "khazarai/Qwen3.5-2B-Qwen3.6-plus-Distilled",
    device_map={"": 0}
)

question = """
An 8-year-old boy is brought to the pediatrician by his mother with nausea, vomiting, and decreased frequency of urination. He has acute lymphoblastic leukemia for which he received the 1st dose of chemotherapy 5 days ago. His leukocyte count was 60,000/mm3 before starting chemotherapy. The vital signs include: pulse 110/min, temperature 37.0°C (98.6°F), and blood pressure 100/70 mm Hg. The physical examination shows bilateral pedal edema. Which of the following serum studies and urinalysis findings will be helpful in confirming the diagnosis of this condition? ?
{'A': 'Hyperkalemia, hyperphosphatemia, hypocalcemia, and extremely elevated creatine kinase (MM)', 'B': 'Hyperkalemia, hyperphosphatemia, hypocalcemia, hyperuricemia, urine supernatant pink, and positive for heme', 'C': 'Hyperuricemia, hyperkalemia, hyperphosphatemia, lactic acidosis, and urate crystals in the urine', 'D': 'Hyperuricemia, hyperkalemia, hyperphosphatemia, and urinary monoclonal spike', 'E': 'Hyperuricemia, hyperkalemia, hyperphosphatemia, lactic acidosis, and oxalate crystals'}
"""

messages = [
    {
        "role": "user",
        "content": [
            {"type": "text", "text": question}
        ]
    },
]

inputs = tokenizer.apply_chat_template(
    messages,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt",
    enable_thinking = True,
).to(model.device)

outputs = model.generate(**inputs, max_new_tokens=2048, do_sample = True, top_p=0.95, top_k=20, temperature=1.0, min_p=0.0, repetition_penalty=1.0)
print(tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:]))