add sample files

tabs/arena_sxs.py

@@ -2,6 +2,11 @@ import gradio as gr
 from utils.pdf_utils import update_page, load_pdf
 from extractors.model_runner import models, run_extract_parallel
 
+sample_files = {
+    "Portfolio Statement": "data/SamplePortfolioStatement.pdf",
+    "Gaming Revenue Report": "data/GamingRevenueReport.pdf",
+}
+
 def update_dropdowns(model_a_choice, model_b_choice):
     if model_a_choice == model_b_choice:
         options_b = [m for m in models if m != model_a_choice]
@@ -11,8 +16,7 @@ def update_dropdowns(model_a_choice, model_b_choice):
         return gr.update(choices=options_b, value=model_b_choice)
 
 def sync_models(model_a_choice, model_b_choice):
-    updated_dropdown_b = update_dropdowns(model_a_choice, model_b_choice)
-    return updated_dropdown_b
+    return update_dropdowns(model_a_choice, model_b_choice)
 
 def dummy_function_a(model_a_choice):
     return f"Model A selected: {model_a_choice}"
@@ -30,20 +34,27 @@ def update_vote_button(output):
 def clear_outputs():
     return "", ""
 
+def load_selected_pdf(file_path):
+    img, page_info, page_num = load_pdf(file_path)
+    return img, page_info, page_num
+
 def arena_sxs():
     with gr.Blocks() as arena_block:
         gr.Markdown("# Rules")
         gr.Markdown("- Upload a PDF file to extract with two chosen models (e.g., Llama, Unstructured, ChatGPT, Claude) and vote for the better one!")
         gr.Markdown("- You can upload multiple files until you identify a winner.")
 
-        gr.Markdown("## 1. Upload a file.")
+        gr.Markdown("## 1. Upload a file or choose a sample.")
         gr.Markdown("Only PDF files supported.")
         with gr.Row():
             with gr.Column(scale=2):
-                pdf = gr.File(type="filepath", label="Upload PDF", file_types=[".pdf"])
+                with gr.Row():
+                    pdf_radio = gr.Radio(choices=list(sample_files.keys()), label="Choose a Sample PDF")
+                    pdf_file = gr.File(type="filepath", label="Upload PDF", file_types=[".pdf"])
                 pdf_image = gr.Image(label="PDF Page")
                 page_info = gr.Textbox(label="")
                 current_page = gr.State(value=0)
+                file_path_state = gr.State(value=None)  # Unified file path state
                 with gr.Row():
                     prev_button = gr.Button("Previous")
                     next_button = gr.Button("Next")
@@ -63,37 +74,55 @@ def arena_sxs():
         with gr.Row():
             extract_button = gr.Button("Parse", interactive=False)
 
-        # with gr.Row():
-        #     vote_model_a_button = gr.Button("👈 A is better", interactive=False)
-        #     vote_model_b_button = gr.Button("👉 B is better", interactive=False)
-        #     vote_tie_button = gr.Button("🤝 Tie", interactive=False)
-        #     vote_bad_button = gr.Button("👎 Both are bad", interactive=False)
-
         # Handlers
-        pdf.change(load_pdf, inputs=[pdf], outputs=[pdf_image, page_info, current_page])
-        pdf.change(fn=update_button, inputs=pdf, outputs=extract_button)
-        model_a.change(sync_models, inputs=[model_a, model_b], outputs=model_b)
-        model_a.change(dummy_function_a, inputs=model_a)
-        model_b.change(dummy_function_b, inputs=model_b)
-        # output_a.change(fn=update_vote_button, inputs=output_a, outputs=[vote_model_a_button, vote_model_b_button, vote_tie_button, vote_bad_button])
-        # output_b.change(fn=update_vote_button, inputs=output_b, outputs=[vote_model_a_button, vote_model_b_button, vote_tie_button, vote_bad_button])
+        def update_pdf_from_source(file_path):
+            img, page_info, page_num = load_pdf(file_path)
+            return img, page_info, page_num
 
-        # # Button Handlers
+        def update_pdf_and_button(file=None, radio=None):
+            file_path = file if file else sample_files.get(radio)
+            if file_path:
+                img, page_info, page_num = update_pdf_from_source(file_path)
+                return img, page_info, 0, file_path, update_button(file_path)  # Include file_path in the outputs
+            return None, "No file selected", 0, None, update_button(None)  # Include file_path in the outputs
+
+        def handle_file_change(file):
+            return update_pdf_and_button(file=file)
+
+        def handle_radio_change(radio):
+            return update_pdf_and_button(radio=radio)
+
+
+        # Handle file uploads
+        pdf_file.change(
+            fn=handle_file_change,
+            inputs=pdf_file,
+            outputs=[pdf_image, page_info, current_page, file_path_state, extract_button]
+        )
+
+        # Handle sample PDF selection from the radio
+        pdf_radio.change(
+            fn=handle_radio_change,
+            inputs=pdf_radio,
+            outputs=[pdf_image, page_info, current_page, file_path_state, extract_button]
+        )
+
+        # Button interactions
         prev_button.click(
             fn=lambda file_path, page: update_page(file_path, page, -1),
-            inputs=[pdf, current_page],
-            outputs=[pdf_image, page_info, current_page],
+            inputs=[file_path_state, current_page],
+            outputs=[pdf_image, page_info, current_page]
         )
 
         next_button.click(
             fn=lambda file_path, page: update_page(file_path, page, 1),
-            inputs=[pdf, current_page],
-            outputs=[pdf_image, page_info, current_page],
+            inputs=[file_path_state, current_page],
+            outputs=[pdf_image, page_info, current_page]
         )
 
         extract_button.click(
             fn=run_extract_parallel,
-            inputs=[model_a, model_b, pdf],
+            inputs=[model_a, model_b, file_path_state],  # Use the unified file path state
             outputs=[output_a, output_b]
         )
 
@@ -102,9 +131,21 @@ def arena_sxs():
             outputs=[output_a, output_b]
         )
 
+    return arena_block
+
+
+        # with gr.Row():
+        #     vote_model_a_button = gr.Button("👈 A is better", interactive=False)
+        #     vote_model_b_button = gr.Button("👉 B is better", interactive=False)
+        #     vote_tie_button = gr.Button("🤝 Tie", interactive=False)
+        #     vote_bad_button = gr.Button("👎 Both are bad", interactive=False)
+        # output_a.change(fn=update_vote_button, inputs=output_a, outputs=[vote_model_a_button, vote_model_b_button, vote_tie_button, vote_bad_button])
+        # output_b.change(fn=update_vote_button, inputs=output_b, outputs=[vote_model_a_button, vote_model_b_button, vote_tie_button, vote_bad_button])
+
+        # # Button Handlers
+
         # vote_model_a_button.click(fn=lambda model_a, model_b: vote_for_model(model_a, Vote.GOOD, model_b, Vote.NEUTRAL), inputs=[model_a, model_b], outputs=[output_a, output_b])
         # vote_model_b_button.click(fn=lambda model_a, model_b: vote_for_model(model_a, Vote.NEUTRAL, model_b, Vote.GOOD), inputs=[model_a, model_b], outputs=[output_a, output_b])
         # vote_tie_button.click(fn=lambda model_a, model_b: vote_for_model(model_a, Vote.NEUTRAL, model_b, Vote.NEUTRAL), inputs=[model_a, model_b], outputs=[output_a, output_b])
         # vote_bad_button.click(fn=lambda model_a, model_b: vote_for_model(model_a, Vote.BAD, model_b, Vote.BAD), inputs=[model_a, model_b], outputs=[output_a, output_b])
 
-    return arena_block

utils/pdf_utils.py

@@ -1,22 +1,43 @@
-import fitz
+import fitz  # PyMuPDF
 from PIL import Image
 
-def update_page(file_path, page_num, direction):
-    if page_num is None:
-        page_num = 0
-    new_page_num = page_num + direction
-    img, actual_page_num, total_pages = get_pdf_page(file_path, new_page_num)
-    return img, f"Page {actual_page_num + 1} of {total_pages}", actual_page_num
+def update_page(file_path, current_page, direction):
+    if not file_path:
+        return None, "No file selected", current_page  # If no file is selected, do nothing
+
+    # Calculate new page index
+    new_page = current_page + direction
+
+    # Load the initial page to get the total pages
+    _, _, total_pages = load_pdf(file_path)
+
+    # Ensure the new page is within bounds
+    if new_page < 0:
+        new_page = 0
+    elif new_page >= total_pages:
+        new_page = total_pages - 1
+
+    # Load the updated page
+    img, page_info, _ = get_pdf_page(file_path, new_page)
+
+    return img, f"Page {new_page + 1} of {total_pages}", new_page
+
 
 def get_pdf_page(file_path, page_num):
-    doc = fitz.open(file_path)
-    page_count = len(doc)
-    page_num = max(0, min(page_num, page_count - 1))  # Ensure page_num is within bounds
-    page = doc.load_page(page_num)
-    pix = page.get_pixmap()
-    img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
-    return img, page_num, page_count
+    try:
+        doc = fitz.open(file_path)
+        page_count = len(doc)
+        page_num = max(0, min(page_num, page_count - 1))  # Ensure page_num is within bounds
+        page = doc.load_page(page_num)
+        pix = page.get_pixmap()
+        img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+        return img, page_num, page_count
+    except Exception as e:
+        print(f"Error loading PDF page: {e}")
+        return None, page_num, 0
+
 
 def load_pdf(file_path):
     img, page_num, total_pages = get_pdf_page(file_path, 0)
-    return img, f"Page {page_num + 1} of {total_pages}", page_num
+    return img, f"Page {page_num + 1} of {total_pages}", total_pages
+