add repetition penalty

app.py

@@ -160,6 +160,7 @@ def evaluate(
     input=None,
     temperature=0.7,
     max_tokens=384,
+    repetition_penalty=1.0,
 ):
     num_beams: int = 1
     top_p: float = 1.0
@@ -186,13 +187,17 @@ def evaluate(
                 )
             except Exception as e:
                 print(e)
-        return f"please reduce the input length. Currently, {len(inputs['input_ids'][0])} tokens are used.", gr.update(interactive=True), gr.update(interactive=True)
+        return (
+            f"please reduce the input length. Currently, {len(inputs['input_ids'][0])} tokens are used.",
+            gr.update(interactive=True),
+            gr.update(interactive=True),
+        )
     input_ids = inputs["input_ids"].to(device)
     generation_config = GenerationConfig(
         temperature=temperature,
         top_p=top_p,
         top_k=top_k,
-        repetition_penalty=1.5,
+        repetition_penalty=repetition_penalty,
         num_beams=num_beams,
         pad_token_id=tokenizer.pad_token_id,
         eos_token=tokenizer.eos_token_id,
@@ -203,7 +208,7 @@ def evaluate(
             generation_config=generation_config,
             return_dict_in_generate=True,
             output_scores=True,
-            max_new_tokens=max_tokens-len(input_ids),
+            max_new_tokens=max_tokens - len(input_ids),
         )
     s = generation_output.sequences[0]
     output = tokenizer.decode(s, skip_special_tokens=True)
@@ -292,6 +297,14 @@ with gr.Blocks(
                 interactive=True,
                 label="Max length (Pre-prompt + instruction + input + output))",
             )
+            repetition_penalty = gr.Slider(
+                minimum=1.0,
+                maximum=5.0,
+                value=1.2,
+                step=0.05,
+                interactive=True,
+                label="Repetition penalty",
+            )
 
     with gr.Column(elem_id="user_consent_container") as user_consent_block:
         # Get user consent
@@ -334,14 +347,14 @@ with gr.Blocks(
     inputs.submit(no_interactive, [], [submit_button, clear_button])
     inputs.submit(
         evaluate,
-        [instruction, inputs, temperature, max_tokens],
+        [instruction, inputs, temperature, max_tokens, repetition_penalty],
         [outputs, submit_button, clear_button],
     )
     submit_button.click(no_interactive, [], [submit_button, clear_button])
     submit_button.click(
         evaluate,
         [instruction, inputs, temperature, max_tokens],
-        [outputs, submit_button, clear_button],
+        [outputs, submit_button, clear_button, repetition_penalty],
     )
     clear_button.click(reset_textbox, [], [instruction, inputs, outputs], queue=False)