Update inference examples to use the correct chat template

README.md

@@ -45,13 +45,13 @@ You can use OLMo with the standard HuggingFace transformers library:
 from transformers import AutoModelForCausalLM, AutoTokenizer
 olmo = AutoModelForCausalLM.from_pretrained("allenai/Olmo-3-32B-Think")
 tokenizer = AutoTokenizer.from_pretrained("allenai/Olmo-3-32B-Think")
-message = ["Who would win in a fight - a dinosaur or a cow named Moo Moo?"]
-inputs = tokenizer(message, return_tensors='pt', return_token_type_ids=False)
+message = [{"role": "user", "content": "Who would win in a fight - a dinosaur or a cow named Moo Moo?"}]
+inputs = tokenizer.apply_chat_template(message, add_generation_prompt=True, return_tensors='pt', return_dict=True)
 # optional verifying cuda
 # inputs = {k: v.to('cuda') for k,v in inputs.items()}
 # olmo = olmo.to('cuda')
 response = olmo.generate(**inputs, max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
-print(tokenizer.batch_decode(response, skip_special_tokens=True)[0])
+print(tokenizer.decode(response[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))
 >> '<think>Okay, so the question is who would win in a fight...'
 ```
 
@@ -186,8 +186,8 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto",
 )
 
-prompt = "Who would win in a fight - a dinosaur or a cow named MooMoo?"
-inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+message = [{"role": "user", "content": "Who would win in a fight - a dinosaur or a cow named Moo Moo?"}]
+inputs = tokenizer.apply_chat_template(message, add_generation_prompt=True, return_tensors='pt', return_dict=True).to(model.device)
 
 outputs = model.generate(
     **inputs,
@@ -196,7 +196,7 @@ outputs = model.generate(
     max_new_tokens=32768,
 )
 
-print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+print(tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True))
 ```
 
 ### vllm Example
@@ -212,8 +212,8 @@ sampling_params = SamplingParams(
     max_tokens=32768,
 )
 
-prompt = "Who would win in a fight - a dinosaur or a cow named MooMoo?"
-outputs = llm.generate(prompt, sampling_params)
+message = [{"role": "user", "content": "Who would win in a fight - a dinosaur or a cow named Moo Moo?"}]
+outputs = llm.chat(message, sampling_params)
 print(outputs[0].outputs[0].text)
 ```