Upload app.py

17-8-2024
v1.0

app.py

@@ -0,0 +1,407 @@
+
+# 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").to(torch_device)
+
+
+# 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").to(torch_device)
+    
+    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").to(torch_device)    
+        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").to(torch_device)    
+        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).to(torch_device)
+        #print (model_selected + " loaded")
+    
+    if model_selected == "Gemma 2":
+        tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b")
+        model = AutoModelForCausalLM.from_pretrained("google/gemma-2b").to(torch_device)
+
+
+
+#--------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()
+