Slider for repetition_penalty

app.py

@@ -14,13 +14,8 @@ checkpoint = "CohereForAI/aya-101"
 tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, torch_dtype=torch.bfloat16, low_cpu_mem_usage=True, device_map=device)
 
-#Set a the value of the repetition penalty
-#The higher the value, the less repetitive the generated text will be
-#Note that `repetition_penalty` has to be a strictly positive float
-repetition_penalty = 1.8
-
 @spaces.GPU
-def aya(text, max_new_tokens):
+def aya(text, max_new_tokens, repetition_penalty):
     model.to(device)
     inputs = tokenizer.encode(text, return_tensors="pt").to(device)
     outputs = model.generate(inputs, max_new_tokens=max_new_tokens, repetition_penalty=repetition_penalty)
@@ -33,9 +28,10 @@ def main():
         gr.Markdown(description)
         input_text = gr.Textbox(label="🗣️Input Text")
         max_new_tokens_slider = gr.Slider(minimum=150, maximum=1648, step=1, value=250, label="Size of your inputs and answer")
+        repetition_penalty_slider = gr.Slider(minimum=1.0, maximum=4.0, step=0.1, value=1.8, label="Repetition Penalty")
         submit_button = gr.Button("Use🌐Aya")        
         output_text = gr.Textbox(label="🌐Aya", interactive=False)
-        submit_button.click(fn=aya, inputs=[input_text, max_new_tokens_slider], outputs=output_text)
+        submit_button.click(fn=aya, inputs=[input_text, max_new_tokens_slider, repetition_penalty_slider], outputs=output_text)
     
     demo.launch()