Merge in dev

.github/workflows/scripts.yml

@@ -0,0 +1,29 @@
+name: Test CLI scripts
+
+on: [push]
+
+env:
+  TORCH_DEVICE: "cpu"
+  OCR_ENGINE: "surya"
+
+jobs:
+  tests:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python 3.11
+        uses: actions/setup-python@v4
+        with:
+          python-version: 3.11
+      - name: Install python dependencies
+        run: |
+          pip install poetry
+          poetry install
+      - name: Download benchmark data
+        run: |
+          wget -O benchmark_data.zip "https://drive.google.com/uc?export=download&id=1NHrdYatR1rtqs2gPVfdvO0BAvocH8CJi"
+          unzip -o benchmark_data.zip
+      - name: Test single script
+        run: poetry run marker_single benchmark_data/pdfs/switch_trans.pdf --page_range 0
+      - name: Test convert script
+        run: poetry run marker benchmark_data/pdfs --max_files 1 --workers 1 --page_range 0

marker/builders/document.py

@@ -43,7 +43,8 @@ class DocumentBuilder(BaseBuilder):
                 page_id=p,
                 lowres_image=lowres_images[i],
                 highres_image=highres_images[i],
-                polygon=provider.get_page_bbox(p)
+                polygon=provider.get_page_bbox(p),
+                refs=provider.get_page_refs(p)
             ) for i, p in enumerate(provider.page_range)
         ]
         DocumentClass: Document = get_block_class(BlockTypes.Document)

marker/converters/pdf.py

@@ -1,13 +1,10 @@
 import os
-
-from marker.processors.llm.llm_handwriting import LLMHandwritingProcessor
-
 os.environ["TOKENIZERS_PARALLELISM"] = "false"  # disables a tokenizers warning
 
 import inspect
 from collections import defaultdict
-from typing import Annotated, Any, Dict, List, Optional, Type, Tuple
 from functools import cache
+from typing import Annotated, Any, Dict, List, Optional, Type, Tuple
 
 from marker.processors import BaseProcessor
 from marker.processors.llm.llm_table_merge import LLMTableMergeProcessor
@@ -33,6 +30,7 @@ from marker.processors.llm.llm_image_description import LLMImageDescriptionProce
 from marker.processors.llm.llm_table import LLMTableProcessor
 from marker.processors.llm.llm_text import LLMTextProcessor
 from marker.processors.page_header import PageHeaderProcessor
+from marker.processors.reference import ReferenceProcessor
 from marker.processors.sectionheader import SectionHeaderProcessor
 from marker.processors.table import TableProcessor
 from marker.processors.text import TextProcessor
@@ -42,6 +40,7 @@ from marker.schema import BlockTypes
 from marker.schema.blocks import Block
 from marker.schema.registry import register_block_class
 from marker.util import strings_to_classes
+from marker.processors.llm.llm_handwriting import LLMHandwritingProcessor
 
 
 class PdfConverter(BaseConverter):
@@ -80,6 +79,7 @@ class PdfConverter(BaseConverter):
         LLMImageDescriptionProcessor,
         LLMEquationProcessor,
         LLMHandwritingProcessor,
+        ReferenceProcessor,
         DebugProcessor,
     )
 

marker/processors/reference.py

@@ -0,0 +1,55 @@
+import numpy as np
+
+from marker.processors import BaseProcessor
+from marker.schema import BlockTypes
+from marker.schema.blocks import Reference
+from marker.schema.document import Document
+from marker.schema.groups.list import ListGroup
+from marker.schema.groups.table import TableGroup
+from marker.schema.registry import get_block_class
+from marker.schema.groups.picture import PictureGroup
+from marker.schema.groups.figure import FigureGroup
+
+
+class ReferenceProcessor(BaseProcessor):
+    """
+    A processor for adding references to the document.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+    def __call__(self, document: Document):
+        ReferenceClass: Reference = get_block_class(BlockTypes.Reference)
+
+        for page in document.pages:
+            refs = page.refs
+            ref_starts = np.array([ref.coord for ref in refs])
+
+            blocks = []
+            for block_id in page.structure:
+                block = page.get_block(block_id)
+                if isinstance(block, (ListGroup, FigureGroup, TableGroup)):
+                    blocks.extend([page.get_block(b) for b in block.structure])
+                else:
+                    blocks.append(block)
+            blocks = [b for b in blocks if not b.ignore_for_output]
+
+            block_starts = np.array([block.polygon.bbox[:2] for block in blocks])
+
+            if not (len(refs) and len(block_starts)):
+                continue
+
+            distances = np.linalg.norm(block_starts[:, np.newaxis, :] - ref_starts[np.newaxis, :, :], axis=2)
+            for ref_idx in range(len(ref_starts)):
+                block_idx = np.argmin(distances[:, ref_idx])
+                block = blocks[block_idx]
+
+                ref_block = page.add_full_block(ReferenceClass(
+                    ref=refs[ref_idx].ref,
+                    polygon=block.polygon,
+                    page_id=page.page_id
+                ))
+                if block.structure is None:
+                    block.structure = []
+                block.structure.insert(0, ref_block.id)

marker/providers/pdf.py

@@ -1,12 +1,13 @@
 import atexit
 import ctypes
 import re
-from typing import Annotated, List, Optional, Set
+from typing import Annotated, Dict, List, Optional, Set
 
 import pypdfium2 as pdfium
 import pypdfium2.raw as pdfium_c
 from ftfy import fix_text
 from pdftext.extraction import dictionary_output
+from pdftext.schema import Reference
 from PIL import Image
 from pypdfium2 import PdfiumError
 
@@ -75,6 +76,7 @@ class PdfProvider(BaseProvider):
 
         self.doc: pdfium.PdfDocument = pdfium.PdfDocument(self.filepath)
         self.page_lines: ProviderPageLines = {i: [] for i in range(len(self.doc))}
+        self.page_refs: Dict[int, List[Reference]] = {i: [] for i in range(len(self.doc))}
 
         if self.page_range is None:
             self.page_range = range(len(self.doc))
@@ -210,7 +212,6 @@ class PdfProvider(BaseProvider):
                                 page_id=page_id,
                                 text_extraction_method="pdftext",
                                 url=span.get("url"),
-                                anchors=span.get("anchors"),
                             )
                         )
                     polygon = PolygonBox.from_bbox(line["bbox"], ensure_nonzero_area=True)
@@ -222,6 +223,7 @@ class PdfProvider(BaseProvider):
                     )
             if self.check_line_spans(lines):
                 page_lines[page_id] = lines
+            self.page_refs[page_id] = page["refs"]
 
         return page_lines
 
@@ -326,6 +328,9 @@ class PdfProvider(BaseProvider):
     def get_page_lines(self, idx: int) -> List[ProviderOutput]:
         return self.page_lines[idx]
 
+    def get_page_refs(self, idx: int):
+        return self.page_refs[idx]
+
     @staticmethod
     def _get_fontname(font) -> str:
         font_name = ""

marker/renderers/__init__.py

@@ -15,7 +15,6 @@ from marker.util import assign_config
 
 
 class BaseRenderer:
-    remove_blocks: Annotated[Tuple[BlockTypes, ...], "The block types to ignore while rendering."] = (BlockTypes.PageHeader, BlockTypes.PageFooter)
     image_blocks: Annotated[Tuple[BlockTypes, ...], "The block types to consider as images."] = (BlockTypes.Picture, BlockTypes.Figure)
     extract_images: Annotated[bool, "Extract images from the document."] = True
     image_extraction_mode: Annotated[

marker/renderers/html.py

@@ -60,14 +60,12 @@ class HTMLRenderer(BaseRenderer):
                     ref_block_id: BlockId = item.id
                     break
 
-            if ref_block_id.block_type in self.remove_blocks:
-                ref.replace_with('')
-            elif ref_block_id.block_type in self.image_blocks:
+            if ref_block_id.block_type in self.image_blocks:
                 if self.extract_images:
                     image = self.extract_image(document, ref_block_id)
                     image_name = f"{ref_block_id.to_path()}.{settings.OUTPUT_IMAGE_FORMAT.lower()}"
                     images[image_name] = image
-                    ref.replace_with(BeautifulSoup(f"<p><img src='{image_name}'></p>", 'html.parser'))
+                    ref.replace_with(BeautifulSoup(f"<p>{content}<img src='{image_name}'></p>", 'html.parser'))
                 else:
                     # This will be the image description if using llm mode, or empty if not
                     ref.replace_with(BeautifulSoup(f"{content}", 'html.parser'))

marker/schema/__init__.py

@@ -28,6 +28,7 @@ class BlockTypes(str, Enum):
     Document = auto()
     ComplexRegion = auto()
     TableCell = auto()
+    Reference = auto()
 
     def __str__(self):
         return self.name

marker/schema/blocks/__init__.py

@@ -19,3 +19,4 @@ from marker.schema.blocks.text import Text
 from marker.schema.blocks.toc import TableOfContents
 from marker.schema.blocks.complexregion import ComplexRegion
 from marker.schema.blocks.tablecell import TableCell
+from marker.schema.blocks.reference import Reference

marker/schema/blocks/base.py

@@ -12,6 +12,7 @@ if TYPE_CHECKING:
     from marker.schema.document import Document
     from marker.schema.groups.page import PageGroup
 
+
 class BlockMetadata(BaseModel):
     llm_request_count: int = 0
     llm_error_count: int = 0
@@ -78,6 +79,7 @@ class Block(BaseModel):
     text_extraction_method: Optional[Literal['pdftext', 'surya', 'gemini']] = None
     structure: List[BlockId] | None = None  # The top-level page structure, which is the block ids in order
     ignore_for_output: bool = False  # Whether this block should be ignored in output
+    replace_output_newlines: bool = False  # Whether to replace newlines with spaces in output
     source: Literal['layout', 'heuristics', 'processor'] = 'layout'
     top_k: Optional[Dict[BlockTypes, float]] = None
     metadata: BlockMetadata | None = None
@@ -187,6 +189,10 @@ class Block(BaseModel):
         template = ""
         for c in child_blocks:
             template += f"<content-ref src='{c.id}'></content-ref>"
+
+        if self.replace_output_newlines:
+            template = "<p>" + template.replace("\n", " ") + "</p>"
+
         return template
 
     def assign_section_hierarchy(self, section_hierarchy):

marker/schema/blocks/basetable.py

@@ -24,13 +24,16 @@ class BaseTable(Block):
 
 
     def assemble_html(self, document, child_blocks: List[BlockOutput], parent_structure=None):
+        # Filter out the table cells, so they don't render twice
+        selected_blocks = [b for b in child_blocks if b.id.block_type != BlockTypes.TableCell]
+        template = super().assemble_html(document, selected_blocks, parent_structure)
+
         if self.html:
             # LLM processor
-            return self.html
+            return template + self.html
         elif len(child_blocks) > 0 and child_blocks[0].id.block_type == BlockTypes.TableCell:
             # Table processor
-            return self.format_cells(document, child_blocks)
+            return template + self.format_cells(document, child_blocks)
         else:
             # Default text lines and spans
-            template = super().assemble_html(document, child_blocks, parent_structure)
             return f"<p>{template}</p>"

marker/schema/blocks/caption.py

@@ -5,8 +5,10 @@ from marker.schema.blocks import Block
 class Caption(Block):
     block_type: BlockTypes = BlockTypes.Caption
     block_description: str = "A text caption that is directly above or below an image or table. Only used for text describing the image or table.  "
+    replace_output_newlines: bool = True
 
     def assemble_html(self, document, child_blocks, parent_structure):
         template = super().assemble_html(document, child_blocks, parent_structure)
         template = template.replace("\n", " ")
         return f"<p>{template}</p>"
+

marker/schema/blocks/equation.py

@@ -11,7 +11,9 @@ class Equation(Block):
 
     def assemble_html(self, document, child_blocks, parent_structure=None):
         if self.latex:
-            html_out = f"<p block-type='{self.block_type}'>"
+            child_ref_blocks = [block for block in child_blocks if block.id.block_type == BlockTypes.Reference]
+            html_out = super().assemble_html(child_ref_blocks, parent_structure)
+            html_out += f"<p block-type='{self.block_type}'>"
 
             try:
                 latex = self.parse_latex(html.escape(self.latex))
@@ -44,9 +46,9 @@ class Equation(Block):
             ("$$", "block"),
             ("$", "inline")
         ]
-        
-        text = text.replace("\n", "<br>") # we can't handle \n's inside <p> properly if we don't do this
-       
+
+        text = text.replace("\n", "<br>")  # we can't handle \n's inside <p> properly if we don't do this
+
         i = 0
         stack = []
         result = []
@@ -73,7 +75,7 @@ class Equation(Block):
             else:  # No delimiter match
                 buffer += text[i]
                 i += 1
-                
+
         if buffer:
             result.append({"class": "text", "content": buffer})
-        return result
+        return result

marker/schema/blocks/figure.py

@@ -8,7 +8,8 @@ class Figure(Block):
     block_description: str = "A chart or other image that contains data."
 
     def assemble_html(self, document, child_blocks, parent_structure):
+        child_ref_blocks = [block for block in child_blocks if block.id.block_type == BlockTypes.Reference]
+        html = super().assemble_html(document, child_ref_blocks, parent_structure)
         if self.description:
-            return f"<p role='img' data-original-image-id='{self.id}'>Image {self.id} description: {self.description}</p>"
-        else:
-            return ""
+            html += f"<p role='img' data-original-image-id='{self.id}'>Image {self.id} description: {self.description}</p>"
+        return html

marker/schema/blocks/footnote.py

@@ -5,9 +5,11 @@ from marker.schema.blocks import Block
 class Footnote(Block):
     block_type: BlockTypes = BlockTypes.Footnote
     block_description: str = "A footnote that explains a term or concept in the document."
+    replace_output_newlines: bool = True
 
     def assemble_html(self, document, child_blocks, parent_structure):
         template = super().assemble_html(document, child_blocks, parent_structure)
         template = template.replace("\n", " ")
 
         return f"<p>{template}</p>"
+

marker/schema/blocks/handwriting.py

@@ -6,6 +6,7 @@ class Handwriting(Block):
     block_type: BlockTypes = BlockTypes.Handwriting
     block_description: str = "A region that contains handwriting."
     html: str | None = None
+    replace_output_newlines: bool = True
 
     def assemble_html(self, document, child_blocks, parent_structure):
         if self.html:
@@ -14,3 +15,4 @@ class Handwriting(Block):
             template = super().assemble_html(document, child_blocks, parent_structure)
             template = template.replace("\n", " ")
             return f"<p>{template}</p>"
+

marker/schema/blocks/pagefooter.py

@@ -5,6 +5,8 @@ from marker.schema.blocks import Block
 class PageFooter(Block):
     block_type: str = BlockTypes.PageFooter
     block_description: str = "Text that appears at the bottom of a page, like a page number."
+    replace_output_newlines: bool = True
+    ignore_for_output: bool = True
 
     def assemble_html(self, document, child_blocks, parent_structure):
         if self.ignore_for_output:
@@ -13,3 +15,4 @@ class PageFooter(Block):
         template = super().assemble_html(document, child_blocks, parent_structure)
         template = template.replace("\n", " ")
         return f"<p>{template}</p>"
+

marker/schema/blocks/pageheader.py

@@ -5,6 +5,8 @@ from marker.schema.blocks import Block
 class PageHeader(Block):
     block_type: BlockTypes = BlockTypes.PageHeader
     block_description: str = "Text that appears at the top of a page, like a page title."
+    replace_output_newlines: bool = True
+    ignore_for_output: bool = True
 
     def assemble_html(self, document, child_blocks, parent_structure):
         if self.ignore_for_output:
@@ -13,3 +15,4 @@ class PageHeader(Block):
         template = super().assemble_html(document, child_blocks, parent_structure)
         template = template.replace("\n", " ")
         return f"<p>{template}</p>"
+

marker/schema/blocks/reference.py

@@ -0,0 +1,11 @@
+from marker.schema import BlockTypes
+from marker.schema.blocks import Block
+
+
+class Reference(Block):
+    block_type: BlockTypes = BlockTypes.Reference
+    ref: str
+
+    def assemble_html(self, child_blocks, parent_structure=None):
+        template = super().assemble_html(child_blocks, parent_structure)
+        return f"<span id='{self.ref}'>{template}</span>"

marker/schema/groups/page.py

@@ -3,6 +3,7 @@ from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 from PIL import Image
 
+from pdftext.schema import Reference
 from marker.providers import ProviderOutput
 from marker.schema import BlockTypes
 from marker.schema.blocks import Block, BlockId, Text
@@ -23,6 +24,7 @@ class PageGroup(Group):
     excluded_block_types: Sequence[BlockTypes] = (BlockTypes.Line, BlockTypes.Span,)
     maximum_assignment_distance: float = 20  # pixels
     block_description: str = "A single page in the document."
+    refs: List[Reference] | None = None
 
     def incr_block_id(self):
         if self.block_id is None:

marker/schema/registry.py

@@ -6,7 +6,7 @@ from marker.schema.blocks import Block, Caption, Code, Equation, Figure, \
     Footnote, Form, Handwriting, InlineMath, \
     ListItem, PageFooter, PageHeader, Picture, \
     SectionHeader, Table, TableOfContents, \
-    Text, ComplexRegion, TableCell
+    Text, ComplexRegion, TableCell, Reference
 from marker.schema.document import Document
 from marker.schema.groups import FigureGroup, ListGroup, PageGroup, \
     PictureGroup, TableGroup
@@ -51,6 +51,7 @@ register_block_class(BlockTypes.Text, Text)
 register_block_class(BlockTypes.TableOfContents, TableOfContents)
 register_block_class(BlockTypes.ComplexRegion, ComplexRegion)
 register_block_class(BlockTypes.TableCell, TableCell)
+register_block_class(BlockTypes.Reference, Reference)
 register_block_class(BlockTypes.Document, Document)
 
 assert len(BLOCK_REGISTRY) == len(BlockTypes)

marker/schema/text/span.py

@@ -25,7 +25,6 @@ class Span(Block):
     formats: List[Literal['plain', 'math', 'chemical', 'bold', 'italic']]
     has_superscript: bool = False
     url: Optional[str] = None
-    anchors: Optional[List[str]] = None
 
     @property
     def bold(self):
@@ -75,6 +74,4 @@ class Span(Block):
         elif self.math:
             text = f"<math display='inline'>{text}</math>"
 
-        if self.anchors:
-            text = "".join(f"<span id='{anchor}'/>" for anchor in self.anchors) + text
         return text

tests/builders/test_pdf_links.py

@@ -1,3 +1,5 @@
+import re
+
 import pytest
 
 from marker.converters.pdf import PdfConverter
@@ -8,9 +10,8 @@ from marker.schema.document import Document
 
 @pytest.mark.filename("arxiv_test.pdf")
 @pytest.mark.output_format("markdown")
-@pytest.mark.config({"page_range": [1]})
 def test_pdf_links(pdf_document: Document, pdf_converter: PdfConverter, temp_pdf):
-    first_page = pdf_document.pages[0]
+    first_page = pdf_document.pages[1]
 
     for section_header_span in first_page.contained_blocks(pdf_document, (BlockTypes.Span,)):
         if "II." in section_header_span.text:
@@ -22,11 +23,13 @@ def test_pdf_links(pdf_document: Document, pdf_converter: PdfConverter, temp_pdf
     section_header_block = first_page.contained_blocks(pdf_document, (BlockTypes.SectionHeader,))[0]
     assert section_header_block.raw_text(pdf_document) == 'II. THEORETICAL FRAMEWORK\n'
 
-    section_header_span = section_header_block.contained_blocks(pdf_document, (BlockTypes.Span,))[0]
-    assert section_header_span.anchors == ['page-1-0']
+    assert first_page.refs[0].ref == "page-1-0"
 
     markdown_output: MarkdownOutput = pdf_converter(temp_pdf.name)
     markdown = markdown_output.markdown
 
     assert '[II.](#page-1-0)' in markdown
     assert '<span id="page-1-0"/>II. THEORETICAL FRAMEWORK' in markdown
+
+    for ref in set([f'<span id="page-{m[0]}-{m[1]}"/>' for m in re.findall(r'\]\(#page-(\d+)-(\d+)\)', markdown)]):
+        assert ref in markdown, f"Reference {ref} not found in markdown"