app.py

# 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("openai-community/gpt2")
        model = AutoModelForCausalLM.from_pretrained("openai-community/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("openai-community/gpt2")
    model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2", pad_token_id=tokenizer.eos_token_id)

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