Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| from __future__ import annotations | |
| import logging | |
| import os | |
| from pathlib import Path | |
| import gradio as gr | |
| from document_processing import DocumentError, prepare_document | |
| from parser_client import ( | |
| ParserClient, | |
| ParserServiceError, | |
| api_key_name_for_environment, | |
| ) | |
| from rate_limit import RateLimiter | |
| from ui_helpers import ( | |
| EMPTY_TABLES, | |
| EMPTY_TEXT, | |
| configuration_required_html, | |
| download_html, | |
| file_strip_html, | |
| header_html, | |
| output_progress_html, | |
| output_values, | |
| page_context_html, | |
| page_limit_html, | |
| privacy_html, | |
| progress_html, | |
| settings_html, | |
| status_html, | |
| upload_intro_html, | |
| viewer_html, | |
| zoom_label_html, | |
| ) | |
| logging.basicConfig( | |
| level=os.environ.get("LOG_LEVEL", "INFO"), | |
| format="%(asctime)s %(levelname)s %(message)s", | |
| ) | |
| LOGGER = logging.getLogger("cohere-parse") | |
| RATE_LIMITER = RateLimiter(limit=5, window_seconds=10 * 60) | |
| ASSETS_DIR = Path(__file__).parent / "assets" | |
| DEFAULT_ZOOM = 50 | |
| gr.set_static_paths(paths=[ASSETS_DIR]) | |
| def _request_key(request: gr.Request | None) -> str: | |
| if request: | |
| forwarded = request.headers.get("x-forwarded-for", "") | |
| if forwarded: | |
| return forwarded.split(",", 1)[0].strip() | |
| real_ip = request.headers.get("x-real-ip", "").strip() | |
| if real_ip: | |
| return real_ip | |
| if request.client and request.client.host: | |
| return request.client.host | |
| return "unknown" | |
| def _view_values( | |
| state: dict, | |
| show_boxes: bool, | |
| render_tables: bool, | |
| zoom: int, | |
| ) -> tuple: | |
| text, tables, raw = output_values(state, render_tables=render_tables) | |
| page_warning = page_limit_html(state) | |
| return ( | |
| viewer_html(state, show_boxes=show_boxes, zoom=zoom), | |
| file_strip_html(state), | |
| gr.HTML(value=page_warning, visible=bool(page_warning)), | |
| gr.Textbox(value=str(state["current_page"] + 1)), | |
| page_context_html(state), | |
| zoom_label_html(zoom), | |
| text, | |
| tables, | |
| raw, | |
| download_html(state), | |
| ) | |
| def prepare_upload( | |
| file_path: str | None, | |
| progress=gr.Progress(), | |
| ) -> tuple: | |
| if not file_path: | |
| raise gr.Error("Choose a PDF or image to continue.") | |
| progress(0, desc="Preparing document") | |
| try: | |
| document = prepare_document( | |
| file_path, | |
| delete_source=True, | |
| enforce_temp_location=True, | |
| on_progress=lambda current, total: progress( | |
| (current, total), | |
| desc=f"Rendering page {current} of {total}", | |
| ), | |
| ) | |
| except DocumentError as exc: | |
| raise gr.Error(str(exc)) from exc | |
| state = document.to_session() | |
| view = _view_values(state, False, True, DEFAULT_ZOOM) | |
| progress(1, desc="Document ready") | |
| LOGGER.info( | |
| "document prepared kind=%s bytes=%s total_pages=%s processed_pages=%s", | |
| state["kind"], | |
| state["size"], | |
| state["total_pages"], | |
| len(state["pages"]), | |
| ) | |
| return ( | |
| state, | |
| gr.Column(visible=False), | |
| gr.Column(visible=True), | |
| *view, | |
| settings_html(state), | |
| gr.HTML(value="", visible=False), | |
| gr.Button(interactive=True), | |
| gr.HTML(value="", visible=False), | |
| ) | |
| def begin_upload(file_path: str | None): | |
| if not file_path: | |
| return gr.HTML(value="", visible=False) | |
| return gr.HTML( | |
| value=progress_html("Preparing document", upload=True), | |
| visible=True, | |
| ) | |
| def clear_upload_progress(): | |
| return gr.HTML(value="", visible=False) | |
| def begin_run(state: dict | None) -> tuple: | |
| if not state or not state.get("pages"): | |
| raise gr.Error("Upload a document first.") | |
| next_page = next( | |
| ( | |
| index + 1 | |
| for index, result in enumerate(state["results"]) | |
| if result is None | |
| ), | |
| state["current_page"] + 1, | |
| ) | |
| return ( | |
| gr.HTML(value="", visible=False), | |
| gr.Button(interactive=False), | |
| output_progress_html(next_page, len(state["pages"])), | |
| gr.Tabs(visible=True), | |
| ) | |
| def clear_run_progress() -> tuple: | |
| return ( | |
| gr.HTML(value="", visible=False), | |
| gr.Button(interactive=True), | |
| EMPTY_TEXT, | |
| gr.Tabs(visible=False), | |
| ) | |
| def run_document( | |
| state: dict | None, | |
| show_boxes: bool, | |
| render_tables: bool, | |
| zoom: int, | |
| request: gr.Request, | |
| progress=gr.Progress(), | |
| ): | |
| if not state or not state.get("pages"): | |
| raise gr.Error("Upload a document first.") | |
| allowed, retry_after = RATE_LIMITER.take(_request_key(request)) | |
| if not allowed: | |
| raise gr.Error( | |
| f"Too many documents have been submitted. Try again in {retry_after} seconds." | |
| ) | |
| try: | |
| client = ParserClient() | |
| except ParserServiceError as exc: | |
| raise gr.Error(str(exc)) from exc | |
| page_count = len(state["pages"]) | |
| parsed_any = False | |
| for index, page in enumerate(state["pages"]): | |
| if state["results"][index] is not None: | |
| continue | |
| parsed_any = True | |
| state["current_page"] = index | |
| progress( | |
| (index, page_count), | |
| desc=f"Parsing page {index + 1} of {page_count}", | |
| ) | |
| try: | |
| result = client.parse_page(page["png"]) | |
| except ParserServiceError as exc: | |
| LOGGER.warning( | |
| "parse failed page=%s completed=%s error=%s", | |
| index + 1, | |
| sum(item is not None for item in state["results"]), | |
| exc, | |
| ) | |
| view = _view_values(state, show_boxes, render_tables, zoom) | |
| yield ( | |
| state, | |
| *view, | |
| gr.HTML( | |
| value=status_html(state, error=str(exc)), | |
| visible=True, | |
| ), | |
| gr.Button(interactive=True), | |
| ) | |
| return | |
| state["results"][index] = result | |
| LOGGER.info( | |
| "page parsed page=%s elapsed_ms=%s chars=%s boxes=%s request_id=%s", | |
| index + 1, | |
| result["elapsed_ms"], | |
| result["output_chars"], | |
| len(result["boxes"]), | |
| result["request_id"], | |
| ) | |
| view = _view_values(state, show_boxes, render_tables, zoom) | |
| yield ( | |
| state, | |
| *view, | |
| gr.HTML( | |
| value=status_html( | |
| state, | |
| running=index + 1 < page_count, | |
| running_page=index + 2, | |
| ), | |
| visible=True, | |
| ), | |
| gr.Button(interactive=False), | |
| ) | |
| if not parsed_any and all(state["results"]): | |
| view = _view_values(state, show_boxes, render_tables, zoom) | |
| yield ( | |
| state, | |
| *view, | |
| gr.HTML(value=status_html(state), visible=True), | |
| gr.Button(interactive=False), | |
| ) | |
| progress(1, desc="Document parsed") | |
| def select_page( | |
| state: dict | None, | |
| page_value: str, | |
| show_boxes: bool, | |
| render_tables: bool, | |
| zoom: int, | |
| ) -> tuple: | |
| if not state: | |
| return gr.skip() | |
| try: | |
| requested_page = int(page_value) | |
| except (TypeError, ValueError): | |
| requested_page = state["current_page"] + 1 | |
| index = max(0, min(len(state["pages"]) - 1, requested_page - 1)) | |
| state["current_page"] = index | |
| view = _view_values(state, show_boxes, render_tables, zoom) | |
| return state, *view | |
| def move_page( | |
| state: dict | None, | |
| delta: int, | |
| show_boxes: bool, | |
| render_tables: bool, | |
| zoom: int, | |
| ) -> tuple: | |
| if not state: | |
| return gr.skip() | |
| index = max( | |
| 0, | |
| min(len(state["pages"]) - 1, state["current_page"] + delta), | |
| ) | |
| state["current_page"] = index | |
| view = _view_values(state, show_boxes, render_tables, zoom) | |
| return state, *view | |
| def refresh_viewer(state: dict | None, show_boxes: bool, zoom: int): | |
| if not state: | |
| return gr.skip() | |
| return viewer_html(state, show_boxes=show_boxes, zoom=zoom) | |
| def refresh_tables(state: dict | None, render_tables: bool): | |
| if not state: | |
| return EMPTY_TABLES | |
| return output_values(state, render_tables=render_tables)[1] | |
| def change_zoom( | |
| state: dict | None, | |
| zoom: int, | |
| delta: int, | |
| show_boxes: bool, | |
| ) -> tuple: | |
| next_zoom = max(10, min(100, zoom + delta)) | |
| if not state: | |
| return next_zoom, gr.skip(), zoom_label_html(next_zoom) | |
| return ( | |
| next_zoom, | |
| viewer_html(state, show_boxes=show_boxes, zoom=next_zoom), | |
| zoom_label_html(next_zoom), | |
| ) | |
| def reset_workspace() -> tuple: | |
| return ( | |
| None, | |
| None, | |
| gr.Column(visible=True), | |
| gr.Column(visible=False), | |
| "", | |
| "", | |
| gr.HTML(value="", visible=False), | |
| EMPTY_TEXT, | |
| EMPTY_TABLES, | |
| "", | |
| "", | |
| gr.HTML(value="", visible=False), | |
| gr.Textbox(value=""), | |
| "", | |
| zoom_label_html(DEFAULT_ZOOM), | |
| gr.Button(interactive=True), | |
| False, | |
| True, | |
| DEFAULT_ZOOM, | |
| gr.Tabs(visible=False), | |
| ) | |
| def build_demo() -> gr.Blocks: | |
| with gr.Blocks( | |
| title="Cohere Parse", | |
| fill_height=True, | |
| fill_width=True, | |
| delete_cache=(60, 600), | |
| analytics_enabled=False, | |
| ) as demo: | |
| session = gr.State(value=None, time_to_live=15 * 60) | |
| zoom_level = gr.State(value=DEFAULT_ZOOM) | |
| gr.HTML(header_html(), elem_id="header-shell") | |
| try: | |
| required_api_key = api_key_name_for_environment() | |
| except ParserServiceError: | |
| required_api_key = None | |
| if not required_api_key or not os.environ.get(required_api_key, "").strip(): | |
| gr.HTML(configuration_required_html(required_api_key)) | |
| with gr.Column(elem_id="upload-screen") as upload_panel: | |
| with gr.Column(elem_id="upload-content"): | |
| gr.HTML(upload_intro_html()) | |
| upload = gr.File( | |
| label="Upload file", | |
| file_types=[".pdf", ".png", ".jpg", ".jpeg", ".webp"], | |
| file_count="single", | |
| type="filepath", | |
| elem_id="document-upload", | |
| ) | |
| upload_progress = gr.HTML( | |
| value="", | |
| visible=False, | |
| elem_id="upload-progress", | |
| ) | |
| gr.HTML(privacy_html(), elem_id="privacy-note-shell") | |
| with gr.Column(visible=False, elem_id="workspace") as workspace: | |
| with gr.Row(equal_height=True, elem_id="workspace-grid"): | |
| with gr.Column(scale=3, min_width=420, elem_id="document-pane"): | |
| file_strip = gr.HTML(elem_id="file-strip-shell") | |
| page_warning = gr.HTML(visible=False, elem_id="page-warning") | |
| with gr.Row(elem_id="viewer-toolbar"): | |
| previous = gr.Button( | |
| "←", | |
| size="sm", | |
| elem_classes=[ | |
| "icon-button", | |
| "page-arrow", | |
| "previous-arrow", | |
| ], | |
| ) | |
| page_select = gr.Textbox( | |
| value="", | |
| show_label=False, | |
| container=False, | |
| max_lines=1, | |
| elem_classes="page-select", | |
| ) | |
| page_total = gr.HTML("", elem_classes="page-total") | |
| next_page = gr.Button( | |
| "→", | |
| size="sm", | |
| elem_classes=[ | |
| "icon-button", | |
| "page-arrow", | |
| "next-arrow", | |
| ], | |
| ) | |
| gr.HTML( | |
| '<span class="toolbar-spacer"></span>', | |
| elem_classes="toolbar-spacer-shell", | |
| ) | |
| zoom_out = gr.Button( | |
| "Zoom out", | |
| size="sm", | |
| elem_classes=["icon-button", "zoom-button", "zoom-out"], | |
| ) | |
| zoom_value = gr.HTML( | |
| zoom_label_html(DEFAULT_ZOOM), | |
| elem_classes="zoom-value-shell", | |
| ) | |
| zoom_in = gr.Button( | |
| "Zoom in", | |
| size="sm", | |
| elem_classes=["icon-button", "zoom-button", "zoom-in"], | |
| ) | |
| viewer = gr.HTML(elem_classes="document-viewer") | |
| with gr.Column(scale=2, min_width=380, elem_id="result-pane"): | |
| settings = gr.HTML() | |
| run_status = gr.HTML(visible=False, elem_id="run-status-shell") | |
| with gr.Column(elem_id="parse-controls"): | |
| show_boxes = gr.Checkbox( | |
| value=False, | |
| label="Bounding boxes", | |
| ) | |
| render_tables = gr.Checkbox( | |
| value=True, | |
| label="Rendered tables", | |
| ) | |
| with gr.Row(elem_id="workspace-actions"): | |
| start_over = gr.Button( | |
| "Start over", | |
| variant="secondary", | |
| elem_classes="start-over-button", | |
| ) | |
| download = gr.HTML(elem_classes="download-container") | |
| run_button = gr.Button( | |
| "Run", | |
| variant="primary", | |
| elem_classes="run-button", | |
| ) | |
| with gr.Tabs(visible=False, elem_id="output-tabs") as output_tabs: | |
| with gr.Tab("Text-only output"): | |
| text_output = gr.HTML( | |
| value=EMPTY_TEXT, | |
| elem_classes="text-output", | |
| ) | |
| with gr.Tab("Tables"): | |
| table_output = gr.HTML( | |
| value=EMPTY_TABLES, | |
| elem_classes="visual-output", | |
| ) | |
| with gr.Tab("Raw"): | |
| raw_output = gr.Code( | |
| value="", | |
| language="markdown", | |
| interactive=False, | |
| show_label=False, | |
| elem_classes="raw-output", | |
| ) | |
| prepare_outputs = [ | |
| session, | |
| upload_panel, | |
| workspace, | |
| viewer, | |
| file_strip, | |
| page_warning, | |
| page_select, | |
| page_total, | |
| zoom_value, | |
| text_output, | |
| table_output, | |
| raw_output, | |
| download, | |
| settings, | |
| run_status, | |
| run_button, | |
| upload_progress, | |
| ] | |
| upload_start = upload.upload( | |
| fn=begin_upload, | |
| inputs=upload, | |
| outputs=upload_progress, | |
| api_visibility="private", | |
| queue=False, | |
| show_progress="hidden", | |
| ) | |
| prepare_event = upload_start.then( | |
| fn=prepare_upload, | |
| inputs=upload, | |
| outputs=prepare_outputs, | |
| api_visibility="private", | |
| concurrency_limit=2, | |
| show_progress="hidden", | |
| ) | |
| prepare_event.failure( | |
| fn=clear_upload_progress, | |
| inputs=None, | |
| outputs=upload_progress, | |
| api_visibility="private", | |
| queue=False, | |
| show_progress="hidden", | |
| ) | |
| view_outputs = [ | |
| session, | |
| viewer, | |
| file_strip, | |
| page_warning, | |
| page_select, | |
| page_total, | |
| zoom_value, | |
| text_output, | |
| table_output, | |
| raw_output, | |
| download, | |
| ] | |
| run_outputs = [ | |
| session, | |
| viewer, | |
| file_strip, | |
| page_warning, | |
| page_select, | |
| page_total, | |
| zoom_value, | |
| text_output, | |
| table_output, | |
| raw_output, | |
| download, | |
| run_status, | |
| run_button, | |
| ] | |
| run_start = run_button.click( | |
| fn=begin_run, | |
| inputs=session, | |
| outputs=[run_status, run_button, text_output, output_tabs], | |
| api_visibility="private", | |
| queue=False, | |
| show_progress="hidden", | |
| ) | |
| run_event = run_start.then( | |
| fn=run_document, | |
| inputs=[session, show_boxes, render_tables, zoom_level], | |
| outputs=run_outputs, | |
| api_visibility="private", | |
| concurrency_limit=2, | |
| concurrency_id="parser-api", | |
| show_progress="hidden", | |
| ) | |
| run_event.failure( | |
| fn=clear_run_progress, | |
| inputs=None, | |
| outputs=[run_status, run_button, text_output, output_tabs], | |
| api_visibility="private", | |
| queue=False, | |
| show_progress="hidden", | |
| ) | |
| page_select.submit( | |
| fn=select_page, | |
| inputs=[session, page_select, show_boxes, render_tables, zoom_level], | |
| outputs=view_outputs, | |
| api_visibility="private", | |
| queue=False, | |
| ) | |
| previous.click( | |
| fn=lambda state, boxes, tables, scale: move_page( | |
| state, -1, boxes, tables, scale | |
| ), | |
| inputs=[session, show_boxes, render_tables, zoom_level], | |
| outputs=view_outputs, | |
| api_visibility="private", | |
| queue=False, | |
| ) | |
| next_page.click( | |
| fn=lambda state, boxes, tables, scale: move_page( | |
| state, 1, boxes, tables, scale | |
| ), | |
| inputs=[session, show_boxes, render_tables, zoom_level], | |
| outputs=view_outputs, | |
| api_visibility="private", | |
| queue=False, | |
| ) | |
| show_boxes.change( | |
| fn=refresh_viewer, | |
| inputs=[session, show_boxes, zoom_level], | |
| outputs=viewer, | |
| api_visibility="private", | |
| queue=False, | |
| ) | |
| zoom_in.click( | |
| fn=lambda state, scale, boxes: change_zoom(state, scale, 10, boxes), | |
| inputs=[session, zoom_level, show_boxes], | |
| outputs=[zoom_level, viewer, zoom_value], | |
| api_visibility="private", | |
| queue=False, | |
| ) | |
| zoom_out.click( | |
| fn=lambda state, scale, boxes: change_zoom(state, scale, -10, boxes), | |
| inputs=[session, zoom_level, show_boxes], | |
| outputs=[zoom_level, viewer, zoom_value], | |
| api_visibility="private", | |
| queue=False, | |
| ) | |
| render_tables.change( | |
| fn=refresh_tables, | |
| inputs=[session, render_tables], | |
| outputs=table_output, | |
| api_visibility="private", | |
| queue=False, | |
| ) | |
| start_over.click( | |
| fn=reset_workspace, | |
| inputs=None, | |
| outputs=[ | |
| session, | |
| upload, | |
| upload_panel, | |
| workspace, | |
| viewer, | |
| file_strip, | |
| page_warning, | |
| text_output, | |
| table_output, | |
| raw_output, | |
| download, | |
| run_status, | |
| page_select, | |
| page_total, | |
| zoom_value, | |
| run_button, | |
| show_boxes, | |
| render_tables, | |
| zoom_level, | |
| output_tabs, | |
| ], | |
| cancels=[run_event], | |
| api_visibility="private", | |
| queue=False, | |
| ) | |
| return demo | |
| demo = build_demo() | |
| demo.queue( | |
| max_size=20, | |
| default_concurrency_limit=2, | |
| api_open=False, | |
| ) | |
| if __name__ == "__main__": | |
| demo.launch( | |
| css_paths="style.css", | |
| max_file_size="20mb", | |
| enable_monitoring=False, | |
| strict_cors=True, | |
| blocked_paths=[ | |
| str(Path(__file__).parent / ".env"), | |
| str(Path(__file__).parent / ".git"), | |
| str(Path(__file__).parent / "__pycache__"), | |
| ], | |
| footer_links=[], | |
| ssr_mode=False, | |
| ) | |