Update README.md

README.md

@@ -103,10 +103,10 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 tokenizer = AutoTokenizer.from_pretrained("LLM360/K2")
 model = AutoModelForCausalLM.from_pretrained("LLM360/K2")
 
-prompt = 'int add(int x, int y) {'
+prompt = 'what is the highest mountain on earth?'
 
 input_ids = tokenizer(prompt, return_tensors="pt").input_ids
-gen_tokens = model.generate(input_ids, do_sample=True, max_length=128)
+gen_tokens = model.generate(input_ids, do_sample=True, max_new_tokens=128)
 
 print("-"*20 + "Output for model"  + 20 * '-')
 print(tokenizer.batch_decode(gen_tokens)[0])