from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
repo = "Davitotty1/Teleste-Learner-4B"
tok = AutoTokenizer.from_pretrained(repo)
model = AutoModelForCausalLM.from_pretrained(
repo, dtype=torch.float16, device_map="auto"
)
system = (
"You adapt to the current request. Infer the user's goal, the hidden rules, "
"and the output contract from this conversation only. If examples are present, "
"the mapping in those examples is the law. If a later message changes the rules, "
"the new rules replace the old ones. Check the answer against the inferred "
"contract before you finish. Do not keep a default job."
)
messages = [
{"role": "system", "content": system},
{"role": "user", "content": "examples: walrus→12, turtle→12, pig→6. now sloth → ?"},
]
prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tok(prompt, return_tensors="pt").to(model.device)
out = model.generate(**inputs, max_new_tokens=256, temperature=0.6, top_p=0.95, top_k=20)
print(tok.decode(out[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))