update

app.py

@@ -6,7 +6,14 @@ from typing import List
 
 MAX_BASE_LLM_NUM = 20
 MIN_BASE_LLM_NUM = 3
-DESCRIPTIONS = """\
+SOURCE_MAX_LENGTH = 256
+DEFAULT_SOURCE_MAX_LENGTH = 128
+CANDIDATE_MAX_LENGTH = 256
+DEFAULT_CANDIDATE_MAX_LENGTH = 128
+FUSER_MAX_NEW_TOKENS = 512
+DEFAULT_FUSER_MAX_NEW_TOKENS = 256
+DESCRIPTIONS = """# LLM-BLENDER
+
 LLM-Blender is an innovative ensembling framework to attain consistently superior performance by leveraging the diverse strengths of multiple open-source large language models (LLMs). LLM-Blender cut the weaknesses through ranking and integrate the strengths through fusing generation to enhance the capability of LLMs.
 """
 EXAMPLES_DATASET = load_dataset("llm-blender/mix-instruct", split='validation', streaming=True)
@@ -21,7 +28,6 @@ for example in SHUFFLED_EXAMPLES_DATASET.take(100):
     CANDIDATE_EXAMPLES[example['instruction']+example['input']] = example['candidates']
 
 # Download ranker checkpoint
-os.system("ls -l /home/user/.local/lib/python3.10/site-packages/llm_blender")
 if not os.path.exists("pairranker-deberta-v3-large.zip"):
     os.system("gdown https://drive.google.com/uc?id=1EpvFu_qYY0MaIu0BAAhK-sYKHVWtccWg")
 if not os.path.exists("pairranker-deberta-v3-large"):
@@ -35,13 +41,13 @@ ranker_config.ranker_type = "pairranker"
 ranker_config.model_type = "deberta"
 ranker_config.model_name = "microsoft/deberta-v3-large" # ranker backbone
 ranker_config.load_checkpoint = "./pairranker-deberta-v3-large" # ranker checkpoint <your checkpoint path>
-ranker_config.source_maxlength = 128
-ranker_config.candidate_maxlength = 128
+ranker_config.source_maxlength = DEFAULT_SOURCE_MAX_LENGTH
+ranker_config.candidate_maxlength = DEFAULT_CANDIDATE_MAX_LENGTH
 ranker_config.n_tasks = 1 # number of singal that has been used to train the ranker. This checkpoint is trained using BARTScore only, thus being 1.
 fuser_config = llm_blender.GenFuserConfig()
 fuser_config.model_name = "llm-blender/gen_fuser_3b" # our pre-trained fuser
 fuser_config.max_length = 1024
-fuser_config.candidate_maxlength = 128
+fuser_config.candidate_maxlength = DEFAULT_CANDIDATE_MAX_LENGTH
 blender_config = llm_blender.BlenderConfig()
 blender_config.device = "cpu" # blender ranker and fuser device
 blender = llm_blender.Blender(blender_config, ranker_config, fuser_config)
@@ -74,7 +80,7 @@ def update_base_llm_dropdown_along_examples(dummy_text):
     ex_llm_outputs = {f"LLM-{i+1}": candidates[i]['text'] for i in range(len(candidates))}
     return ex_llm_outputs
     
-def check_save_ranker_inputs(inst, input, llm_outputs):
+def check_save_ranker_inputs(inst, input, llm_outputs, blender_config):
     if not inst and not input:
         raise gr.Error("Please enter instruction or input context")
     
@@ -87,23 +93,29 @@ def check_save_ranker_inputs(inst, input, llm_outputs):
         "candidates": list(llm_outputs.values()),
     }
 
-def check_fuser_inputs(blender_state, top_k_for_fuser, ranks):
+def check_fuser_inputs(blender_state, blender_config, ranks):
     pass
 
-def llms_rank(inst, input, llm_outputs):
+def llms_rank(inst, input, llm_outputs, blender_config):
     candidates = list(llm_outputs.values())
-    
-    return blender.rank(instructions=[inst], inputs=[input], candidates=[candidates])[0]
+    rank_params = {
+        "source_max_length": blender_config['source_max_length'],
+        "candidate_max_length": blender_config['candidate_max_length'],
+    }
+    ranks = blender.rank(instructions=[inst], inputs=[input], candidates=[candidates])[0]
+    return [ranks, ",  ".join([f"LLM-{i+1}: {rank}" for i, rank in enumerate(ranks)])]
 
-def display_ranks(ranks):
-    return ",  ".join([f"LLM-{i+1}: {rank}" for i, rank in enumerate(ranks)])
 
-def llms_fuse(blender_state, top_k_for_fuser, ranks):
+def llms_fuse(blender_state, blender_config, ranks):
     inst = blender_state['inst']
     input = blender_state['input']
     candidates = blender_state['candidates']
+    top_k_for_fuser = blender_config['top_k_for_fuser']
+    fuse_params = blender_config.copy()
+    del fuse_params["top_k_for_fuser"]
     top_k_candidates = get_topk_candidates_from_ranks([ranks], [candidates], top_k=top_k_for_fuser)[0]
-    return blender.fuse(instructions=[inst], inputs=[input], candidates=[top_k_candidates])[0]
+    fuser_outputs = blender.fuse(instructions=[inst], inputs=[input], candidates=[top_k_candidates], **fuse_params)[0]
+    return [fuser_outputs, fuser_outputs]
 
 def display_fuser_output(fuser_output):
     return fuser_output
@@ -111,16 +123,18 @@ def display_fuser_output(fuser_output):
         
 with gr.Blocks(theme='ParityError/Anime') as demo:
     gr.Markdown(DESCRIPTIONS)
+    gr.Markdown("## Input and Base LLMs")
     with gr.Row():
         with gr.Column():
             inst_textbox = gr.Textbox(lines=1, label="Instruction", placeholder="Enter instruction here", show_label=True)
             input_textbox = gr.Textbox(lines=4, label="Input Context", placeholder="Enter input context here", show_label=True)
         with gr.Column():
             saved_llm_outputs = gr.State(value={})
-            selected_base_llm_name_dropdown = gr.Dropdown(label="Base LLM",
-                choices=[f"LLM-{i+1}" for i in range(MIN_BASE_LLM_NUM)], value="LLM-1", show_label=True)
-            selected_base_llm_output = gr.Textbox(lines=4, label="LLM-1 (Click Save to save current content)",
-                placeholder="Enter LLM-1 output here", show_label=True)
+            with gr.Group():
+                selected_base_llm_name_dropdown = gr.Dropdown(label="Base LLM",
+                    choices=[f"LLM-{i+1}" for i in range(MIN_BASE_LLM_NUM)], value="LLM-1", show_label=True)
+                selected_base_llm_output = gr.Textbox(lines=4, label="LLM-1 (Click Save to save current content)",
+                    placeholder="Enter LLM-1 output here", show_label=True)
             with gr.Row():
                 base_llm_outputs_save_button = gr.Button('Save', variant='primary')
                 
@@ -136,28 +150,67 @@ with gr.Blocks(theme='ParityError/Anime') as demo:
                 )
     
     blender_state = gr.State(value={})
-    with gr.Tab("Ranking outputs"):
-        saved_rank_outputs = gr.State(value=[])
-        rank_outputs = gr.Textbox(lines=4, label="Ranking outputs", placeholder="Ranking outputs", show_label=True)
-    with gr.Tab("Fusing outputs"):
-        saved_fuse_outputs = gr.State(value=[])
+    saved_rank_outputs = gr.State(value=[])
+    saved_fuse_outputs = gr.State(value=[])
+    gr.Markdown("## Blender Outputs")
+    with gr.Group():
+        rank_outputs = gr.Textbox(lines=1, label="Ranking outputs", placeholder="Ranking outputs", show_label=True)
         fuser_outputs = gr.Textbox(lines=4, label="Fusing outputs", placeholder="Fusing outputs", show_label=True)
     with gr.Row():
-        rank_button = gr.Button('Rank LLM Outputs', variant='primary',
-            scale=1, min_width=0)
-        fuse_button = gr.Button('Fuse Top-K ranked outputs', variant='primary',
-            scale=1, min_width=0)
-        clear_button = gr.Button('Clear Blender', variant='primary',
-            scale=1, min_width=0)
+        rank_button = gr.Button('Rank LLM Outputs', variant='primary')
+        fuse_button = gr.Button('Fuse Top-K ranked outputs', variant='primary')
+        clear_button = gr.Button('Clear Blender Outputs', variant='primary')
+    blender_config = gr.State(value={
+        "source_max_length": DEFAULT_SOURCE_MAX_LENGTH,
+        "candidate_max_length": DEFAULT_CANDIDATE_MAX_LENGTH,
+        "top_k_for_fuser": 3,
+        "max_new_tokens": DEFAULT_FUSER_MAX_NEW_TOKENS,
+        "temperature": 0.7,
+        "top_p": 1.0,
+    })
         
     with gr.Accordion(label='Advanced options', open=False):
-        
+        source_max_length = gr.Slider(
+            label='Max length of Instruction + Input',
+            minimum=1,
+            maximum=SOURCE_MAX_LENGTH,
+            step=1,
+            value=DEFAULT_SOURCE_MAX_LENGTH,
+        )
+        candidate_max_length = gr.Slider(
+            label='Max length of LLM-Output Candidate',
+            minimum=1,
+            maximum=CANDIDATE_MAX_LENGTH,
+            step=1,
+            value=DEFAULT_CANDIDATE_MAX_LENGTH,
+        )
         top_k_for_fuser = gr.Slider(
-            label='Top k for fuser',
+            label='Top-k ranked candidates to fuse',
             minimum=1,
             maximum=3,
             step=1,
-            value=1,
+            value=3,
+        )
+        max_new_tokens = gr.Slider(
+            label='Max new tokens fuser can generate',
+            minimum=1,
+            maximum=FUSER_MAX_NEW_TOKENS,
+            step=1,
+            value=DEFAULT_FUSER_MAX_NEW_TOKENS,
+        )
+        temperature = gr.Slider(
+            label='Temperature of fuser generation',
+            minimum=0.1,
+            maximum=2.0,
+            step=0.1,
+            value=0.7,
+        )
+        top_p = gr.Slider(
+            label='Top-p of fuser generation',
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=1.0,
         )
     
     examples_dummy_textbox = gr.Textbox(lines=1, label="", placeholder="", show_label=False, visible=False)     
@@ -211,30 +264,22 @@ with gr.Blocks(theme='ParityError/Anime') as demo:
 
     rank_button.click(
         fn=check_save_ranker_inputs,
-        inputs=[inst_textbox, input_textbox, saved_llm_outputs],
+        inputs=[inst_textbox, input_textbox, saved_llm_outputs, blender_config],
         outputs=blender_state,
     ).success(
         fn=llms_rank,
-        inputs=[inst_textbox, input_textbox, saved_llm_outputs],
-        outputs=[saved_rank_outputs],
-    ).then(
-        fn=display_ranks,
-        inputs=[saved_rank_outputs],
-        outputs=rank_outputs,
+        inputs=[inst_textbox, input_textbox, saved_llm_outputs, blender_config],
+        outputs=[saved_rank_outputs, rank_outputs],
     )
     
     fuse_button.click(
         fn=check_fuser_inputs,
-        inputs=[blender_state, top_k_for_fuser, saved_rank_outputs],
+        inputs=[blender_state, blender_config, saved_rank_outputs],
         outputs=[],
     ).success(
         fn=llms_fuse,
-        inputs=[blender_state, top_k_for_fuser, saved_rank_outputs],
-        outputs=[saved_fuse_outputs],
-    ).then(
-        fn=display_fuser_output,
-        inputs=[saved_fuse_outputs],
-        outputs=fuser_outputs,
+        inputs=[blender_state, blender_config, saved_rank_outputs],
+        outputs=[saved_fuse_outputs, fuser_outputs],
     )
     
     clear_button.click(
@@ -243,6 +288,38 @@ with gr.Blocks(theme='ParityError/Anime') as demo:
         outputs=[rank_outputs, fuser_outputs, blender_state, saved_rank_outputs],
     )
     
+    # update blender config
+    source_max_length.change(
+        fn=lambda x, y: y.update({"source_max_length": x}) or y,
+        inputs=[source_max_length, blender_config],
+        outputs=blender_config,
+    )
+    candidate_max_length.change(
+        fn=lambda x, y: y.update({"candidate_max_length": x}) or y,
+        inputs=[candidate_max_length, blender_config],
+        outputs=blender_config,
+    )
+    top_k_for_fuser.change(
+        fn=lambda x, y: y.update({"top_k_for_fuser": x}) or y,
+        inputs=[top_k_for_fuser, blender_config],
+        outputs=blender_config,
+    )
+    max_new_tokens.change(
+        fn=lambda x, y: y.update({"max_new_tokens": x}) or y,
+        inputs=[max_new_tokens, blender_config],
+        outputs=blender_config,
+    )
+    temperature.change(
+        fn=lambda x, y: y.update({"temperature": x}) or y,
+        inputs=[temperature, blender_config],
+        outputs=blender_config,
+    )
+    top_p.change(
+        fn=lambda x, y: y.update({"top_p": x}) or y,
+        inputs=[top_p, blender_config],
+        outputs=blender_config,
+    )
+