Update app.py

app.py

@@ -69,7 +69,7 @@ def truncate_to_last_words(text, n):
     )
     return result
 
-def generate_text(model_name, input_text, decoding_method, max_length, use_last_words, num_last_words):
+def generate_text(model_name, input_text, decoding_method, max_length, use_last_words, num_last_words, repetition_penalty):
     global stop_generation
     stop_generation = False
 
@@ -90,7 +90,8 @@ def generate_text(model_name, input_text, decoding_method, max_length, use_last_
         generation_args = {
             "max_new_tokens": max_length,
             "pad_token_id": tokenizer.eos_token_id,
-            "do_sample": decoding_method == "stochastic"
+            "do_sample": decoding_method == "stochastic",
+            "repetition_penalty": repetition_penalty
         }
 
         if decoding_method == "stochastic":
@@ -139,6 +140,11 @@ with gr.Blocks(title=title) as demo:
                 label="Max Tokens"
             )
 
+            repetition_penalty = gr.Slider(
+                minimum=1.0, maximum=2.0, value=1.2, step=0.1,
+                label="Repetition Penalty (1.0=no penalty, higher=less repetition)"
+            )
+
             use_last_words = gr.Checkbox(
                 label="Use Last N Words",
                 value=False
@@ -177,7 +183,7 @@ with gr.Blocks(title=title) as demo:
 
     generate_btn.click(
         fn=generate_text,
-        inputs=[model_dropdown, input_text, decoding_method, max_length, use_last_words, num_last_words],
+        inputs=[model_dropdown, input_text, decoding_method, max_length, use_last_words, num_last_words, repetition_penalty],
         outputs=output_text
     )
 
@@ -185,5 +191,4 @@ with gr.Blocks(title=title) as demo:
 
 # Launch
 if __name__ == "__main__":
-    demo.launch()
-
+    demo.launch()