app.py

# Gradio Params Playground

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
import gradio as gr
import os

token = os.environ.get("HF_TOKEN")

# Load default model as GPT2 and other models


tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")
model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")

tokenizer_gpt2 = AutoTokenizer.from_pretrained("openai-community/gpt2")
model_gpt2 = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")

tokenizer_gemma = AutoTokenizer.from_pretrained("google/gemma-2b")
model_gemma = AutoModelForCausalLM.from_pretrained("google/gemma-2b", token=token)

tokenizer_Mistral = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
model_Mistral = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")

# 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 = float(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 = float(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 = float(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 = float(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 = tokenizer_gpt2
        model = model_gpt2
        #print (model_selected + " loaded")
    
    if model_selected == "Gemma 2":
        tokenizer = tokenizer_gemma
        model = model_gemma
    
    if model_selected == "Mistral":
        tokenizer = tokenizer_Mistral
        model = model_Mistral



#--------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]
        
    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", "Mistral"], 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=1):
                
            text = gr.Textbox(
            label="Prompt",
            autoscroll=True,
            value="It's a rainy day today"
            )

            out_markdown = gr.Textbox(label="Output", autoscroll=True)

            button = gr.Button("Generate")  
              
            cleared = gr.Button ("Clear")
            cleared.click (fn=clear, inputs=[], outputs=[out_markdown])

        
        
        
        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.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]
    )


    with gr.Row():

        gr.Markdown (
            """
            ##
            # About Params Playground
            
            A space to tweak, test and learn generative model parameters for text output.
            
            ## Strategies:
        
            Given some text as input, a decoder-only model hunts for a continuation using various search strategies. (Whether the continuation makes sense or not is for us to determine.)
            
            Example:
            
            *Input: Today is a rainy day,*
            
            Option 1: and [probability score: 0.62]
            
            Option 2: so [probability score: 0.21]
            
            Option 3: ! [probability score: 0.73]
            
            
            ### 1. Greedy Search: 
            Picks up the next word/token carrying the highest probability score. The most well trodden path. Default for GPT2.
            
            In this illustrative example, since "!" has the highest probability score, a greedy strategy will output: Today is a rainy day!
            
            
            ### 2. Random Sampling: 
            Picks up any random path or trail of tokens to traverse and continue the input. To turn sampling on, use ```do_sample=True```
            
            *Temperature* - Increasing the temperature allows words with lesser probabilities to show up in the output. At ```temperature = 0```, search becomes 'greedy' for words with high probabilities.
            
            *Top_K*: Creates a small list of paths [tokens or words] to choose  from. In the above example, if set to 2, only Option 1 and 3 - the two top ranking tokens in terms of probabilities, will be available for random sampling.
            
            *Top_P*: Creates a small list of tokens based on the sum of their probability scores which should not exceed the Top P value. In the above example, if set to 0.80, only Option 3 will be available. If set to 1.5, Options 1 and 3 will be available. This metric can be used to make the output factually correct when the input is expecting facts like: "The capital of XYZ is [next token]"
            
            When used with temperature: Reducing temperature makes the search greedy.
                     
            
            ### 3. Simple Beam search: 
            Selects the branches (beams) going towards other heavy laden branch of fruits, to find the heaviest set among the branches in all. Akin to greedy search, but finds the total heaviest or largest route.
            
            If ```num_beams = 2```, every branch will divide into the top two scoring tokens at each step, and so on till the search ends.
            
            *Early Stopping*: Makes the search stop when a pre-determined criteria for ending the search is satisfied.
                      
            
            ### 4. Diversity Beam search: 
            Divided beams into groups of beams, and applies the diversity penalty. This makes the output more diverse and interesting. 
            
            *Group Diversity Penalty*: Used to instruct the next beam group to ignore the words/tokens already selected by previous groups.
                        
            
            ### 5. Contrastive search: 
            Uses the entire input context to create more interesting outputs.
            
            *Penalty Alpha*: When ```penalty_alpha=0```, search becomes greedy.
            
            Refer: https://huggingface.co/blog/introducing-csearch
                        
            
            ### Other parameters:
                        
            - Length penalty: Used to force the model to meet the expected output length.
                      
            - Repetition penalty: Used to force the model to avoid repetition.
                        
            - No repeat n-gram size: Used to force the model not to repeat the n-size set of words. Avoid setting to 1, as this forces no two words to be identical.
            
            
            **References**: 
            
            1. https://huggingface.co/blog/how-to-generate
            
            2. https://huggingface.co/docs/transformers/generation_strategies#decoding-strategies
            
            3. https://huggingface.co/docs/transformers/main/en/main_classes/text_generation

            

            
            """
            
        )

demo.launch()