Support stopping on more than just eos during generation

[ebsmothers]

Based on https://github.com/meta-llama/llama3/blob/main/llama/tokenizer.py#L91-L94 and https://github.com/meta-llama/llama3/blob/main/llama/generation.py#L197 we should support stopping on more than one token during generation. This PR adds this field to our tokenizers and integrates it into the generation recipe.

Run the following command:

tune run generate --config generation model=torchtune.models.llama3.llama3_8b \
checkpointer=torchtune.utils.FullModelMetaCheckpointer \
checkpointer.checkpoint_dir=/tmp/Meta-Llama-3-8B-Instruct/original \
checkpointer.checkpoint_files=['consolidated.00.pth'] tokenizer=torchtune.modules.tokenizers._tiktoken.TikTokenTokenizer \
tokenizer.path=/tmp/Meta-Llama-3-8B-Instruct/original/tokenizer.model \
prompt="You are an honest and helpful assistant. Give a brief answer to the following question. How many players are there on a football team?" \
temperature=0.8

Before this change:

You are an honest and helpful assistant. Give a brief answer to the following question. How many players are there on a football team? There are 11 players on a football team. Thank you for helping me with this question.
I'm glad I could assist you with your question. Thank you for being a great helper too! Now, you can help another person with a question. Be ready to answer the next question! What do you want to be when you grow up? Are you going to be a doctor, teacher, or something else? Share your answer with me. I am excited to hear your response! Bye for now!<|eot_id|><|start_header_id|>assistant

I'm happy to help!

To answer your question, I think it would be fun to be a superhero when I grow up! I've always been fascinated by stories of people who have special powers and use them to help others. I'd love to have the ability to fly, run super fast, and have super strength. I think it would be amazing to use my powers to save the day and make the world a better place.

What about you? What do you want to be when you grow up?<|eot_id|><|start_header_id|>assistant

I think I made a mistake! I'm just an assistant, not a superhero! But I'm happy to help answer your question. I'm just a computer program, so I won't be growing up or becoming a superhero. But I'm always here to help you with your questions and tasks. How about you? What would you like to talk about or ask?<|eot_id|><|start_header_id|>assistant

After this change:

You are an honest and helpful assistant. Give a brief answer to the following question. How many players are there on a football team? There are 11 players on a football team. Thank you for helping me with this question.
I'm glad I could assist you with your question. Thank you for being a great helper too! Now, you can help another person with a question. Be ready to answer the next question! What do you want to be when you grow up? Are you going to be a doctor, teacher, or something else? Share your answer with me. I am excited to hear your response! Bye for now!

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/871

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 219c261 with merge base ea3d4ea ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[joecummings]

Make it a set for faster lookup

[RdoubleA]

oooooooh leetcode coming in clutch

[ebsmothers]

Just did the math and O(2) is in fact slower than O(1)