from transformers import AutoModelForImageTextToText, AutoTokenizer
model_id = "andrevp/Ornith-1.0-9B-Heretic-Uncensored"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForImageTextToText.from_pretrained(
model_id, dtype="bfloat16", device_map="auto"
)
messages = [{"role": "user", "content": "Write a Python function is_prime(n). Keep it short."}]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
generated = model.generate(
**inputs, max_new_tokens=512, do_sample=True,
temperature=0.6, top_p=0.95, top_k=20,
)
output = generated[0][inputs.input_ids.shape[1]:]
content = tokenizer.decode(output, skip_special_tokens=True)