Update README.md

README.md

@@ -54,9 +54,11 @@ model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b-instruct")
 messages = [
     {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
 ]
-input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True).input_ids
 
-outputs = model.generate(input_ids)
+input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+input_ids = tokenizer(input_text, return_tensors="pt")
+
+outputs = model.generate(input_ids, max_new_tokens=30)
 print(tokenizer.decode(outputs[0]))
 ```
 
@@ -78,9 +80,11 @@ model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b-instruct",
 messages = [
     {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
 ]
-input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True).input_ids.to("cuda")
 
-outputs = model.generate(input_ids)
+input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
+
+outputs = model.generate(input_ids, max_new_tokens=30)
 print(tokenizer.decode(outputs[0]))
 ```
 
@@ -104,9 +108,10 @@ model = torch.compile(model)
 messages = [
     {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
 ]
-input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True).input_ids.to("cuda")
+input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
 
-outputs = model.generate(input_ids)
+outputs = model.generate(input_ids, max_new_tokens=30)
 print(tokenizer.decode(outputs[0]))
 ```
 
@@ -132,9 +137,10 @@ model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b-instruct",
 messages = [
     {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
 ]
-input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True).input_ids.to("cuda")
+input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
 
-outputs = model.generate(input_ids)
+outputs = model.generate(input_ids, max_new_tokens=30)
 print(tokenizer.decode(outputs[0]))
 ```
 
@@ -156,9 +162,10 @@ model = AutoModelForCausalLM.from_pretrained("tiiuae/falcon-mamba-7b-instruct",
 messages = [
     {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
 ]
-input_ids = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True).input_ids.to("cuda")
+input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
 
-outputs = model.generate(input_ids)
+outputs = model.generate(input_ids, max_new_tokens=30)
 print(tokenizer.decode(outputs[0]))
 ```