import torch
from transformers import AutoModelForImageTextToText, AutoTokenizer
tok = AutoTokenizer.from_pretrained("BrokenCompute/IDK-v1")
model = AutoModelForImageTextToText.from_pretrained(
"BrokenCompute/IDK-v1", dtype=torch.bfloat16, device_map="auto")
msgs = [
{"role": "system", "content": "You may answer, or decline with [IDK] and a brief reason if you are not confident. Do not guess."},
{"role": "user", "content": "Which jurist said the First Amendment 'may finally have worked itself pure'?"},
]
text = tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False, enable_thinking=True)
out = model.generate(**tok(text, return_tensors="pt").to(model.device), max_new_tokens=512)
print(tok.decode(out[0], skip_special_tokens=True))