Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| from transformers import AutoTokenizer | |
| from transformers import pipeline | |
| from utils import format_moves | |
| import pandas as pd | |
| import tensorflow as tf | |
| model_checkpoint = "distilgpt2" | |
| tokenizer = AutoTokenizer.from_pretrained(model_checkpoint) | |
| generate = pipeline("text-generation", | |
| model="arjunpatel/distilgpt2-finetuned-pokemon-moves", | |
| tokenizer=tokenizer) | |
| # load in the model | |
| seed_text = "This move is called " | |
| tf.random.set_seed(0) | |
| def update_history(df, move_name, move_desc, generation, parameters): | |
| new_row = [{"Move Name": move_name, | |
| "Move Description": move_desc, | |
| "Generation Type": generation, | |
| "Parameters": parameters}] | |
| return pd.concat([df, pd.DataFrame(new_row)]) | |
| def create_move(move, history): | |
| generated_move = format_moves(generate(seed_text + move, num_return_sequences=1)) | |
| return generated_move, update_history(history, move, generated_move, | |
| "baseline", "None") | |
| def create_greedy_search_move(move, history): | |
| generated_move = format_moves(generate(seed_text + move, do_sample=False)) | |
| return generated_move, update_history(history, move, generated_move, | |
| "greedy", "None") | |
| def create_beam_search_move(move, num_beams, history): | |
| generated_move = format_moves(generate(seed_text + move, num_beams=num_beams, | |
| num_return_sequences=1, | |
| do_sample=False, early_stopping=True)) | |
| return generated_move, update_history(history, move, generated_move, | |
| "beam", {"num_beams": 2}) | |
| def create_sampling_search_move(move, do_sample, temperature, history): | |
| generated_move = format_moves(generate(seed_text + move, do_sample=do_sample, temperature=float(temperature), | |
| num_return_sequences=1, topk=0)) | |
| return generated_move, update_history(history, move, generated_move, | |
| "temperature", {"do_sample": do_sample, | |
| "temperature": temperature}) | |
| def create_top_search_move(move, topk, topp, history): | |
| generated_move = format_moves(generate( | |
| seed_text + move, | |
| do_sample=True, | |
| num_return_sequences=1, | |
| top_k=topk, | |
| top_p=topp, | |
| force_word_ids=tokenizer.encode("The user", return_tensors='tf'))) | |
| return generated_move, update_history(history, move, generated_move, | |
| "top", {"top k": topk, | |
| "top p": topp}) | |
| demo = gr.Blocks() | |
| with demo: | |
| gr.Markdown("<h1><center>What's that Pokemon Move?</center></h1>") | |
| gr.Markdown( | |
| """This Gradio demo allows you to generate Pokemon Move descriptions given a name, and learn more about text | |
| decoding methods in the process! Each tab aims to explain each generation methodology available for the | |
| model. The dataframe below allows you to keep track of each move generated, to compare!""") | |
| gr.Markdown("<h3> How does text generation work? <h3>") | |
| gr.Markdown("""Roughly, text generation models accept an input sequence of words (or parts of words, | |
| known as tokens. | |
| These models then output a corresponding set of words or tokens. Given the input, the model | |
| estimates the probability of another possible word or token appearing right after the given sequence. In | |
| other words, the model estimates conditional probabilities and ranks them in order to generate sequences | |
| . """) | |
| gr.Markdown("Enter a two to three word Pokemon Move name of your imagination below, with each word capitalized!") | |
| gr.Markdown("<h3> Move Generation <h3>") | |
| with gr.Tabs(): | |
| with gr.TabItem("Standard Generation"): | |
| gr.Markdown( | |
| """The default parameters for distilgpt2 work well to generate moves. Use this tab to have fun and as | |
| a baseline for your experiments.""") | |
| with gr.Row(): | |
| text_input_baseline = gr.Textbox(label="Move", | |
| placeholder="Type a two or three word move name here! Try \"Wonder " | |
| "Shield\"!") | |
| text_output_baseline = gr.Textbox(label="Move Description", | |
| placeholder="Leave this blank!") | |
| text_button_baseline = gr.Button("Create my move!") | |
| with gr.TabItem("Greedy Search Decoding"): | |
| gr.Markdown(""" | |
| Greedy search is a decoding method that relies on finding words that has the highest estimated | |
| probability of following the sequence thus far. | |
| Therefore, the model \"greedily\" grabs the highest | |
| probability word and continues generating the sentence. | |
| This has the side effect of finding sequences that are reasonable, but avoids sequences that are | |
| less probable but way more interesting. | |
| Try the other decoding methods to get sentences with more variety! | |
| """) | |
| with gr.Row(): | |
| text_input_greedy = gr.Textbox(label="Move") | |
| text_output_greedy = gr.Textbox(label="Move Description") | |
| text_button_greedy = gr.Button("Create my move!") | |
| with gr.TabItem("Beam Search"): | |
| gr.Markdown("""Beam search is an improvement on Greedy Search. Instead of directly grabbing the word that | |
| maximizes probability, we conduct a search with B number of candidates. We then try to find the next word | |
| that would most likely follow each beam, and we grab the top B candidates of that search. This may | |
| eliminate one of the original beams we started with, and that's okay! That is how the algorithm decides | |
| on an optimal candidate. Eventually, the beam sequence terminate or are eliminated due to being too | |
| improbable. | |
| Increasing the number of beams will increase model generation time, but also result in a more thorough | |
| search. Decreasing the number of beams will decrease decoding time, but it may not find an optimal | |
| sentence. | |
| Play around with the num_beams parameter to experiment! """ | |
| ) | |
| with gr.Row(): | |
| num_beams = gr.Slider(minimum=2, maximum=10, value=2, step=1, | |
| label="Number of Beams") | |
| text_input_beam = gr.Textbox(label="Move") | |
| text_output_beam = gr.Textbox(label="Move Description") | |
| text_button_beam = gr.Button("Create my move!") | |
| with gr.TabItem("Sampling and Temperature Search"): | |
| gr.Markdown( | |
| """Greedy Search and Beam Search were both good at finding sequences that are likely to follow our | |
| input text, but when generating cool move descriptions, we want some more variety! | |
| Instead of choosing the word or token that is most likely to follow a given sequence, we can instead | |
| ask the model to sample across the probability distribution of likely words. | |
| It's kind of like walking into the tall grass and finding a Pokemon encounter. | |
| There are different encounter rates, which allow | |
| for the most common mons to appear (looking at you, Zubat), but also account for surprise, like shinys! | |
| We might even want to go further, though. We can rescale the probability distributions directly | |
| instead, allowing for rare words to temporarily become more frequently. We do this using the | |
| temperature parameter. | |
| Turn the temperature up, and rare tokens become very likely! Cool down, and we approach more sensible | |
| output. | |
| Experiment with turning sampling on and off, and by varying temperature below!. | |
| """) | |
| with gr.Row(): | |
| temperature = gr.Slider(minimum=0.3, maximum=4.0, value=1.0, step=0.1, | |
| label="Temperature") | |
| text_input_temp = gr.Textbox(label="Move") | |
| with gr.Row(): | |
| sample_boolean = gr.Checkbox(label="Enable Sampling?") | |
| text_output_temp = gr.Textbox(label="Move Description") | |
| text_button_temp = gr.Button("Create my move!") | |
| with gr.TabItem("Top K and Top P Sampling"): | |
| gr.Markdown( | |
| """When we want more control over the words we get to sample from, we turn to Top K and Top P | |
| decoding methods! | |
| The Top K sampling method selects the K most probable words given a sequence, and then samples from | |
| that subset, rather than the whole vocabulary. This effectively cuts out low probability words. | |
| Top P also reduces the available vocabulary to sample from, but instead of choosing the number of | |
| words or tokens in advance, we sort the vocabulary from most to least likely word, and we | |
| grab the smallest set of words that sum to P. This allows for the number of words we look at to | |
| change while sampling, instead of being fixed. | |
| We can even use both methods at the same time! To disable Top K, set it to 0 using the slider. | |
| To disable Top P, set it to 1""") | |
| with gr.Row(): | |
| topk = gr.Slider(minimum=0, maximum=200, value=0, step=5, | |
| label="Top K") | |
| text_input_top = gr.Textbox(label="Move") | |
| with gr.Row(): | |
| topp = gr.Slider(minimum=0.10, maximum=1, value=1, step=0.05, | |
| label="Top P") | |
| text_output_top = gr.Textbox(label="Move Description") | |
| text_button_top = gr.Button("Create my move!") | |
| with gr.Box(): | |
| gr.Markdown("<h3> Generation History <h3>") | |
| # Displays a dataframe with the history of moves generated, with parameters | |
| history = gr.Dataframe(headers=["Move Name", "Move Description", "Generation Type", "Parameters"]) | |
| with gr.Box(): | |
| gr.Markdown("<h3>How did you make this?<h3>") | |
| gr.Markdown(""" | |
| I collected the dataset from [Serebii] (https://www.serebii.net) , a news source and aggregator of Pokemon info. | |
| I then added a seed phrase "This move is called" just before each move in order to assist the model in | |
| generation. | |
| I then followed HuggingFace's handy language_modeling.ipynb for fine-tuning distillgpt2 on this tiny dataset, | |
| and it surprisingly worked! | |
| I learned all about text generation using the book [Natural Language Processing with Transformers] ( | |
| https://www.oreilly.com/library/view/natural-language-processing/9781098103231/) by Lewis Tunstall, | |
| Leandro von Werra and Thomas Wolf, as well as [this fantastic article] ( | |
| https://huggingface.co/blog/how-to-generate) by Patrick von Platen. Thanks to all of these folks for creating | |
| these learning materials, and thanks to the Hugging Face team for developing this product! """) | |
| text_button_baseline.click(create_move, inputs=[text_input_baseline, history], | |
| outputs=[text_output_baseline, history]) | |
| text_button_greedy.click(create_greedy_search_move, inputs=[text_input_greedy, history], | |
| outputs=[text_output_greedy, history]) | |
| text_button_temp.click(create_sampling_search_move, inputs=[text_input_temp, sample_boolean, temperature, history], | |
| outputs=[text_output_temp, history]) | |
| text_button_beam.click(create_beam_search_move, inputs=[text_input_beam, num_beams, history], | |
| outputs=[text_output_beam, history]) | |
| text_button_top.click(create_top_search_move, inputs=[text_input_top, topk, topp, history], | |
| outputs=[text_output_top, history]) | |
| demo.launch(share=True) | |