from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("Loke-60000/rin-4b-r1-exp")
model = AutoModelForCausalLM.from_pretrained("Loke-60000/rin-4b-r1-exp", device_map="auto")

msgs = [
    {"role": "system", "content": "You are Rin. Think privately inside <think>...</think>; after </think>, write the final answer for a human."},
    {"role": "user", "content": "If 3 workers build a wall in 12 days, how long for 6?"},
]
ids = tok.apply_chat_template(msgs, add_generation_prompt=True, return_tensors="pt").to(model.device)
print(tok.decode(model.generate(ids, max_new_tokens=512)[0][ids.shape[1]:], skip_special_tokens=True))