Update app.py

app.py

@@ -8,13 +8,14 @@ model = AutoModelForCausalLM.from_pretrained("unsloth/Llama-3.2-1B")
 # Use a pipeline for text generation
 text_gen_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Text generation function with additional parameters
-def generate_text(prompt, max_length=50, temperature=0.7, top_p=0.9, top_k=50):
+# Text generation function with repetition penalty
+def generate_text(prompt, max_length=50, temperature=0.7, top_p=0.9, top_k=50, repetition_penalty=1.2):
     generated_text = text_gen_pipeline(prompt, 
                                        max_length=max_length, 
                                        temperature=temperature, 
                                        top_p=top_p, 
                                        top_k=top_k, 
+                                       repetition_penalty=repetition_penalty,  # Penalty added
                                        num_return_sequences=1)
     return generated_text[0]['generated_text']
 
@@ -37,6 +38,9 @@ with gr.Blocks() as demo:
     # Slider for top_k (controls diversity)
     top_k_input = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top-k (sampling diversity)")
     
+    # Slider for repetition penalty
+    repetition_penalty_input = gr.Slider(minimum=1.0, maximum=2.0, value=1.2, step=0.1, label="Repetition Penalty")
+    
     # Output box for the generated text
     output_text = gr.Textbox(label="Generated Text")
     
@@ -44,7 +48,9 @@ with gr.Blocks() as demo:
     generate_button = gr.Button("Generate")
     
     # Action on button click
-    generate_button.click(generate_text, inputs=[prompt_input, max_length_input, temperature_input, top_p_input, top_k_input], outputs=output_text)
+    generate_button.click(generate_text, 
+                          inputs=[prompt_input, max_length_input, temperature_input, top_p_input, top_k_input, repetition_penalty_input], 
+                          outputs=output_text)
 
 # Launch the app
 demo.launch()