Update app.py

app.py

@@ -120,11 +120,11 @@ def update_placeholder(role):
 
 
 def reset_components():
-    return ['pla', 'gnd', 'ver', 'ans'], 0, 256
+    return ['pla', 'gnd', 'ver', 'ans'], 5, 0, 256
 
 
 @spaces.GPU
-def main(video, prompt, role, temperature, max_new_tokens):
+def main(video, prompt, role, max_candidates, temperature, max_new_tokens):
     global model, processor, device
 
     history = []
@@ -356,8 +356,7 @@ def main(video, prompt, role, temperature, max_new_tokens):
         response += '\n| ID | Start Time | End Time | Confidence |'
         response += '\n| :-: | :-: | :-: | :-: |'
 
-        # using top-5 predictions
-        for i, (p, c) in enumerate(zip(pred[:5], conf[:5])):
+        for i, (p, c) in enumerate(zip(pred[:max_candidates], conf[:max_candidates])):
             response += f'\n| {i} | {seconds_to_hms(p[0])} | {seconds_to_hms(p[1])} | {c:.2f} |'
 
         response += f'\n\nTherefore, the target moment might happens from <span style="color:red">{seconds_to_hms(pred[0][0])}</span> to <span style="color:red">{seconds_to_hms(pred[0][1])}</span>.'
@@ -381,9 +380,8 @@ def main(video, prompt, role, temperature, max_new_tokens):
 
         start_time = time.perf_counter()
 
-        # using top-5 predictions
         prob = []
-        for i, cand in enumerate(pred[:5]):
+        for i, cand in enumerate(pred[:max_candidates]):
             s0, e0 = parse_span(cand, duration, 2)
             offset = (e0 - s0) / 2
             s1, e1 = parse_span([s0 - offset, e0 + offset], duration)
@@ -593,6 +591,13 @@ def build_demo():
                     role.change(update_placeholder, role, prompt)
 
                     with gr.Accordion(label='Hyperparameters', open=False):
+                        max_candidates = gr.Slider(
+                            1,
+                            100,
+                            value=5,
+                            interactive=True,
+                            label='Max Candidate Moments',
+                            info='The maximum number of candidate moments in Grounder (Default: 5)')
                         temperature = gr.Slider(
                             0,
                             1,
@@ -616,12 +621,12 @@ def build_demo():
                     random_btn.click(random_sample, None, [video, prompt, role])
 
                     reset_btn = gr.ClearButton([video, prompt, chat], value='🗑️ Reset')
-                    reset_btn.click(reset_components, None, [role, temperature, max_new_tokens])
+                    reset_btn.click(reset_components, None, [role, max_candidates, temperature, max_new_tokens])
 
                     submit_btn = gr.Button(value='🚀 Submit', variant='primary')
-                    submit_ctx = submit_btn.click(disable_btns, None, [random_btn, reset_btn, submit_btn])
-                    submit_ctx = submit_ctx.then(main, [video, prompt, role, temperature, max_new_tokens], chat)
-                    submit_ctx.then(enable_btns, None, [random_btn, reset_btn, submit_btn])
+                    ctx = submit_btn.click(disable_btns, None, [random_btn, reset_btn, submit_btn])
+                    ctx = ctx.then(main, [video, prompt, role, max_candidates, temperature, max_new_tokens], chat)
+                    ctx.then(enable_btns, None, [random_btn, reset_btn, submit_btn])
 
                 gr.Examples(examples=EXAMPLES, inputs=[video, prompt, role], examples_per_page=3)