Improve structured extraction alpha

README.md

@@ -86,7 +86,7 @@ First, some configuration:
 I've included a streamlit app that lets you interactively try marker with some basic options.  Run it with:
 
 ```shell
-pip install streamlit
+pip install streamlit streamlit-ace
 marker_gui
 ```
 
@@ -249,6 +249,32 @@ You can also run this via the CLI with
 marker_single FILENAME --converter_cls marker.converters.ocr.OCRConverter
 ```
 
+### Structured Extraction (alpha)
+
+You can run structured extraction via the `ExtractionConverter`.  This requires an llm service to be setup first (see [here](#llm-services) for details).  You'll get a JSON output with the extracted values.
+
+```python
+from marker.converters.extraction import ExtractionConverter
+from marker.models import create_model_dict
+from marker.config.parser import ConfigParser
+from pydantic import BaseModel
+
+class Links(BaseModel):
+    links: list[str]
+    
+schema = Links.model_json_schema()
+config_parser = ConfigParser({
+    "page_schema": schema
+})
+
+converter = ExtractionConverter(
+    artifact_dict=create_model_dict(),
+    config=config_parser.generate_config_dict(),
+    llm_service=config_parser.get_llm_service(),
+)
+rendered = converter("FILEPATH")
+```
+
 # Output Formats
 
 ## Markdown

marker/converters/extraction.py

@@ -1,4 +1,3 @@
-import json
 import re
 
 from marker.builders.document import DocumentBuilder
@@ -6,12 +5,16 @@ from marker.builders.line import LineBuilder
 from marker.builders.ocr import OcrBuilder
 from marker.builders.structure import StructureBuilder
 from marker.converters.pdf import PdfConverter
-from marker.extractors.page import PageExtractor
+from marker.extractors.page import PageExtractor, json_schema_to_base_model
 from marker.providers.registry import provider_from_filepath
 
-from marker.renderers.extraction import ExtractionOutput
+from marker.renderers.extraction import ExtractionMerger
 from marker.renderers.markdown import MarkdownRenderer
 
+from marker.logger import get_logger
+
+logger = get_logger()
+
 
 class ExtractionConverter(PdfConverter):
     pattern: str = r"{\d+\}-{48}\n\n"
@@ -33,11 +36,19 @@ class ExtractionConverter(PdfConverter):
 
         return document, provider
 
-    def __call__(self, filepath: str):
+    def __call__(self, filepath: str) -> str:
         self.config["paginate_output"] = True  # Ensure we can split the output properly
         self.config["output_format"] = (
             "markdown"  # Output must be markdown for extraction
         )
+        try:
+            json_schema_to_base_model(self.config["page_schema"])
+        except Exception as e:
+            logger.error(f"Could not parse page schema: {e}")
+            raise ValueError(
+                "Could not parse your page schema. Please check the schema format."
+            )
+
         document, provider = self.build_document(filepath)
         renderer = self.resolve_dependencies(MarkdownRenderer)
         output = renderer(document)
@@ -53,10 +64,13 @@ class ExtractionConverter(PdfConverter):
             )
 
         extractor = self.resolve_dependencies(PageExtractor)
+        merger = self.resolve_dependencies(ExtractionMerger)
+
+        pnums = provider.page_range
+        all_json = {}
+        for page, page_md, pnum in zip(document.pages, output_pages, pnums):
+            extracted_json = extractor(document, page, page_md.strip())
+            all_json[pnum] = extracted_json
 
-        all_json = []
-        for page, page_md in zip(document.pages, output_pages):
-            extracted_model = extractor(document, page, page_md.strip())
-            extracted_json = extracted_model.model_dump_json()
-            all_json.append(extracted_json)
-        return ExtractionOutput(json=json.dumps(all_json, indent=4, ensure_ascii=False))
+        merged = merger(all_json)
+        return merged

marker/extractors/__init__.py

@@ -1,4 +1,7 @@
-from typing import Annotated, Sequence
+from typing import Annotated, Sequence, Optional
+
+from pydantic import BaseModel
+
 from marker.schema import BlockTypes
 from marker.schema.document import Document
 from marker.schema.groups import PageGroup
@@ -8,6 +11,12 @@ from marker.services import BaseService
 from marker.util import assign_config
 
 
+class ExtractionResult(BaseModel):
+    extracted_data: dict | list
+    value_confidence: int
+    existence_confidence: int
+
+
 class BaseExtractor:
     """
     An extractor that uses a provided service to extract structured data from documents.
@@ -38,5 +47,7 @@ class BaseExtractor:
             remove_blocks=remove_blocks,
         )
 
-    def __call__(self, document: Document, *args, **kwargs):
+    def __call__(
+        self, document: Document, *args, **kwargs
+    ) -> Optional[ExtractionResult]:
         raise NotImplementedError

marker/extractors/page.py

@@ -4,7 +4,7 @@ from pydantic import create_model, BaseModel, Field, ValidationError
 from typing import Annotated, Type, Optional, Any, Dict
 from enum import Enum
 
-from marker.extractors import BaseExtractor
+from marker.extractors import BaseExtractor, ExtractionResult
 from marker.schema.document import Document
 from marker.schema.groups.page import PageGroup
 
@@ -118,7 +118,8 @@ Some guidelines:
 3. Analyze the JSON schema.
 4. Write a short description of the fields in the schema, and the associated values in the image.
 5. Extract the data in the schema that can be found in the image and output the data in JSON format.
-6. Output a confidence score from 1 to 5, where 1 is very low confidence in your extracted values, and 5 is very high confidence in your extracted values.
+6. Output an existence confidence score 1 to 5, where 1 is very low confidence that the values exist on the page, and 5 is very high confidence that the values exist on the page.
+7. Output a value confidence score from 1 to 5, where 1 is very low confidence that the values are correct, and 5 is very high confidence that the values are correct.
 
 **Example:**
 Input:
@@ -158,7 +159,8 @@ Description: The schema has a list of cars, each with a make, sales, and color.
 }
 ```
 
-Confidence: 5
+Existence confidence: 5
+Value confidence: 5
 
 **Input:**
 
@@ -175,7 +177,7 @@ Schema
 
     def __call__(
         self, document: Document, page: PageGroup, page_markdown: str, **kwargs
-    ) -> Optional[BaseModel]:
+    ) -> Optional[ExtractionResult]:
         page_image = self.extract_image(document, page)
         if not self.page_schema:
             raise ValueError(
@@ -189,23 +191,37 @@ Schema
         ).replace("{schema}", json.dumps(optional_schema))
         response = self.llm_service(prompt, page_image, page, PageExtractionSchema)
 
-        if not response or "extracted_json" not in response:
+        if not response or any(
+            [
+                key not in response
+                for key in [
+                    "extracted_json",
+                    "existence_confidence",
+                    "value_confidence",
+                ]
+            ]
+        ):
             page.update_metadata(llm_error_count=1)
-            return
+            return None
 
         extracted_json = response["extracted_json"]
 
         OptionalPageModel = json_schema_to_base_model(optional_schema)
         try:
-            parsed_json = OptionalPageModel.model_validate_json(extracted_json)
+            OptionalPageModel.model_validate_json(extracted_json)
         except ValidationError as e:
             print(f"Validation error with extracted data: {e}")
-            return
+            return None
 
-        return parsed_json
+        return ExtractionResult(
+            extracted_data=json.loads(extracted_json),
+            existence_confidence=response["existence_confidence"],
+            value_confidence=response["value_confidence"],
+        )
 
 
 class PageExtractionSchema(BaseModel):
     description: str
     extracted_json: str
-    confidence: int
+    existence_confidence: int
+    value_confidence: int

marker/renderers/extraction.py

@@ -1,5 +1,63 @@
+from dataclasses import dataclass
+from typing import Dict
+
 from pydantic import BaseModel
 
+from marker.extractors import ExtractionResult
+
+
+@dataclass
+class MergeData:
+    confidence_exists_1: float
+    confidence_exists_2: float
+    confidence_value_1: float
+    confidence_value_2: float
+
+
+def merge_keys(json: dict | list, json2: dict, merge_data: MergeData):
+    if isinstance(json, list):
+        json.extend(json2)
+
+    elif isinstance(json, dict):
+        for key in json:
+            if isinstance(json[key], dict):
+                merge_keys(json[key], json2[key], merge_data)
+            elif isinstance(json[key], list):
+                json[key] = json[key] + json2[key]
+            else:
+                if (
+                    merge_data.confidence_exists_2 > 3
+                    and merge_data.confidence_value_2 > 3
+                    and json2[key]
+                ):
+                    json[key] = json2[key]
+
+                if not json[key] and json2[key]:
+                    json[key] = json2[key]
+
 
 class ExtractionOutput(BaseModel):
-    json: str
+    pages: Dict[int, ExtractionResult]
+    json: dict
+
+
+class ExtractionMerger:
+    def __init__(self):
+        pass
+
+    def __call__(self, outputs: Dict[int, ExtractionResult]):
+        pnums = sorted(list(outputs.keys()))
+        merged_result = outputs[pnums[0]].extracted_data.copy()
+        confidence_exists = outputs[pnums[0]].existence_confidence
+        confidence_value = outputs[pnums[0]].value_confidence
+
+        for pnum in pnums[1:]:
+            merge_data = MergeData(
+                confidence_exists_1=confidence_exists,
+                confidence_exists_2=outputs[pnum].existence_confidence,
+                confidence_value_1=confidence_value,
+                confidence_value_2=outputs[pnum].value_confidence,
+            )
+            merge_keys(merged_result, outputs[pnum].extracted_data, merge_data)
+
+        return ExtractionOutput(pages=outputs, json=merged_result)

marker/scripts/common.py

@@ -165,8 +165,10 @@ def get_root_class(schema_code: str) -> Optional[BaseModel]:
         return None
 
     if "from pydantic" not in schema_code:
+        schema_code = "from pydantic import BaseModel\n" + schema_code
+    if "from typing" not in schema_code:
         schema_code = (
-            "from pydantic import BaseModel\nfrom typing import List, Dict, Optional, Set, Tuple, Union, Any\n\n"
+            "from typing import List, Dict, Optional, Set, Tuple, Union, Any\n\n"
             + schema_code
         )
 

marker/scripts/extraction_app.py

@@ -1,6 +1,7 @@
 import json
 import os
 
+from streamlit_ace import st_ace
 from pydantic import BaseModel
 
 from marker.converters.extraction import ExtractionConverter
@@ -23,7 +24,6 @@ from typing import Any, Dict
 import streamlit as st
 
 from marker.config.parser import ConfigParser
-from marker.output import text_from_rendered
 
 
 def extract_data(fname: str, config: dict, schema: str) -> (str, Dict[str, Any], dict):
@@ -55,9 +55,9 @@ cli_options = parse_args()
 st.markdown("""
 # Marker Extraction Demo
 
-This app will let you try marker, a PDF or image -> Markdown, HTML, JSON converter. It works with any language, and extracts images, tables, equations, etc.
+This app will let you use marker to do structured extraction.
 
-Find the project [here](https://github.com/VikParuchuri/marker).
+Warning: This can execute untrusted code entered into the schema panel.
 """)
 
 in_file: UploadedFile = st.sidebar.file_uploader(
@@ -79,18 +79,18 @@ with col1:
     st.image(pil_image, use_container_width=True)
 
 with col2:
-    schema = st.text_area(
-        "Pyantic schema for extraction",
-        value="""
+    st.write("Enter pydantic schema here")
+    schema = st_ace(
+        value="""from pydantic import BaseModel
 class Schema(BaseModel):
-    pass        
-""",
+    pass""",
+        language="python",
     )
 
 run_marker = st.sidebar.button("Run Extraction")
 
 use_llm = st.sidebar.checkbox(
-    "Use LLM", help="Use LLM for higher quality processing", value=False
+    "Use LLM", help="Use LLM for higher quality text", value=False
 )
 force_ocr = st.sidebar.checkbox("Force OCR", help="Force OCR on all pages", value=False)
 strip_existing_ocr = st.sidebar.checkbox(
@@ -123,8 +123,6 @@ with tempfile.TemporaryDirectory() as tmp_dir:
     )
     rendered = extract_data(temp_pdf, cli_options, schema)
 
-text, ext, images = text_from_rendered(rendered)
-
 with col2:
     st.write("Output JSON")
-    st.json(text)
+    st.json(rendered.model_dump())