Update prompt

README.md

@@ -32,11 +32,18 @@ DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 # Load different ckpts via passing e.g. `revision=kto`
 model = OlmoeForCausalLM.from_pretrained("allenai/OLMoE-1B-7B-0924-Instruct").to(DEVICE)
 tokenizer = AutoTokenizer.from_pretrained("allenai/OLMoE-1B-7B-0924-Instruct")
-message = [{"role": "user", "content": "Explain to me like I'm five what is Bitcoin."}]
-inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt")
-out = model.generate(**inputs, max_length=64)
+messages = [{"role": "user", "content": "Explain to me like I'm five what is Bitcoin."}]
+inputs = tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors="pt").to(DEVICE)
+out = model.generate(inputs, max_length=100)
 print(tokenizer.decode(out[0]))
-# > # Bitcoin is a digital currency that is created and held electronically. No one controls it. Bitcoins aren’t printed, like dollars or euros – they’re produced by people and businesses running computers all around the world, using software that solves mathematical
+"""
+<|endoftext|><|user|>
+Explain to me like I'm five what is Bitcoin.
+<|assistant|>
+Bitcoin is like a special kind of money that you can use to buy things online. But unlike regular money, like dollars or euros, Bitcoin isn't printed by governments or banks. Instead, it's created by a special computer program that helps people keep track of it.
+
+Here's how it works: imagine you have a bunch of toys, and you want to
+"""
 ```
 
 Branches: