Add other generation parameters

app.py

@@ -44,11 +44,11 @@ code_to_name = {
 }
 
 @st.cache
-def generate_sequence(pixel_values, lang_code, num_beams, temperature, top_p):
+def generate_sequence(pixel_values, lang_code, num_beams, temperature, top_p, do_sample, top_k, max_length):
     lang_code = language_mapping[lang_code]
-    output_ids = state.model.generate(input_ids=pixel_values, forced_bos_token_id=tokenizer.lang_code_to_id[lang_code], max_length=64, num_beams=num_beams, temperature=temperature, top_p = top_p)
+    output_ids = state.model.generate(input_ids=pixel_values, forced_bos_token_id=tokenizer.lang_code_to_id[lang_code], max_length=max_length, num_beams=num_beams, temperature=temperature, top_p = top_p, top_k=top_k, do_sample=do_sample)
     print(output_ids)
-    output_sequence = tokenizer.batch_decode(output_ids[0], skip_special_tokens=True, max_length=64)
+    output_sequence = tokenizer.batch_decode(output_ids[0], skip_special_tokens=True, max_length=max_length)
     return output_sequence
 
 def read_markdown(path, parent="./sections/"):
@@ -72,6 +72,9 @@ st.write(
 )
 
 st.sidebar.title("Generation Parameters")
+max_length = st.sidebar.number_input("Max Length", min_value=16, max_value=128, value=64, step=1, help="The maximum length of sequence to be generated.")
+do_sample = st.sidebar.checkbox("Sample", value=False, help="Sample from the model instead of using beam search.")
+top_k = st.sidebar.number_input("Top K", min_value=10, max_value=200, value=50, step=1, help="The number of highest probability vocabulary tokens to keep for top-k-filtering.")
 num_beams = st.sidebar.number_input(label="Number of Beams", min_value=2, max_value=10, value=4, step=1, help="Number of beams to be used in beam search.")
 temperature = st.sidebar.select_slider(label="Temperature", options = list(np.arange(0.0,1.1, step=0.1)), value=1.0, help ="The value used to module the next token probabilities.", format_func=lambda x: f"{x:.2f}")
 top_p = st.sidebar.select_slider(label = "Top-P", options = list(np.arange(0.0,1.1, step=0.1)),value=1.0, help="Nucleus Sampling : If set to float < 1, only the most probable tokens with probabilities that add up to :obj:`top_p` or higher are kept for generation.", format_func=lambda x: f"{x:.2f}")
@@ -150,7 +153,7 @@ lang_id = new_col2.selectbox(
 sequence = ['']
 if new_col2.button("Generate Caption", help="Generate a caption in the specified language."):
     with st.spinner("Generating Sequence..."):
-        sequence = generate_sequence(transformed_image, lang_id, num_beams, temperature, top_p)
+        sequence = generate_sequence(transformed_image, lang_id, num_beams, temperature, top_p, do_sample, top_k, max_length)
 # print(sequence)
 
 if sequence!=['']: