Spaces:
Sleeping
Sleeping
| # Gradio Params Playground | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import torch | |
| import gradio as gr | |
| # Load default model as GPT2 | |
| tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2") | |
| model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2") | |
| # Define functions | |
| global chosen_strategy | |
| def generate(input_text, number_steps, number_beams, number_beam_groups, diversity_penalty, length_penalty, num_return_sequences, temperature, no_repeat_ngram_size, repetition_penalty, early_stopping, beam_temperature, top_p, top_k,penalty_alpha,top_p_box,top_k_box,strategy_selected,model_selected): | |
| chosen_strategy = strategy_selected | |
| inputs = tokenizer(input_text, return_tensors="pt") | |
| if chosen_strategy == "Sampling": | |
| top_p_flag = top_p_box | |
| top_k_flag = top_k_box | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=number_steps, | |
| return_dict_in_generate=False, | |
| temperature=temperature, | |
| top_p=top_p if top_p_flag else None, | |
| top_k=top_k if top_k_flag else None, | |
| no_repeat_ngram_size = no_repeat_ngram_size, | |
| repetition_penalty = repetition_penalty if (repetition_penalty > 0) else None, | |
| output_scores=False, | |
| do_sample=True | |
| ) | |
| return tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| elif chosen_strategy == "Beam Search": | |
| beam_temp_flag = beam_temperature | |
| early_stop_flag = early_stopping | |
| inputs = tokenizer(input_text, return_tensors="pt") | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=number_steps, | |
| num_beams=number_beams, | |
| num_return_sequences=min(num_return_sequences, number_beams), | |
| return_dict_in_generate=False, | |
| length_penalty=length_penalty, | |
| temperature=temperature if beam_temp_flag else None, | |
| no_repeat_ngram_size = no_repeat_ngram_size, | |
| repetition_penalty = repetition_penalty if (repetition_penalty > 0) else None, | |
| early_stopping = True if early_stop_flag else False, | |
| output_scores=False, | |
| do_sample=True if beam_temp_flag else False | |
| ) | |
| beam_options_list = [] | |
| for i, beam_output in enumerate(outputs): | |
| beam_options_list.append (tokenizer.decode(beam_output, skip_special_tokens=True)) | |
| options = "\n\n - Option - \n".join(beam_options_list) | |
| return ("Beam Search Generation" + "\n" + "-" * 10 + "\n" + options) | |
| #print ("Option {}: {}\n".format(i, tokenizer.decode(beam_output, skip_special_tokens=True))) | |
| elif chosen_strategy == "Diversity Beam Search": | |
| early_stop_flag = early_stopping | |
| if number_beam_groups == 1: | |
| number_beam_groups = 2 | |
| if number_beam_groups > number_beams: | |
| number_beams = number_beam_groups | |
| inputs = tokenizer(input_text, return_tensors="pt") | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=number_steps, | |
| num_beams=number_beams, | |
| num_beam_groups=number_beam_groups, | |
| diversity_penalty=float(diversity_penalty), | |
| num_return_sequences=min(num_return_sequences, number_beams), | |
| return_dict_in_generate=False, | |
| length_penalty=length_penalty, | |
| no_repeat_ngram_size = no_repeat_ngram_size, | |
| repetition_penalty = repetition_penalty if (repetition_penalty > 0) else None, | |
| early_stopping = True if early_stop_flag else False, | |
| output_scores=False, | |
| ) | |
| beam_options_list = [] | |
| for i, beam_output in enumerate(outputs): | |
| beam_options_list.append (tokenizer.decode(beam_output, skip_special_tokens=True)) | |
| options = "\n\n ------ Option ------- \n".join(beam_options_list) | |
| return ("Diversity Beam Search Generation" + "\n" + "-" * 10 + "\n" + options) | |
| elif chosen_strategy == "Contrastive": | |
| top_k_flag = top_k_box | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=number_steps, | |
| return_dict_in_generate=False, | |
| temperature=temperature, | |
| penalty_alpha=penalty_alpha, | |
| top_k=top_k if top_k_flag else None, | |
| no_repeat_ngram_size = no_repeat_ngram_size, | |
| repetition_penalty = repetition_penalty if (repetition_penalty > 0) else None, | |
| output_scores=False, | |
| do_sample=True | |
| ) | |
| return tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| #--------ON SELECTING MODEL------------------------ | |
| def load_model(model_selected): | |
| if model_selected == "gpt2": | |
| tokenizer = AutoTokenizer.from_pretrained("gpt2") | |
| model = AutoModelForCausalLM.from_pretrained("gpt2", pad_token_id=tokenizer.eos_token_id) | |
| #print (model_selected + " loaded") | |
| if model_selected == "Gemma 2": | |
| tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b") | |
| model = AutoModelForCausalLM.from_pretrained("google/gemma-2b") | |
| #--------ON SELECT NO. OF RETURN SEQUENCES---------- | |
| def change_num_return_sequences(n_beams, num_return_sequences): | |
| if (num_return_sequences > n_beams): | |
| return gr.Slider( | |
| label="Number of sequences", minimum=1, maximum=n_beams, step=1, value=n_beams) | |
| return gr.Slider ( | |
| label="Number of sequences", minimum=1, maximum=n_beams, step=1, value=num_return_sequences) | |
| #--------ON CHANGING NO OF BEAMS------------------ | |
| def popualate_beam_groups (n_beams): | |
| global chosen_strategy | |
| no_of_beams = n_beams | |
| No_beam_group_list = [] #list for beam group selection | |
| for y in range (2, no_of_beams+1): | |
| if no_of_beams % y == 0: #perfectly divisible | |
| No_beam_group_list.append (y) #add to list, use as list for beam group selection | |
| if chosen_strategy == "Diversity Beam Search": | |
| return {beam_groups: gr.Dropdown(No_beam_group_list, value=max(No_beam_group_list), label="Beam groups", info="Divide beams into equal groups", visible=True), | |
| num_return_sequences: gr.Slider(maximum=no_of_beams) | |
| } | |
| if chosen_strategy == "Beam Search": | |
| return {beam_groups: gr.Dropdown(No_beam_group_list, value=max(No_beam_group_list), label="Beam groups", info="Divide beams into equal groups", visible=False), | |
| num_return_sequences: gr.Slider(maximum=no_of_beams) | |
| } | |
| #-----------ON SELECTING TOP P / TOP K-------------- | |
| def top_p_switch(input_p_box): | |
| value = input_p_box | |
| if value: | |
| return {top_p: gr.Slider(visible = True)} | |
| else: | |
| return {top_p: gr.Slider(visible = False)} | |
| def top_k_switch(input_k_box): | |
| value = input_k_box | |
| if value: | |
| return {top_k: gr.Slider(visible = True)} | |
| else: | |
| return {top_k: gr.Slider(visible = False)} | |
| #-----------ON SELECTING BEAM TEMPERATURE-------------- | |
| def beam_temp_switch (input): | |
| value = input | |
| if value: | |
| return {temperature: gr.Slider (visible=True)} | |
| else: | |
| return {temperature: gr.Slider (visible=False)} | |
| #-----------ON COOOSING STRATEGY: HIDE/DISPLAY PARAMS ----------- | |
| def select_strategy(input_strategy): | |
| global chosen_strategy | |
| chosen_strategy = input_strategy | |
| if chosen_strategy == "Beam Search": | |
| return {n_beams: gr.Slider(visible=True), | |
| num_return_sequences: gr.Slider(visible=True), | |
| beam_temperature: gr.Checkbox(visible=True), | |
| early_stopping: gr.Checkbox(visible=True), | |
| length_penalty: gr.Slider(visible=True), | |
| beam_groups: gr.Dropdown(visible=False), | |
| diversity_penalty: gr.Slider(visible=False), | |
| temperature: gr.Slider (visible=False), | |
| top_k: gr.Slider(visible=False), | |
| top_p: gr.Slider(visible=False), | |
| top_k_box: gr.Checkbox(visible = False), | |
| top_p_box: gr.Checkbox(visible = False), | |
| penalty_alpha: gr.Slider (visible=False) | |
| } | |
| if chosen_strategy == "Sampling": | |
| if top_k_box == True: | |
| {top_k: gr.Slider(visible = True)} | |
| if top_p_box == True: | |
| {top_p: gr.Slider(visible = True)} | |
| return { | |
| temperature: gr.Slider (visible=True), | |
| top_p: gr.Slider(visible=False), | |
| top_k: gr.Slider(visible=False), | |
| n_beams: gr.Slider(visible=False), | |
| beam_groups: gr.Dropdown(visible=False), | |
| diversity_penalty: gr.Slider(visible=False), | |
| num_return_sequences: gr.Slider(visible=False), | |
| beam_temperature: gr.Checkbox(visible=False), | |
| early_stopping: gr.Checkbox(visible=False), | |
| length_penalty: gr.Slider(visible=False), | |
| top_p_box: gr.Checkbox(visible = True, value=False), | |
| top_k_box: gr.Checkbox(visible = True, value=False), | |
| penalty_alpha: gr.Slider (visible=False) | |
| } | |
| if chosen_strategy == "Diversity Beam Search": | |
| return {n_beams: gr.Slider(visible=True), | |
| beam_groups: gr.Dropdown(visible=True), | |
| diversity_penalty: gr.Slider(visible=True), | |
| num_return_sequences: gr.Slider(visible=True), | |
| length_penalty: gr.Slider(visible=True), | |
| beam_temperature: gr.Checkbox(visible=False), | |
| early_stopping: gr.Checkbox(visible=True), | |
| temperature: gr.Slider (visible=False), | |
| top_k: gr.Slider(visible=False), | |
| top_p: gr.Slider(visible=False), | |
| top_k_box: gr.Checkbox(visible = False), | |
| top_p_box: gr.Checkbox(visible = False), | |
| penalty_alpha: gr.Slider (visible=False), | |
| } | |
| if chosen_strategy == "Contrastive": | |
| if top_k_box: | |
| {top_k: gr.Slider(visible = True)} | |
| return { | |
| temperature: gr.Slider (visible=True), | |
| penalty_alpha: gr.Slider (visible=True), | |
| top_p: gr.Slider(visible=False), | |
| #top_k: gr.Slider(visible = True) if top_k_box | |
| #top_k: gr.Slider(visible=False), | |
| n_beams: gr.Slider(visible=False), | |
| beam_groups: gr.Dropdown(visible=False), | |
| diversity_penalty: gr.Slider(visible=False), | |
| num_return_sequences: gr.Slider(visible=False), | |
| beam_temperature: gr.Checkbox(visible=False), | |
| early_stopping: gr.Checkbox(visible=False), | |
| length_penalty: gr.Slider(visible=False), | |
| top_p_box: gr.Checkbox(visible = False), | |
| top_k_box: gr.Checkbox(visible = True) | |
| } | |
| def clear(): | |
| print ("") | |
| #------------------MAIN BLOCKS DISPLAY--------------- | |
| with gr.Blocks() as demo: | |
| No_beam_group_list = [2] | |
| text = gr.Textbox( | |
| label="Prompt", | |
| value="It's a rainy day today", | |
| ) | |
| tokenizer = AutoTokenizer.from_pretrained("gpt2") | |
| model = AutoModelForCausalLM.from_pretrained("gpt2", pad_token_id=tokenizer.eos_token_id, cache_dir=cache_dir).to(torch_device) | |
| with gr.Row(): | |
| with gr.Column (scale=0, min_width=200) as Models_Strategy: | |
| model_selected = gr.Radio (["gpt2", "Gemma 2"], label="ML Model", value="gpt2") | |
| strategy_selected = gr.Radio (["Sampling", "Beam Search", "Diversity Beam Search","Contrastive"], label="Search strategy", value = "Sampling", interactive=True) | |
| with gr.Column (scale=0, min_width=250) as Beam_Params: | |
| n_steps = gr.Slider( | |
| label="Number of steps/tokens", minimum=1, maximum=100, step=1, value=20 | |
| ) | |
| n_beams = gr.Slider( | |
| label="Number of beams", minimum=2, maximum=100, step=1, value=4, visible=False | |
| ) | |
| #----------------Dropdown----------------- | |
| beam_groups = gr.Dropdown(No_beam_group_list, value=2, label="Beam groups", info="Divide beams into equal groups", visible=False | |
| ) | |
| diversity_penalty = gr.Slider( | |
| label="Group diversity penalty", minimum=0.1, maximum=2, step=0.1, value=0.8, visible=False | |
| ) | |
| num_return_sequences = gr.Slider( | |
| label="Number of return sequences", minimum=1, maximum=3, step=1, value=2, visible=False | |
| ) | |
| temperature = gr.Slider( | |
| label="Temperature", minimum=0.1, maximum=3, step=0.1, value=0.6, visible = True | |
| ) | |
| top_k = gr.Slider( | |
| label="Top_K", minimum=1, maximum=50, step=1, value=5, visible = False | |
| ) | |
| top_p = gr.Slider( | |
| label="Top_P", minimum=0.1, maximum=3, step=0.1, value=0.3, visible = False | |
| ) | |
| penalty_alpha = gr.Slider( | |
| label="Contrastive penalty α", minimum=0.1, maximum=2, step=0.1, value=0.6, visible=False | |
| ) | |
| top_p_box = gr.Checkbox(label="Top P", info="Turn on Top P", visible = True, interactive=True) | |
| top_k_box = gr.Checkbox(label="Top K", info="Turn on Top K", visible = True, interactive=True) | |
| early_stopping = gr.Checkbox(label="Early stopping", info="Stop with heuristically chosen good result", visible = False, interactive=True) | |
| beam_temperature = gr.Checkbox(label="Beam Temperature", info="Turn on sampling", visible = False, interactive=True) | |
| with gr.Column(scale=0, min_width=200): | |
| length_penalty = gr.Slider( | |
| label="Length penalty", minimum=-3, maximum=3, step=0.5, value=0, info="'+' more, '-' less no. of words", visible = False, interactive=True | |
| ) | |
| no_repeat_ngram_size = gr.Slider( | |
| label="No repeat n-gram phrase size", minimum=0, maximum=8, step=1, value=4, info="Not to repeat 'n' words" | |
| ) | |
| repetition_penalty = gr.Slider( | |
| label="Repetition penalty", minimum=0, maximum=3, step=1, value=float(0), info="Prior context based penalty for unique text" | |
| ) | |
| #----------ON SELECTING/CHANGING: RETURN SEEQUENCES/NO OF BEAMS/BEAM GROUPS/TEMPERATURE-------- | |
| model_selected.change( | |
| fn=load_model, inputs=[model_selected], outputs=[] | |
| ) | |
| #num_return_sequences.change( | |
| #fn=change_num_return_sequences, inputs=[n_beams,num_return_sequences], outputs=num_return_sequences | |
| #) | |
| n_beams.change( | |
| fn=popualate_beam_groups, inputs=[n_beams], outputs=[beam_groups,num_return_sequences] | |
| ) | |
| strategy_selected.change(fn=select_strategy, inputs=strategy_selected, outputs=[n_beams,beam_groups,length_penalty,diversity_penalty,num_return_sequences,temperature,early_stopping,beam_temperature,penalty_alpha,top_p,top_k,top_p_box,top_k_box]) | |
| beam_temperature.change (fn=beam_temp_switch, inputs=beam_temperature, outputs=temperature) | |
| top_p_box.change (fn=top_p_switch, inputs=top_p_box, outputs=top_p) | |
| top_k_box.change (fn=top_k_switch, inputs=top_k_box, outputs=top_k) | |
| #-------------GENERATE BUTTON------------------- | |
| button = gr.Button("Generate") | |
| out_markdown = gr.Textbox() | |
| button.click( | |
| fn = generate, | |
| inputs=[text, n_steps, n_beams, beam_groups, diversity_penalty, length_penalty, num_return_sequences, temperature, no_repeat_ngram_size, repetition_penalty, early_stopping, beam_temperature, top_p, top_k,penalty_alpha,top_p_box,top_k_box,strategy_selected,model_selected], | |
| outputs=[out_markdown] | |
| ) | |
| cleared = gr.Button ("Clear") | |
| cleared.click (fn=clear, inputs=[], outputs=[out_markdown]) | |
| demo.launch() | |