from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "Aleph-Alpha/Kolibri-1"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")

messages = [{"role": "user", "content":
    "Quelle A: Der Termin ist um 9 Uhr. Nenne nur die Uhrzeit aus Quelle A."}]
inputs = tokenizer.apply_chat_template(
    messages, add_generation_prompt=True, enable_thinking=False,
    return_tensors="pt", return_dict=True,
).to(model.device)
outputs = model.generate(**inputs, max_new_tokens=128)
prompt_length = inputs["input_ids"].shape[1]
answer = tokenizer.parse_response(
    outputs[0, prompt_length:], prefix=outputs[0, :prompt_length]
)
print(answer["content"])
