Update app.py

app.py

@@ -69,6 +69,7 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
         top_k = gr.Slider(minimum=10, maximum=100, value=50, step=5, interactive=True, label="Top-k", info="Controls the number of highest probability tokens to consider for each step.")
         top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.70, step=0.05, interactive=True, label="Top-p", info="Controls the cumulative probability of the generated tokens.")
         temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.1, step=0.1, interactive=True, label="Temperature", info="Controls the randomness of the generated tokens.")
+        repetition_penalty = gr.Slider(minimum=1, maximum=2, value=1.5, step=0.1, interactive=True, label="Repetition Penalty", info="Higher values help the model to avoid repetition in text generation.")
         max_length = gr.Slider(minimum=10, maximum=500, value=100, step=10, interactive=True, label="Max Length", info="Controls the maximum length of the generated text.")
         smaple_from = gr.Slider(minimum=2, maximum=10, value=2, step=1, interactive=True, label="Sample From", info="Controls the number of generations that the reward model will sample from.")
 
@@ -78,7 +79,7 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
     def user(user_message, chat_history):
         return gr.update(value=user_message, interactive=True), chat_history + [["👤 " + user_message, None]]
 
-    def generate_response(user_msg, top_p, temperature, top_k, max_length, smaple_from, safety, chat_history):
+    def generate_response(user_msg, top_p, temperature, top_k, max_length, smaple_from, repetition_penalty, safety, chat_history):
 
         inputs = tokenizer(tokenizer.bos_token + user_msg + tokenizer.eos_token, return_tensors="pt").to(model.device)
 
@@ -86,6 +87,7 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
             bos_token_id=tokenizer.bos_token_id,
             pad_token_id=tokenizer.pad_token_id,
             eos_token_id=tokenizer.eos_token_id,
+            repetition_penalty=repetition_penalty,
             do_sample=True,
             early_stopping=True,
             top_k=top_k,
@@ -145,7 +147,7 @@ with gr.Blocks(theme='freddyaboulton/dracula_revamped') as demo:
             yield chat_history
 
     response = msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
-        generate_response, [msg, top_p, temperature, top_k, max_length, smaple_from, safety, chatbot], chatbot
+        generate_response, [msg, top_p, temperature, top_k, max_length, smaple_from, repetition_penalty, safety, chatbot], chatbot
     )
     response.then(lambda: gr.update(interactive=True), None, [msg], queue=False)
     msg.submit(lambda x: gr.update(value=''), None,[msg])