Update app.py

app.py

@@ -6,7 +6,6 @@ import io
 import base64
 from datasets import load_dataset
 
-
 max_token_budget = 512
 
 min_pixels = 1 * 28 * 28
@@ -77,7 +76,6 @@ def doc_to_messages(text, slides):
   text = processor.apply_chat_template(
       messages, tokenize=False, add_generation_prompt=True
   )
-  print(text)
   image_inputs, video_inputs = process_vision_info(messages)
   inputs = processor(
       text=[text],
@@ -93,6 +91,8 @@ def doc_to_messages(text, slides):
 current_doc_index = 0
 annotations = []
 
+choices = [f"{i} | {ds['title'][i]}" for i in range(len(ds))]
+
 def load_document(index):
     """Load a specific document from the dataset"""
     if 0 <= index < len(ds):
@@ -103,19 +103,25 @@ def load_document(index):
             doc["abstract"],
             create_interleaved_html(segments_doc, doc["slides"], scale=0.7),
             doc_to_messages(segments_doc, doc["slides"]).input_ids.shape[1],
+            choices[index],
         )
-    return ("", "", "", "")
+    return ("", "", "", "", "")
 
 def get_next_document():
     """Get the next document in the dataset"""
     global current_doc_index
-    current_doc_index = (current_doc_index + 1) % len(ds)
-    return load_document(current_doc_index)
+    return choices[(current_doc_index + 1) % len(ds)]
 
 def get_prev_document():
     """Get the previous document in the dataset"""
     global current_doc_index
-    current_doc_index = (current_doc_index - 1) % len(ds)
+    return choices[(current_doc_index - 1) % len(ds)]
+
+def get_selected_document(arg):
+    """Get the selected document from the dataset"""
+    global current_doc_index
+    index = int(arg.split(" | ")[0])
+    current_doc_index = index
     return load_document(current_doc_index)
 
 
@@ -123,36 +129,34 @@ theme = gr.themes.Ocean()
 
 with gr.Blocks(theme=theme) as demo:
     gr.Markdown("# Slide Presentation Visualization Tool")
+    pres_selection_dd = gr.Dropdown(label="Presentation", value=choices[0], choices=choices)
     with gr.Row():
         with gr.Column():
             body = gr.HTML(max_height=400)
 
-        # Function to update the interleaved view
-        def update_interleaved_view(title, abstract, body, token_count):
-            return body
-
         with gr.Column():
             title = gr.Textbox(label="Title", interactive=False, max_lines=1)
             abstract = gr.Textbox(label="Abstract", interactive=False, max_lines=8)
             token_count = gr.Textbox(label=f"Token Count (Qwen2-VL with under {max_token_budget} tokens per image)", interactive=False, max_lines=1)
 
-    title.change(
-        fn=update_interleaved_view,
-        inputs=[title, abstract, body, token_count],
-        outputs=body,
-        )
     # Load first document
-    title_val, abstract_val, body_val, token_count_val = load_document(current_doc_index)
+    title_val, abstract_val, body_val, token_count_val, choices_val = load_document(current_doc_index)
     title.value = title_val
     abstract.value = abstract_val
     body.value = body_val
     token_count.value = str(token_count_val)
+    pres_selection_dd.value = choices_val
 
+    pres_selection_dd.change(
+        fn=get_selected_document,
+        inputs=pres_selection_dd,
+        outputs=[title, abstract, body, token_count, pres_selection_dd],
+    )
 
     with gr.Row():
       prev_button = gr.Button("Previous Document")
-      prev_button.click(fn=get_prev_document, inputs=[], outputs=[title, abstract, body, token_count])
+      prev_button.click(fn=get_prev_document, inputs=[], outputs=[pres_selection_dd])
       next_button = gr.Button("Next Document")
-      next_button.click(fn=get_next_document, inputs=[], outputs=[title, abstract, body, token_count])
+      next_button.click(fn=get_next_document, inputs=[], outputs=[pres_selection_dd])
 
 demo.launch()