Update app.py

app.py

@@ -16,11 +16,28 @@ st.title("Image Captioning with ViT & GePpeTto 🇮🇹")
 
 st.sidebar.markdown("## Generation parameters")
 max_length = st.sidebar.number_input("Max length", value=20, min_value=1)
-num_beams = st.sidebar.number_input("Beam size", value=5, min_value=1)
-early_stopping = st.sidebar.checkbox("Early stopping", value=True)
-no_repeat_ngram_size= st.sidebar.number_input("no repeat ngrams size", value=2, min_value=1)
+no_repeat_ngram_size = st.sidebar.number_input("no repeat ngrams size", value=2, min_value=1)
 num_return_sequences = st.sidebar.number_input("Generated sequences", value=3, min_value=1)
 
+gen_mode = st.sidebar.selectbox("Generation mode", ["beam search", "sampling"])
+if gen_mode == "beam_search": 
+    num_beams = st.sidebar.number_input("Beam size", value=5, min_value=1)
+    early_stopping = st.sidebar.checkbox("Early stopping", value=True)
+    gen_params = {
+        "num_beams": num_beams,
+        "early_stopping": early_stopping
+    }
+else:
+    do_sample = True
+    top_k = st.sidebar.number_input("topk", value=30, min_value=0)
+    top_p = st.sidebar.number_input("topk", value=0, min_value=0)  
+    temperature = st.sidebar.number_input("topk", value=0.7, min_value=0)
+    gen_params = {
+        "do_sample": do_sample,
+        "top_k": top_k,
+        "top_p": top_p,
+        "temperature": temperature
+    }
 
 def generate_caption(url):
     image = Image.open(requests.get(url, stream=True).raw).convert("RGB")
@@ -28,10 +45,9 @@ def generate_caption(url):
     generated_ids = model.generate(
         inputs["pixel_values"],
         max_length=20,
-        num_beams=5,
-        early_stopping=True,
         no_repeat_ngram_size=2,
         num_return_sequences=3,
+        **gen_params
     )
     captions = tokenizer.batch_decode(
         generated_ids,