Update README.md

README.md

@@ -29,14 +29,13 @@ llm = pipeline(
     "text-generation",
     model=model_id,
     torch_dtype=torch.float16,
-    max_new_tokens=512,
     device_map="auto",
 )
 
 def generate(input_text):
   messages = [{"role": "user", "content": input_text}]
   prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-  outputs = llm(prompt)
+  outputs = llm(prompt, max_new_tokens=512,)
   return outputs[0]["generated_text"][len(prompt):]
 
 generate("Explain quantum tunneling in simple terms.")