Completed model comparison + added private models support + custom params support

app.py

@@ -30,17 +30,19 @@ MODEL_CLASSES = {
     "Custom Model" : (AutoModelForCausalLM, AutoTokenizer),
     }
 
-def load_model(model_name, custom_model_path): 
+CHOICES = sorted(list(CHECKPOINTS.keys())[:3])
+
+def load_model(model_name, custom_model_path, token): 
     try:
       model_class, tokenizer_class = MODEL_CLASSES[model_name]
       model_path = CHECKPOINTS[model_name]
     
     except KeyError:
       model_class, tokenizer_class = MODEL_CLASSES['Custom Model']
-      model_path = custom_model_path
+      model_path = custom_model_path or model_name 
     
-    model = model_class.from_pretrained(model_path)
-    tokenizer = tokenizer_class.from_pretrained(model_path)
+    model = model_class.from_pretrained(model_path, use_auth_token=token)
+    tokenizer = tokenizer_class.from_pretrained(model_path, use_auth_token=token)
 
     tokenizer.pad_token = tokenizer.eos_token
     model.config.pad_token_id = model.config.eos_token_id
@@ -67,6 +69,7 @@ def adjust_length_to_model(length, max_sequence_length):
     return length
 
 def generate(model_name,
+             token,
              custom_model_path,
              input_sentence,
              length = 75,
@@ -88,7 +91,7 @@ def generate(model_name,
     set_seed(seed, n_gpu)
 
     # Load model
-    model, tokenizer = load_model(model_name, custom_model_path)
+    model, tokenizer = load_model(model_name, custom_model_path, token)
     model.to(device)
 
     #length = adjust_length_to_model(length, max_sequence_length=model.config.max_position_embeddings)
@@ -126,6 +129,7 @@ def generate(model_name,
 
     return generated_sequences[0]
 
+
 def show_mode(mode):
   if mode == 'Single Model':
     return (
@@ -174,7 +178,7 @@ def show_search_bar(value):
             gr.update(visible=False)
             )
 
-def search_model(model_name):
+def search_model(model_name, token):
   api = HfApi()
 
   model_args = ModelSearchArguments()
@@ -182,7 +186,7 @@ def search_model(model_name):
       task=model_args.pipeline_tag.TextGeneration,
       library=model_args.library.PyTorch)
 
-  results = api.list_models(filter=filt, search=model_name)
+  results = api.list_models(filter=filt, search=model_name, use_auth_token=token)
   model_list = [model.modelId for model in results]
 
   return gr.update(visible=True,
@@ -190,6 +194,12 @@ def search_model(model_name):
                    label='Choose the model',
                    )
 
+def show_api_key_textbox(checkbox):
+  if checkbox:
+    return gr.update(visible=True)
+  else:
+    return gr.update(visible=False)
+
 def forward_model_choice(model_choice_path):
   return (model_choice_path,
           model_choice_path)
@@ -200,16 +210,30 @@ def auto_complete(input, generated):
   completed_prompt = {"text": output, "entities": output_spans}   
   return completed_prompt
 
-def process_user_input(model, custom_model_path, input):
+def process_user_input(model, 
+                       token, 
+                       custom_model_path, 
+                       input, 
+                       length, 
+                       temperature, 
+                       top_p, 
+                       top_k):
   warning = 'Please enter a valid prompt.'
   if input == None:
     generated = warning
   else:
-    generated = generate(model, custom_model_path, input)
-    generated_with_spans = auto_complete(input, generated)
+    generated = generate(model_name=model,
+                         token=token, 
+                         custom_model_path=custom_model_path, 
+                         input_sentence=input,
+                         length=length,
+                         temperature=temperature,
+                         top_p=top_p,
+                         top_k=top_k)
+    generated_with_spans = auto_complete(input=input, generated=generated)
 
   return (
-      generated_with_spans,
+      gr.update(value=generated_with_spans),
       gr.update(visible=True),
       gr.update(visible=True),
       input,
@@ -264,11 +288,55 @@ def upload_flag(*args):
   if flagging_callback.flag(list(args), flag_option = None):
       return gr.update(visible=True)
 
+def forward_model_choice_multi(model_choice_path):
+  CHOICES.append(model_choice_path)
+  return gr.update(choices = CHOICES)
+
+def process_user_input_multi(models, 
+                             input, 
+                             token,
+                             length,
+                             temperature,
+                             top_p,
+                             top_k):
+  warning = 'Please enter a valid prompt.'
+  if input == None:
+    generated = warning
+  else:
+    generated_dict= {model:generate(model_name=model,
+                                    token=token, 
+                                    custom_model_path=None, 
+                                    input_sentence=input,
+                                    length=length,
+                                    temperature=temperature,
+                                    top_p=top_p,
+                                    top_k=top_k) for model in sorted(models)}
+    generated_with_spans_dict = {model:auto_complete(input, generated) for model,generated in generated_dict.items()}
+
+    update_outputs = [gr.HighlightedText.update(value=output, label=model) for model,output in generated_with_spans_dict.items()]
+    update_hide = [gr.HighlightedText.update(visible=False) for i in range(10-len(models))]
+    return update_outputs + update_hide
+
+def show_choices_multi(models):
+  update_show = [gr.HighlightedText.update(visible=True) for model in sorted(models)]
+  update_hide = [gr.HighlightedText.update(visible=False,value=None, label=None) for i in range(10-len(models))]
+
+  return update_show + update_hide
+
+def show_params(checkbox):
+  if checkbox == True:
+    return gr.update(visible=True)
+  else: 
+    return gr.update(visible=False)
+
 CSS = """
 #inside_group {
   padding-top: 0.6em;
   padding-bottom: 0.6em;
 }
+#pw textarea {
+    -webkit-text-security: disc;
+}
 """
 
 with gr.Blocks(css=CSS) as demo:
@@ -286,9 +354,12 @@ with gr.Blocks(css=CSS) as demo:
                                                  organization = "fsdlredteam",
                                                  private = True )
   
-  gr.Markdown("# Project Interface proposal")
+  gr.Markdown("# FSDL 2022 Red-Teaming Open-Source Models Interface")
+  gr.Markdown("<img src=https://i.imgur.com/ZxbbLUQ.png>")
+
   gr.Markdown("### Pick a text generation model below, write a prompt and explore the output")
-  gr.Markdown("### Or compare multiple models")
+  gr.Markdown("### Or compare the output of multiple models at the same time")
+
   
   choose_mode = gr.Radio(choices=['Single Model', "Multi-Model"],
                          value='Single Model',
@@ -297,6 +368,12 @@ with gr.Blocks(css=CSS) as demo:
                          show_label=False)
 
   with gr.Group() as single_model:
+
+    gr.Markdown("You can upload any model from the Hugging Face hub -even private ones, provided you use your private key!")
+    gr.Markdown("Write your prompt or alternatively use one from the [RealToxicityPrompts](https://allenai.org/data/real-toxicity-prompts) dataset")
+    gr.Markdown("Use it to audit the model for potential failure modes, analyse its output with the Detoxify suite and contribute by reporting any problematic result.")
+    gr.Markdown("Beware ! Generation can take up to a few minutes with very large models.")
+    
     with gr.Row():
 
       with gr.Column(scale=1): # input & prompts dataset exploration 
@@ -315,11 +392,44 @@ with gr.Blocks(css=CSS) as demo:
 
         randomize_button = gr.Button('Show another subset', visible=False, elem_id="inside_group")
 
+        show_params_checkbox_single = gr.Checkbox(label='Set custom params', 
+                                                  interactive=True, 
+                                                  value=False)
+
+        with gr.Box(visible=False) as params_box_single:
+
+          length_single = gr.Slider(label='Output length',
+                                    visible=True,
+                                    interactive=True,
+                                    minimum=50,
+                                    maximum=200,
+                                    value=75)
+          
+          top_k_single = gr.Slider(label='top_k',
+                                   visible=True,
+                                   interactive=True,
+                                   minimum=1,
+                                   maximum=100,
+                                   value=50)
+          
+          top_p_single = gr.Slider(label='top_p',
+                                   visible=True,
+                                   interactive=True,
+                                   minimum=0.1,
+                                   maximum=1,
+                                   value=0.95)
+
+          temperature_single = gr.Slider(label='temperature',
+                                         visible=True,
+                                         interactive=True,
+                                         minimum=0.1,
+                                         maximum=1,
+                                         value=0.7)  
+
             
       with gr.Column(scale=1): # Model choice & output 
         gr.Markdown("### 2. Evaluate output")
 
-        
         model_radio = gr.Radio(choices=list(CHECKPOINTS.keys()),
                               label='Model',
                               interactive=True,
@@ -331,11 +441,19 @@ with gr.Blocks(css=CSS) as demo:
                                 elem_id="inside_group")
         model_drop = gr.Dropdown(visible=False)
 
-        generate_button = gr.Button('Submit your prompt')
+        private_checkbox = gr.Checkbox(visible=True,label="Private Model ?", elem_id="inside_group")
+
+        api_key_textbox = gr.Textbox(label="Enter your AUTH TOKEN below",
+                                      value=None, 
+                                      interactive=True,
+                                      visible=False,
+                                      elem_id="pw")
 
-        output_spans = gr.HighlightedText(visible=True, label="Generated text", elem_id="inside_group")
+        generate_button = gr.Button('Submit your prompt', elem_id="inside_group")
 
-        flag_button = gr.Button("Report output here", visible=False)
+        output_spans = gr.HighlightedText(visible=True, label="Generated text")
+
+        flag_button = gr.Button("Report output here", visible=False, elem_id="inside_group")
 
     with gr.Row(): # Flagging 
       
@@ -373,9 +491,94 @@ with gr.Blocks(css=CSS) as demo:
                                   visible=False,
                                   elem_id="inside_group")
       
-  with gr.Group() as multi_model:
-    gr.Markdown("Model comparison will be here")
+  with gr.Group(visible=False) as multi_model:
+    model_list = list()
+
+    gr.Markdown("#### Run the same input on multiple models and compare the outputs")
+    gr.Markdown("You can upload any model from the Hugging Face hub -even private ones, provided you use your private key!")
+    gr.Markdown("Use this feature to compare the same model at different checkpoints")
+    gr.Markdown('Or to benchmark your model against another one as a reference.')
+    gr.Markdown("Beware ! Generation can take up to a few minutes with very large models.")
+
+    with gr.Row(elem_id="inside_group"):
+      with gr.Column():
+        models_multi = gr.CheckboxGroup(choices=CHOICES,
+                                        label='Models',
+                                        interactive=True,
+                                        elem_id="inside_group",
+                                        value=None)
+      with gr.Column():
+        generate_button_multi = gr.Button('Submit your prompt',elem_id="inside_group")
+
+        show_params_checkbox_multi = gr.Checkbox(label='Set custom params', 
+                                                  interactive=True, 
+                                                  value=False)
+
+        with gr.Box(visible=False) as params_box_multi:
+        
+          length_multi = gr.Slider(label='Output length',
+                                   visible=True,
+                                   interactive=True,
+                                   minimum=50,
+                                   maximum=200,
+                                   value=75)
+          
+          top_k_multi = gr.Slider(label='top_k',
+                                  visible=True,
+                                  interactive=True,
+                                  minimum=1,
+                                  maximum=100,
+                                  value=50)
+          
+          top_p_multi = gr.Slider(label='top_p',
+                                  visible=True,
+                                  interactive=True,
+                                  minimum=0.1,
+                                  maximum=1,
+                                  value=0.95)
+
+          temperature_multi = gr.Slider(label='temperature',
+                                        visible=True,
+                                        interactive=True,
+                                        minimum=0.1,
+                                        maximum=1,
+                                        value=0.7)  
+
+    with gr.Row(elem_id="inside_group"):
+
+      with gr.Column(elem_id="inside_group", scale=1):
+        input_text_multi = gr.Textbox(label="Write your prompt below.", 
+                                      interactive=True, 
+                                      lines=4, 
+                                      elem_id="inside_group")
+    
+      with gr.Column(elem_id="inside_group", scale=1):
+        search_bar_multi = gr.Textbox(label="Search another model", 
+                                      interactive=True, 
+                                      visible=True, 
+                                      elem_id="inside_group")
+    
+        model_drop_multi = gr.Dropdown(visible=False,
+                                       show_progress=True,
+                                       elem_id="inside_group")
+        
+        private_checkbox_multi = gr.Checkbox(visible=True,label="Private Model ?")
+
+        api_key_textbox_multi = gr.Textbox(label="Enter your AUTH TOKEN below",
+                                           value=None, 
+                                           interactive=True,
+                                           visible=False,
+                                           elem_id="pw")
+
+    with gr.Row() as outputs_row:
+      for i in range(10):
+        output_spans_multi = gr.HighlightedText(visible=False, elem_id="inside_group")
+        model_list.append(output_spans_multi)
+
+    
+    gr.Markdown('App made during the FSDL course by Team53: Jean-Antoine, Sajenthan, Sashank, Kemp, Srihari, Astitwa')
 
+  # Single Model 
 
   choose_mode.change(fn=show_mode,
                      inputs=choose_mode,
@@ -398,16 +601,27 @@ with gr.Blocks(css=CSS) as demo:
                      outputs=[model_choice,search_bar])
   
   search_bar.submit(fn=search_model,
-                    inputs=search_bar,
+                    inputs=[search_bar,api_key_textbox],
                     outputs=model_drop,
                     show_progress=True)
   
+  private_checkbox.change(fn=show_api_key_textbox,
+                          inputs=private_checkbox,
+                          outputs=api_key_textbox)
+  
   model_drop.change(fn=forward_model_choice, 
                     inputs=model_drop, 
                     outputs=[model_choice,custom_model_path])
 
   generate_button.click(fn=process_user_input,
-                        inputs=[model_choice, custom_model_path, input_text],
+                        inputs=[model_choice, 
+                                api_key_textbox, 
+                                custom_model_path, 
+                                input_text,
+                                length_single,
+                                temperature_single,
+                                top_p_single,
+                                top_k_single],
                         outputs=[output_spans,
                                  toxi_button,
                                  flag_button, 
@@ -442,7 +656,45 @@ with gr.Blocks(css=CSS) as demo:
                                       user_comment,
                                       flag_choice],
                             outputs=success_message)
+  
+  show_params_checkbox_single.change(fn=show_params,
+                                     inputs=show_params_checkbox_single,
+                                     outputs=params_box_single)
+
+  # Model comparison
+
+  search_bar_multi.submit(fn=search_model,
+                          inputs=[search_bar_multi, api_key_textbox_multi],
+                          outputs=model_drop_multi,
+                          show_progress=True)
+  
+  show_params_checkbox_multi.change(fn=show_params,
+                                    inputs=show_params_checkbox_multi,
+                                    outputs=params_box_multi)
+  
+  private_checkbox_multi.change(fn=show_api_key_textbox,
+                                inputs=private_checkbox_multi,
+                                outputs=api_key_textbox_multi)
+  
+  model_drop_multi.change(fn=forward_model_choice_multi,
+                          inputs=model_drop_multi, 
+                          outputs=[models_multi])
+    
+  models_multi.change(fn=show_choices_multi,
+                      inputs=models_multi,
+                      outputs=model_list)
+    
+  generate_button_multi.click(fn=process_user_input_multi,
+                              inputs=[models_multi, 
+                                      input_text_multi, 
+                                      api_key_textbox_multi,
+                                      length_multi,
+                                      temperature_multi,
+                                      top_p_multi,
+                                      top_k_multi],
+                              outputs=model_list,
+                              show_progress=True) 
                             
 #demo.launch(debug=True)
 if __name__ == "__main__":
-    demo.launch(enable_queue=False)
+    demo.launch(enable_queue=False, debug=True)