Improve structured extraction

README.md

@@ -1,6 +1,6 @@
 # Marker
 
-Marker converts documents to markdown, JSON, and HTML quickly and accurately.
+Marker converts documents to markdown, JSON, chunks, and HTML quickly and accurately.
 
 - Converts PDF, image, PPTX, DOCX, XLSX, HTML, EPUB files in all languages
 - Formats tables, forms, equations, inline math, links, references, and code blocks
@@ -276,6 +276,8 @@ converter = ExtractionConverter(
 rendered = converter("FILEPATH")
 ```
 
+Rendered will have an `original_markdown` field.  If you pass this back in next time you run the converter, as the `existing_markdown` config key, you can skip re-parsing the document.
+
 # Output Formats
 
 ## Markdown
@@ -348,7 +350,7 @@ Note that child blocks of pages can have their own children as well (a tree stru
 
 ## Chunks
 
-Chunks format is similar to JSON, but flattens everything into a single list instead of a tree.  Only the top level blocks from each page show up. It also has the full HTML of each block inside, so you don't need to crawl the tree to reconstruct it.
+Chunks format is similar to JSON, but flattens everything into a single list instead of a tree.  Only the top level blocks from each page show up. It also has the full HTML of each block inside, so you don't need to crawl the tree to reconstruct it.  This enable flexible and easy chunking for RAG.
 
 ## Metadata
 

marker/converters/extraction.py

@@ -1,5 +1,5 @@
-import json
 import re
+from typing import Annotated
 
 from marker.builders.document import DocumentBuilder
 from marker.builders.line import LineBuilder
@@ -8,7 +8,6 @@ from marker.builders.structure import StructureBuilder
 from marker.converters.pdf import PdfConverter
 from marker.extractors.document import DocumentExtractor
 from marker.extractors.page import PageExtractor
-from marker.extractors.util import json_schema_to_base_model
 from marker.providers.registry import provider_from_filepath
 
 from marker.renderers.extraction import ExtractionRenderer, ExtractionOutput
@@ -21,6 +20,9 @@ logger = get_logger()
 
 class ExtractionConverter(PdfConverter):
     pattern: str = r"{\d+\}-{48}\n\n"
+    existing_markdown: Annotated[
+        str, "Markdown that was already converted for extraction."
+    ] = None
 
     def build_document(self, filepath: str):
         provider_cls = provider_from_filepath(filepath)
@@ -44,23 +46,16 @@ class ExtractionConverter(PdfConverter):
         self.config["output_format"] = (
             "markdown"  # Output must be markdown for extraction
         )
+        markdown = self.existing_markdown
 
-        try:
-            json_schema_to_base_model(json.loads(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)
-        self.page_count = len(document.pages)
-        renderer = self.resolve_dependencies(MarkdownRenderer)
-        output = renderer(document)
+        if not markdown:
+            document, provider = self.build_document(filepath)
+            self.page_count = len(document.pages)
+            renderer = self.resolve_dependencies(MarkdownRenderer)
+            output = renderer(document)
+            markdown = output.markdown
 
-        output_pages = re.split(self.pattern, output.markdown)[
-            1:
-        ]  # Split output into pages
+        output_pages = re.split(self.pattern, markdown)[1:]  # Split output into pages
 
         # This needs an LLM service for extraction, this sets it in the extractor
         if not self.artifact_dict["llm_service"]:
@@ -73,8 +68,8 @@ class ExtractionConverter(PdfConverter):
         renderer = self.resolve_dependencies(ExtractionRenderer)
 
         # Inference in parallel
-        notes = page_extractor(document, document.pages, output_pages)
-        document_output = document_extractor(document, notes)
+        notes = page_extractor(output_pages)
+        document_output = document_extractor(notes)
 
-        merged = renderer(document_output)
+        merged = renderer(document_output, markdown)
         return merged

marker/extractors/document.py

@@ -6,7 +6,6 @@ from typing import Annotated, Optional, List
 from marker.extractors import BaseExtractor
 from marker.extractors.page import PageExtractionSchema
 from marker.logger import get_logger
-from marker.schema.document import Document
 
 logger = get_logger()
 
@@ -114,7 +113,6 @@ Schema
 
     def __call__(
         self,
-        document: Document,
         page_notes: List[PageExtractionSchema],
         **kwargs,
     ) -> Optional[DocumentExtractionSchema]:
@@ -123,11 +121,10 @@ Schema
                 "Page schema must be defined for structured extraction to work."
             )
 
-        page = document.pages[0]
         prompt = self.page_extraction_prompt.replace(
             "{{document_notes}}", self.assemble_document_notes(page_notes)
         ).replace("{{schema}}", json.dumps(self.page_schema))
-        response = self.llm_service(prompt, None, page, DocumentExtractionSchema)
+        response = self.llm_service(prompt, None, None, DocumentExtractionSchema)
 
         logger.debug(f"Document extraction response: {response}")
 
@@ -140,7 +137,6 @@ Schema
                 ]
             ]
         ):
-            page.update_metadata(llm_error_count=1)
             return None
 
         json_data = response["document_json"].strip().lstrip("```json").rstrip("```")

marker/extractors/page.py

@@ -8,8 +8,6 @@ from tqdm import tqdm
 
 from marker.extractors import BaseExtractor
 from marker.logger import get_logger
-from marker.schema.document import Document
-from marker.schema.groups.page import PageGroup
 
 logger = get_logger()
 
@@ -100,29 +98,25 @@ Schema
 ```
 """
 
-    def chunk_page_markdown(
-        self, pages: List[PageGroup], page_markdown: List[str]
-    ) -> List[tuple]:
+    def chunk_page_markdown(self, page_markdown: List[str]) -> List[str]:
         """
         Chunk the page markdown into smaller pieces for processing.
         """
-        if len(pages) == 0:
-            return []
 
         chunks = []
-        for i in range(0, len(pages), self.extraction_page_chunk_size):
+        for i in range(0, len(page_markdown), self.extraction_page_chunk_size):
             chunk = page_markdown[i : i + self.extraction_page_chunk_size]
-            chunks.append((pages[i], "\n\n".join(chunk)))
+            chunks.append("\n\n".join(chunk))
 
         return chunks
 
     def inference_single_chunk(
-        self, page: PageGroup, page_markdown: str
+        self, page_markdown: str
     ) -> Optional[PageExtractionSchema]:
         prompt = self.page_extraction_prompt.replace(
             "{{page_md}}", page_markdown
         ).replace("{{schema}}", json.dumps(self.page_schema))
-        response = self.llm_service(prompt, None, page, PageExtractionSchema)
+        response = self.llm_service(prompt, None, None, PageExtractionSchema)
         logger.debug(f"Page extraction response: {response}")
 
         if not response or any(
@@ -134,7 +128,6 @@ Schema
                 ]
             ]
         ):
-            page.update_metadata(llm_error_count=1)
             return None
 
         return PageExtractionSchema(
@@ -144,20 +137,15 @@ Schema
 
     def __call__(
         self,
-        document: Document,
-        pages: List[PageGroup],
         page_markdown: List[str],
         **kwargs,
     ) -> List[PageExtractionSchema]:
-        assert len(page_markdown) == len(pages), (
-            f"Mismatch in page markdown and pages length: {len(page_markdown)} vs {len(pages)}"
-        )
         if not self.page_schema:
             raise ValueError(
                 "Page schema must be defined for structured extraction to work."
             )
 
-        chunks = self.chunk_page_markdown(pages, page_markdown)
+        chunks = self.chunk_page_markdown(page_markdown)
         results = []
         pbar = tqdm(
             desc="Running page extraction",
@@ -167,8 +155,7 @@ Schema
 
         with ThreadPoolExecutor(max_workers=self.max_concurrency) as executor:
             for future in [
-                executor.submit(self.inference_single_chunk, chunk[0], chunk[1])
-                for chunk in chunks
+                executor.submit(self.inference_single_chunk, chunk) for chunk in chunks
             ]:
                 results.append(future.result())  # Raise exceptions if any occurred
                 pbar.update(1)

marker/extractors/util.py

@@ -1,213 +0,0 @@
-from typing import Any, Type, Union, Optional
-from pydantic import BaseModel, Field, create_model, validator
-from enum import Enum
-import re
-from datetime import datetime
-from uuid import UUID
-
-
-def json_schema_to_base_model(
-    schema: dict[str, Any], model_name: str = None
-) -> Type[BaseModel]:
-    """Convert a JSON schema to a Pydantic BaseModel dynamically."""
-
-    # Enhanced type mapping with format support
-    def get_type_from_schema(field_props: dict[str, Any]) -> type:
-        json_type = field_props.get("type", "string")
-        format_type = field_props.get("format")
-
-        # Handle format-specific types
-        if json_type == "string":
-            if format_type == "date-time":
-                return datetime
-            elif format_type == "uuid":
-                return UUID
-            else:
-                return str
-        elif json_type == "integer":
-            return int
-        elif json_type == "number":
-            return float
-        elif json_type == "boolean":
-            return bool
-        elif json_type == "array":
-            return list
-        elif json_type == "object":
-            return dict
-        else:
-            return str  # fallback
-
-    def handle_union_types(field_props: dict[str, Any]) -> type:
-        """Handle anyOf, oneOf, and type arrays."""
-        any_of = field_props.get("anyOf", [])
-        one_of = field_props.get("oneOf", [])
-        type_list = field_props.get("type", [])
-
-        if any_of:
-            types = [get_type_from_schema(schema) for schema in any_of]
-            return Union[tuple(types)]
-        elif one_of:
-            types = [get_type_from_schema(schema) for schema in one_of]
-            return Union[tuple(types)]
-        elif isinstance(type_list, list):
-            types = [get_type_from_schema({"type": t}) for t in type_list]
-            return Union[tuple(types)]
-
-        return None
-
-    def create_validator_from_constraints(field_name: str, field_props: dict[str, Any]):
-        """Create Pydantic validators from JSON schema constraints."""
-        validators = {}
-
-        # String constraints
-        if "minLength" in field_props:
-            min_len = field_props["minLength"]
-
-            def min_length_validator(cls, v):
-                if isinstance(v, str) and len(v) < min_len:
-                    raise ValueError(
-                        f"{field_name} must be at least {min_len} characters"
-                    )
-                return v
-
-            validators[f"{field_name}_min_length"] = validator(
-                field_name, allow_reuse=True
-            )(min_length_validator)
-
-        if "maxLength" in field_props:
-            max_len = field_props["maxLength"]
-
-            def max_length_validator(cls, v):
-                if isinstance(v, str) and len(v) > max_len:
-                    raise ValueError(
-                        f"{field_name} must be at most {max_len} characters"
-                    )
-                return v
-
-            validators[f"{field_name}_max_length"] = validator(
-                field_name, allow_reuse=True
-            )(max_length_validator)
-
-        if "pattern" in field_props:
-            pattern = field_props["pattern"]
-
-            def pattern_validator(cls, v):
-                if isinstance(v, str) and not re.match(pattern, v):
-                    raise ValueError(f"{field_name} must match pattern {pattern}")
-                return v
-
-            validators[f"{field_name}_pattern"] = validator(
-                field_name, allow_reuse=True
-            )(pattern_validator)
-
-        # Numeric constraints
-        if "minimum" in field_props:
-            min_val = field_props["minimum"]
-
-            def min_validator(cls, v):
-                if isinstance(v, (int, float)) and v < min_val:
-                    raise ValueError(f"{field_name} must be at least {min_val}")
-                return v
-
-            validators[f"{field_name}_minimum"] = validator(
-                field_name, allow_reuse=True
-            )(min_validator)
-
-        if "maximum" in field_props:
-            max_val = field_props["maximum"]
-
-            def max_validator(cls, v):
-                if isinstance(v, (int, float)) and v > max_val:
-                    raise ValueError(f"{field_name} must be at most {max_val}")
-                return v
-
-            validators[f"{field_name}_maximum"] = validator(
-                field_name, allow_reuse=True
-            )(max_validator)
-
-        return validators
-
-    def process_field(field_name: str, field_props: dict[str, Any]) -> tuple:
-        """Process a single field from the schema."""
-
-        # Handle const values
-        if "const" in field_props:
-            const_value = field_props["const"]
-            return type(const_value), Field(default=const_value, const=True)
-
-        # Handle enums
-        enum_values = field_props.get("enum")
-        if enum_values:
-            enum_name = f"{field_name.capitalize()}Enum"
-            field_type = Enum(enum_name, {str(v): v for v in enum_values})
-
-        # Handle union types (anyOf, oneOf, type arrays)
-        elif union_type := handle_union_types(field_props):
-            field_type = union_type
-
-        # Handle nested objects
-        elif field_props.get("type") == "object" and "properties" in field_props:
-            nested_model_name = f"{field_name.capitalize()}Model"
-            field_type = json_schema_to_base_model(field_props, nested_model_name)
-
-        # Handle arrays
-        elif field_props.get("type") == "array" and "items" in field_props:
-            item_props = field_props["items"]
-
-            # Handle array of objects
-            if item_props.get("type") == "object" and "properties" in item_props:
-                item_model_name = f"{field_name.capitalize()}ItemModel"
-                item_type = json_schema_to_base_model(item_props, item_model_name)
-            else:
-                item_type = get_type_from_schema(item_props)
-
-            field_type = list[item_type]
-
-        # Handle primitive types
-        else:
-            field_type = get_type_from_schema(field_props)
-
-        # Handle nullable
-        if field_props.get("nullable", False):
-            field_type = Optional[field_type]
-
-        # Determine default value
-        if "default" in field_props:
-            default_value = field_props["default"]
-        elif field_name not in schema.get("required", []):
-            default_value = None
-            if not field_props.get("nullable", False):
-                field_type = Optional[field_type]
-        else:
-            default_value = ...
-
-        # Create field with metadata
-        field_info = Field(
-            default=default_value,
-            description=field_props.get("description", field_props.get("title", "")),
-            title=field_props.get("title"),
-            examples=field_props.get("examples"),
-        )
-
-        return field_type, field_info
-
-    # Process schema
-    properties = schema.get("properties", {})
-    model_fields = {}
-    validators = {}
-
-    # Process each field
-    for field_name, field_props in properties.items():
-        model_fields[field_name] = process_field(field_name, field_props)
-
-        # Add validators for constraints
-        field_validators = create_validator_from_constraints(field_name, field_props)
-        validators.update(field_validators)
-
-    # Create the model
-    model_name = model_name or schema.get("title", "DynamicModel")
-
-    # Create model with validators
-    model_class = create_model(model_name, **model_fields, __validators__=validators)
-
-    return model_class

marker/renderers/extraction.py

@@ -7,11 +7,16 @@ from marker.renderers import BaseRenderer
 class ExtractionOutput(BaseModel):
     analysis: str
     document_json: str
+    original_markdown: str
 
 
 class ExtractionRenderer(BaseRenderer):
-    def __call__(self, output: DocumentExtractionSchema) -> ExtractionOutput:
+    def __call__(
+        self, output: DocumentExtractionSchema, markdown: str
+    ) -> ExtractionOutput:
         # We definitely want to do more complex stuff here soon, so leave it in
         return ExtractionOutput(
-            analysis=output.analysis, document_json=output.document_json
+            analysis=output.analysis,
+            document_json=output.document_json,
+            original_markdown=markdown,
         )

marker/scripts/extraction_app.py

@@ -26,9 +26,12 @@ import streamlit as st
 from marker.config.parser import ConfigParser
 
 
-def extract_data(fname: str, config: dict, schema: str) -> (str, Dict[str, Any], dict):
+def extract_data(
+    fname: str, config: dict, schema: str, markdown: str | None = None
+) -> (str, Dict[str, Any], dict):
     config["pdftext_workers"] = 1
     config["page_schema"] = schema
+    config["existing_markdown"] = markdown
     config_parser = ConfigParser(config)
     config_dict = config_parser.generate_config_dict()
 
@@ -62,12 +65,35 @@ in_file: UploadedFile = st.sidebar.file_uploader(
     type=["pdf", "png", "jpg", "jpeg", "gif", "pptx", "docx", "xlsx", "html", "epub"],
 )
 
-if in_file is None:
-    st.stop()
-
+# Initialize session state variables
 if "rendered_pydantic_schema" not in st.session_state:
     st.session_state.rendered_pydantic_schema = ""
 
+if "markdown" not in st.session_state:
+    st.session_state.markdown = ""
+
+if "current_file_id" not in st.session_state:
+    st.session_state.current_file_id = None
+
+# Detect file changes and clear markdown when new file is uploaded
+if in_file is not None:
+    # Create a unique identifier for the current file
+    current_file_id = f"{in_file.name}_{in_file.size}_{hash(in_file.getvalue())}"
+
+    # Check if this is a new file
+    if st.session_state.current_file_id != current_file_id:
+        st.session_state.current_file_id = current_file_id
+        st.session_state.markdown = ""  # Clear markdown for new file
+else:
+    # No file uploaded, clear the current file ID
+    if st.session_state.current_file_id is not None:
+        st.session_state.current_file_id = None
+        st.session_state.markdown = ""  # Clear markdown when no file
+        st.session_state.rendered_pydantic_schema = ""
+
+if in_file is None:
+    st.stop()
+
 filetype = in_file.type
 
 with col1:
@@ -196,11 +222,14 @@ if run_marker:
         )
 
         try:
-            rendered = extract_data(temp_pdf, cli_options, schema)
+            rendered = extract_data(
+                temp_pdf, cli_options, schema, st.session_state.markdown
+            )
 
             with col2:
                 st.write("## Output JSON")
-                st.json(rendered.model_dump())
+                st.json(rendered.model_dump(exclude=["original_markdown"]))
+                st.session_state.markdown = rendered.original_markdown
 
         except Exception as e:
             st.error(f"❌ Extraction failed: {e}")

marker/services/__init__.py

@@ -37,7 +37,7 @@ class BaseService:
         self,
         prompt: str,
         image: PIL.Image.Image | List[PIL.Image.Image] | None,
-        block: Block,
+        block: Block | None,
         response_schema: type[BaseModel],
         max_retries: int | None = None,
         timeout: int | None = None,

marker/services/claude.py

@@ -74,7 +74,7 @@ class ClaudeService(BaseService):
         self,
         prompt: str,
         image: PIL.Image.Image | List[PIL.Image.Image] | None,
-        block: Block,
+        block: Block | None,
         response_schema: type[BaseModel],
         max_retries: int | None = None,
         timeout: int | None = None,

marker/services/gemini.py

@@ -41,7 +41,7 @@ class BaseGeminiService(BaseService):
         self,
         prompt: str,
         image: PIL.Image.Image | List[PIL.Image.Image] | None,
-        block: Block,
+        block: Block | None,
         response_schema: type[BaseModel],
         max_retries: int | None = None,
         timeout: int | None = None,
@@ -72,7 +72,10 @@ class BaseGeminiService(BaseService):
                 )
                 output = responses.candidates[0].content.parts[0].text
                 total_tokens = responses.usage_metadata.total_token_count
-                block.update_metadata(llm_tokens_used=total_tokens, llm_request_count=1)
+                if block:
+                    block.update_metadata(
+                        llm_tokens_used=total_tokens, llm_request_count=1
+                    )
                 return json.loads(output)
             except APIError as e:
                 if e.code in [429, 443, 503]:

marker/services/ollama.py

@@ -35,7 +35,7 @@ class OllamaService(BaseService):
         self,
         prompt: str,
         image: PIL.Image.Image | List[PIL.Image.Image] | None,
-        block: Block,
+        block: Block | None,
         response_schema: type[BaseModel],
         max_retries: int | None = None,
         timeout: int | None = None,
@@ -68,7 +68,9 @@ class OllamaService(BaseService):
             total_tokens = (
                 response_data["prompt_eval_count"] + response_data["eval_count"]
             )
-            block.update_metadata(llm_request_count=1, llm_tokens_used=total_tokens)
+
+            if block:
+                block.update_metadata(llm_request_count=1, llm_tokens_used=total_tokens)
 
             data = response_data["response"]
             return json.loads(data)

marker/services/openai.py

@@ -78,7 +78,7 @@ class OpenAIService(BaseService):
         self,
         prompt: str,
         image: PIL.Image.Image | List[PIL.Image.Image] | None,
-        block: Block,
+        block: Block | None,
         response_schema: type[BaseModel],
         max_retries: int | None = None,
         timeout: int | None = None,
@@ -117,7 +117,10 @@ class OpenAIService(BaseService):
                 )
                 response_text = response.choices[0].message.content
                 total_tokens = response.usage.total_tokens
-                block.update_metadata(llm_tokens_used=total_tokens, llm_request_count=1)
+                if block:
+                    block.update_metadata(
+                        llm_tokens_used=total_tokens, llm_request_count=1
+                    )
                 return json.loads(response_text)
             except (APITimeoutError, RateLimitError) as e:
                 # Rate limit exceeded

tests/extractors/test_basemodel_gen.py

@@ -1,44 +0,0 @@
-from marker.extractors.util import json_schema_to_base_model
-
-
-def test_model_generator():
-    test_schema = {
-        "title": "UserModel",
-        "type": "object",
-        "properties": {
-            "email": {
-                "type": "string",
-                "format": "email",
-                "description": "User's email address",
-            },
-            "age": {"type": "integer", "minimum": 0, "maximum": 150},
-            "name": {"type": "string", "minLength": 1, "maxLength": 100},
-            "status": {"anyOf": [{"type": "string"}, {"type": "null"}]},
-            "tags": {"type": "array", "items": {"type": "string"}},
-            "preferences": {
-                "type": "object",
-                "properties": {
-                    "theme": {"type": "string", "enum": ["dark", "light"]},
-                    "notifications": {"type": "boolean", "default": True},
-                },
-            },
-            "role": {
-                "type": "string",
-                "enum": ["admin", "user", "guest"],
-                "default": "user",
-            },
-        },
-        "required": ["email", "name"],
-    }
-
-    # Create the model
-    UserModel = json_schema_to_base_model(test_schema)
-    user = UserModel(
-        email="test@example.com",
-        name="John Doe",
-        age=30,
-        tags=["python", "pydantic"],
-        preferences={"theme": "dark"},
-        role="admin",
-    )
-    assert user is not None