From ec57b528ed3f8e7b9604ff4cdf06da3d52c714dd Mon Sep 17 00:00:00 2001
From: Tiago Santana <54704492+SantanaTiago@users.noreply.github.com>
Date: Thu, 13 Mar 2025 16:09:59 +0000
Subject: [PATCH] feat: expose options for new features (#92)

Signed-off-by: Tiago Santana <54704492+SantanaTiago@users.noreply.github.com>
---
 docling_serve/datamodel/convert.py  | 44 +++++++++++++++++++++++++++++
 docling_serve/docling_conversion.py |  4 +++
 docling_serve/gradio_ui.py          | 39 +++++++++++++++++++++++++
 3 files changed, 87 insertions(+)

diff --git a/docling_serve/datamodel/convert.py b/docling_serve/datamodel/convert.py
index f299815..087b160 100644
--- a/docling_serve/datamodel/convert.py
+++ b/docling_serve/datamodel/convert.py
@@ -172,3 +172,47 @@ class ConvertDocumentsOptions(BaseModel):
             examples=[2.0],
         ),
     ] = 2.0
+
+    do_code_enrichment: Annotated[
+        bool,
+        Field(
+            description=(
+                "If enabled, perform OCR code enrichment. "
+                "Boolean. Optional, defaults to false."
+            ),
+            examples=[False],
+        ),
+    ] = False
+
+    do_formula_enrichment: Annotated[
+        bool,
+        Field(
+            description=(
+                "If enabled, perform formula OCR, return Latex code. "
+                "Boolean. Optional, defaults to false."
+            ),
+            examples=[False],
+        ),
+    ] = False
+
+    do_picture_classification: Annotated[
+        bool,
+        Field(
+            description=(
+                "If enabled, classify pictures in documents. "
+                "Boolean. Optional, defaults to false."
+            ),
+            examples=[False],
+        ),
+    ] = False
+
+    do_picture_description: Annotated[
+        bool,
+        Field(
+            description=(
+                "If enabled, describe pictures in documents. "
+                "Boolean. Optional, defaults to false."
+            ),
+            examples=[False],
+        ),
+    ] = False
diff --git a/docling_serve/docling_conversion.py b/docling_serve/docling_conversion.py
index 2d41f5a..6faaf19 100644
--- a/docling_serve/docling_conversion.py
+++ b/docling_serve/docling_conversion.py
@@ -120,6 +120,10 @@ def get_pdf_pipeline_opts(  # noqa: C901
         do_ocr=request.do_ocr,
         ocr_options=ocr_options,
         do_table_structure=request.do_table_structure,
+        do_code_enrichment=request.do_code_enrichment,
+        do_formula_enrichment=request.do_formula_enrichment,
+        do_picture_classification=request.do_picture_classification,
+        do_picture_description=request.do_picture_description,
     )
     pipeline_options.table_structure_options.do_cell_matching = True  # do_cell_matching
     pipeline_options.table_structure_options.mode = TableFormerMode(request.table_mode)
diff --git a/docling_serve/gradio_ui.py b/docling_serve/gradio_ui.py
index 9c328b7..4b6f1ec 100644
--- a/docling_serve/gradio_ui.py
+++ b/docling_serve/gradio_ui.py
@@ -185,6 +185,10 @@ def process_url(
     table_mode,
     abort_on_error,
     return_as_file,
+    do_code_enrichment,
+    do_formula_enrichment,
+    do_picture_classification,
+    do_picture_description,
 ):
     parameters = {
         "http_sources": [{"url": source} for source in input_sources.split(",")],
@@ -199,6 +203,10 @@ def process_url(
             "table_mode": table_mode,
             "abort_on_error": abort_on_error,
             "return_as_file": return_as_file,
+            "do_code_enrichment": do_code_enrichment,
+            "do_formula_enrichment": do_formula_enrichment,
+            "do_picture_classification": do_picture_classification,
+            "do_picture_description": do_picture_description,
         },
     }
     if (
@@ -237,6 +245,10 @@ def process_file(
     table_mode,
     abort_on_error,
     return_as_file,
+    do_code_enrichment,
+    do_formula_enrichment,
+    do_picture_classification,
+    do_picture_description,
 ):
     if not files or len(files) == 0 or files[0] == "":
         logger.error("No files provided.")
@@ -254,6 +266,10 @@ def process_file(
         "table_mode": table_mode,
         "abort_on_error": str(abort_on_error).lower(),
         "return_as_file": str(return_as_file).lower(),
+        "do_code_enrichment": str(do_code_enrichment).lower(),
+        "do_formula_enrichment": str(do_formula_enrichment).lower(),
+        "do_picture_classification": str(do_picture_classification).lower(),
+        "do_picture_description": str(do_picture_description).lower(),
     }
 
     try:
@@ -478,6 +494,21 @@ with gr.Blocks(
             with gr.Column(scale=1):
                 abort_on_error = gr.Checkbox(label="Abort on Error", value=False)
                 return_as_file = gr.Checkbox(label="Return as File", value=False)
+        with gr.Row():
+            with gr.Column():
+                do_code_enrichment = gr.Checkbox(
+                    label="Enable code enrichment", value=False
+                )
+                do_formula_enrichment = gr.Checkbox(
+                    label="Enable formula enrichment", value=False
+                )
+            with gr.Column():
+                do_picture_classification = gr.Checkbox(
+                    label="Enable picture classification", value=False
+                )
+                do_picture_description = gr.Checkbox(
+                    label="Enable picture description", value=False
+                )
 
     # Document output
     with gr.Row(visible=False) as content_output:
@@ -561,6 +592,10 @@ with gr.Blocks(
             table_mode,
             abort_on_error,
             return_as_file,
+            do_code_enrichment,
+            do_formula_enrichment,
+            do_picture_classification,
+            do_picture_description,
         ],
         outputs=[
             output_markdown,
@@ -630,6 +665,10 @@ with gr.Blocks(
             table_mode,
             abort_on_error,
             return_as_file,
+            do_code_enrichment,
+            do_formula_enrichment,
+            do_picture_classification,
+            do_picture_description,
         ],
         outputs=[
             output_markdown,