Add generation configurations to chatbot interface

app.py

@@ -25,7 +25,7 @@ terminators = [
 ]
 
 @spaces.GPU(duration=120)
-def generate_both(system_prompt, input_text, base_chatbot, new_chatbot):
+def generate_both(system_prompt, input_text, base_chatbot, new_chatbot, max_new_tokens=2048, temperature=0.2, top_p=0.9, repetition_penalty=1.1):
     base_text_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
     new_text_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True)
 
@@ -60,22 +60,24 @@ def generate_both(system_prompt, input_text, base_chatbot, new_chatbot):
     base_generation_kwargs = dict(
         input_ids=base_input_ids,
         streamer=base_text_streamer,
-        max_new_tokens=2048,
+        max_new_tokens=max_new_tokens,
         eos_token_id=terminators,
         pad_token_id=tokenizer.eos_token_id,
-        do_sample=True,
-        temperature=0.2,
-        top_p=0.9,
+        do_sample=True if temperature > 0 else False,
+        temperature=temperature,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
     )
     new_generation_kwargs = dict(
         input_ids=new_input_ids,
         streamer=new_text_streamer,
-        max_new_tokens=2048,
+        max_new_tokens=max_new_tokens,
         eos_token_id=terminators,
         pad_token_id=tokenizer.eos_token_id,
-        do_sample=True,
-        temperature=0.2,
-        top_p=0.9,
+        do_sample=True if temperature > 0 else False,
+        temperature=temperature,
+        top_p=top_p,
+        repetition_penalty=repetition_penalty,
     )
 
     base_thread = Thread(target=base_model.generate, kwargs=base_generation_kwargs)
@@ -111,16 +113,21 @@ with gr.Blocks(title="Arabic-ORPO-Llama3") as demo:
         gr.HTML("<center><h1>Arabic Chatbot Comparison</h1></center>")
         system_prompt = gr.Textbox(lines=1, label="System Prompt", value="أنت متحدث لبق باللغة العربية!", rtl=True, text_align="right", show_copy_button=True)
         with gr.Row(variant="panel"):
-            base_chatbot = gr.Chatbot(label=base_model_id, rtl=True, likeable=True, show_copy_button=True)
-            new_chatbot = gr.Chatbot(label=new_model_id, rtl=True, likeable=True, show_copy_button=True)
+            base_chatbot = gr.Chatbot(label=base_model_id, rtl=True, likeable=True, show_copy_button=True, height=500)
+            new_chatbot = gr.Chatbot(label=new_model_id, rtl=True, likeable=True, show_copy_button=True, height=500)
         with gr.Row(variant="panel"):
             with gr.Column(scale=1):
                 submit_btn = gr.Button(value="Generate", variant="primary")
                 clear_btn = gr.Button(value="Clear", variant="secondary")
             input_text = gr.Textbox(lines=1, label="", value="مرحبا", rtl=True, text_align="right", scale=3, show_copy_button=True)
-
-    input_text.submit(generate_both, inputs=[system_prompt, input_text, base_chatbot, new_chatbot], outputs=[base_chatbot, new_chatbot])
-    submit_btn.click(generate_both, inputs=[system_prompt, input_text, base_chatbot, new_chatbot], outputs=[base_chatbot, new_chatbot])
+        with gr.Accordion(label="Generation Configurations", open=False):
+            max_new_tokens = gr.Slider(minimum=128, maximum=4096, value=2048, label="Max New Tokens", step=128)
+            temperature = gr.Slider(minimum=0.0, maximum=1.0, value=0.2, label="Temperature", step=0.01)
+            top_p = gr.Slider(minimum=0.0, maximum=1.0, value=0.9, label="Top-p", step=0.01)
+            repetition_penalty = gr.Slider(minimum=0.1, maximum=2.0, value=1.1, label="Repetition Penalty", step=0.1)
+
+    input_text.submit(generate_both, inputs=[system_prompt, input_text, base_chatbot, new_chatbot, max_new_tokens, temperature, top_p, repetition_penalty], outputs=[base_chatbot, new_chatbot])
+    submit_btn.click(generate_both, inputs=[system_prompt, input_text, base_chatbot, new_chatbot, max_new_tokens, temperature, top_p, repetition_penalty], outputs=[base_chatbot, new_chatbot])
     clear_btn.click(clear, outputs=[base_chatbot, new_chatbot])
 
 demo.launch()