# 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()