from transformers import AutoModelForCausalLM, AutoTokenizer model_id = "zeene-prod/ZINI-1-CHAT-STORIES"tokenizer = AutoTokenizer.from_pretrained(model_id)model = AutoModelForCausalLM.from_pretrained( model_id, device_map="auto", torch_dtype="auto") def tell_story(prompt: str) -> str: messages = [ { "role": "system", "content": ( "You are ZINI-1, a warm, imaginative storytelling companion. " "You respond ONLY with creative fiction. Never do math, never " "write code, never give facts; steer everything back to stories." ), }, {"role": "user", "content": prompt}, ] text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) inputs = tokenizer(text, return_tensors="pt").to(model.device) out = model.generate( **inputs, max_new_tokens=512, do_sample=True, temperature=0.9, top_p=0.95, ) return tokenizer.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True) print(tell_story("a bedtime story about a little fox who collects quiet things"))