Add postprocessor

README.md

@@ -1,6 +1,6 @@
 # Marker
 
-Marker converts PDF, EPUB, and MOBI to Markdown.  It is up to 10x faster than nougat, works across many types of documents, and minimizes the risk of hallucinations significantly.
+Marker converts PDF, EPUB, and MOBI to Markdown.  It is 12x faster than nougat, works across many types of documents, and minimizes the risk of hallucinations significantly.
 
 Features:
 
@@ -115,18 +115,33 @@ METADATA_FILE=../pdf_meta.json NUM_DEVICES=4 NUM_WORKERS=35 bash chunk_convert.s
 
 # Benchmarks
 
-Benchmarking PDF extraction quality is hard.  I've created a test set by finding books and scientific papers that have a pdf version and a latex source.  I can then convert the latex to text, and compare it to the output of marker using edit distance.
+Benchmarking PDF extraction quality is hard.  I've created a test set by finding books and scientific papers that have a pdf version and a latex source.  I can then convert the latex to text, and compare the reference to the output of text extraction methods.
 
-Benchmarks show that marker is up to 10x faster than nougat, and more accurate outside arXiv (nougat is better inside arXiv):
+Benchmarks show that marker is 12x faster than nougat, and more accurate outside arXiv (nougat was trained on arXiv data).
 
+**Speed**
 
+Method      Average Score    Time per doc
+--------  ---------------  --------------
+naive            0.287605      0.149704
+marker           0.62978       33.9778
+nougat           0.63989       395.091
 
+**Accuracy**
+
+First 3 are non-arXiv books, last 3 are arXiv papers.
+
+Method      thinkos.pdf    thinkdsp.pdf    thinkpython.pdf    switch_trans.pdf    crowd.pdf    multicolcnn.pdf
+--------  -------------  --------------  -----------------  ------------------  -----------  -----------------
+naive          0.366817        0.412014           0.468147             0.244739     0.14489           0.0890217
+marker         0.753291        0.787938           0.779262             0.478387     0.446068          0.533737
+nougat         0.638434        0.632723           0.637626             0.690028     0.540994          0.699539
 
 Peak GPU memory usage during the benchmark is `3.3GB` for nougat, and `3.7GB` for marker.
 
 ## Running your own benchmarks
 
-You can benchmark the performance of marker on your machine.  The benchmark consists of 3 scientific papers from arXiv, and 3 textbooks. 
+You can benchmark the performance of marker on your machine.
 
 Run `benchmark.py` like this:
 
@@ -134,7 +149,7 @@ Run `benchmark.py` like this:
 python benchmark.py benchmark_data/pdfs benchmark_data/references report.json --nougat
 ```
 
-This will benchmark marker against other text extraction methods.  It sets up batch sizes for nougat and marker to use a similar amount of GPU RAM for each (4GB).
+This will benchmark marker against other text extraction methods.  It sets up batch sizes for nougat and marker to use a similar amount of GPU RAM for each.
 
 Omit `--nougat` to exclude nougat from the benchmark.  I don't recommend running nougat on CPU, since it is very slow.
 

benchmark.py

@@ -7,6 +7,7 @@ from tqdm import tqdm
 
 from marker.convert import convert_single_pdf
 from marker.logger import configure_logging
+from marker.models import load_all_models
 from marker.ordering import load_ordering_model
 from marker.segmentation import load_layout_model
 from marker.cleaners.equations import load_nougat_model
@@ -48,9 +49,7 @@ if __name__ == "__main__":
     if args.nougat:
         methods.append("nougat")
 
-    layoutlm_model = load_layout_model()
-    nougat_model = load_nougat_model()
-    order_model = load_ordering_model()
+    model_lst = load_all_models()
 
     scores = defaultdict(dict)
     benchmark_files = os.listdir(args.in_folder)
@@ -70,7 +69,7 @@ if __name__ == "__main__":
         for method in methods:
             start = time.time()
             if method == "marker":
-                full_text, out_meta = convert_single_pdf(pdf_filename, layoutlm_model, nougat_model, order_model, parallel=args.marker_parallel)
+                full_text, out_meta = convert_single_pdf(pdf_filename, model_lst, parallel=args.marker_parallel)
             elif method == "nougat":
                 full_text = nougat_prediction(pdf_filename, batch_size=args.nougat_batch_size)
             elif method == "naive":

convert.py

@@ -8,6 +8,7 @@ from tqdm import tqdm
 import math
 
 from marker.convert import convert_single_pdf, get_length_of_text
+from marker.models import load_all_models
 from marker.ordering import load_ordering_model
 from marker.segmentation import load_layout_model
 from marker.cleaners.equations import load_nougat_model
@@ -20,7 +21,7 @@ configure_logging()
 
 
 @ray.remote(num_cpus=settings.RAY_CORES_PER_WORKER, num_gpus=.05 if settings.CUDA else 0)
-def process_single_pdf(fname: str, out_folder: str, nougat_model, layout_model, order_model, metadata: Dict | None=None, min_length: int | None = None):
+def process_single_pdf(fname: str, out_folder: str, model_refs, metadata: Dict | None=None, min_length: int | None = None):
     out_filename = fname.rsplit(".", 1)[0] + ".md"
     out_filename = os.path.join(out_folder, os.path.basename(out_filename))
     out_meta_filename = out_filename.rsplit(".", 1)[0] + "_meta.json"
@@ -35,7 +36,7 @@ def process_single_pdf(fname: str, out_folder: str, nougat_model, layout_model,
             if length < min_length:
                 return
 
-        full_text, out_metadata = convert_single_pdf(fname, layout_model, nougat_model, order_model, metadata=metadata)
+        full_text, out_metadata = convert_single_pdf(fname, model_refs, metadata=metadata)
         if len(full_text.strip()) > 0:
             with open(out_filename, "w+") as f:
                 f.write(full_text)
@@ -94,13 +95,8 @@ if __name__ == "__main__":
         log_to_driver=False
     )
 
-    nougat_model = load_nougat_model()
-    layoutlm_model = load_layout_model()
-    order_model = load_ordering_model()
-
-    nougat_ref = ray.put(nougat_model)
-    layoutlm_ref = ray.put(layoutlm_model)
-    order_ref = ray.put(order_model)
+    model_lst = load_all_models()
+    model_refs = [ray.put(m) if m else None for m in model_lst]
 
     # Dynamically set GPU allocation per task based on GPU ram
     gpu_frac = settings.INFERENCE_RAM // settings.VRAM_PER_TASK if settings.CUDA else 0
@@ -110,9 +106,7 @@ if __name__ == "__main__":
         process_single_pdf.options(num_gpus=gpu_frac).remote(
             filename,
             out_folder,
-            nougat_ref,
-            layoutlm_ref,
-            order_ref,
+            model_refs,
             metadata=metadata.get(os.path.basename(filename)),
             min_length=args.min_length
         ) for filename in files_to_convert

convert_single.py

@@ -2,6 +2,7 @@ import argparse
 
 from marker.convert import convert_single_pdf
 from marker.logger import configure_logging
+from marker.models import load_all_models
 from marker.ordering import load_ordering_model
 from marker.segmentation import load_layout_model
 from marker.cleaners.equations import load_nougat_model
@@ -19,10 +20,8 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     fname = args.filename
-    layoutlm_model = load_layout_model()
-    nougat_model = load_nougat_model()
-    order_model = load_ordering_model()
-    full_text, out_meta = convert_single_pdf(fname, layoutlm_model, nougat_model, order_model, max_pages=args.max_pages, parallel=args.workers)
+    model_lst = load_all_models()
+    full_text, out_meta = convert_single_pdf(fname, model_lst, max_pages=args.max_pages, parallel=args.workers)
 
     with open(args.output, "w+") as f:
         f.write(full_text)

marker/convert.py

@@ -5,6 +5,7 @@ from marker.extract_text import get_text_blocks
 from marker.cleaners.headers import filter_header_footer, filter_common_titles
 from marker.cleaners.equations import replace_equations
 from marker.ordering import order_blocks
+from marker.postprocessors.editor import edit_full_text
 from marker.segmentation import detect_all_block_types
 from marker.cleaners.code import identify_code_blocks, indent_blocks
 from marker.cleaners.bullets import replace_bullets
@@ -56,9 +57,7 @@ def get_length_of_text(fname: str) -> int:
 
 def convert_single_pdf(
         fname: str,
-        layoutlm_model,
-        nougat_model,
-        order_model,
+        model_lst: List,
         max_pages=None,
         metadata: Dict | None=None,
         parallel: int = 1
@@ -96,6 +95,9 @@ def convert_single_pdf(
         print(f"Could not extract any text blocks for {fname}")
         return "", out_meta
 
+    # Unpack models from list
+    nougat_model, layoutlm_model, order_model, edit_model = model_lst
+
     block_types = detect_all_block_types(doc, blocks, layoutlm_model, parallel=parallel)
 
     # Find headers and footers
@@ -135,5 +137,7 @@ def convert_single_pdf(
 
     # Replace bullet characters with a -
     full_text = replace_bullets(full_text)
+    full_text, edit_stats = edit_full_text(full_text, edit_model)
+    out_meta["postprocess_stats"] = {"edit": edit_stats}
 
     return full_text, out_meta

marker/models.py

@@ -0,0 +1,13 @@
+from marker.cleaners.equations import load_nougat_model
+from marker.ordering import load_ordering_model
+from marker.postprocessors.editor import load_editing_model
+from marker.segmentation import load_layout_model
+
+
+def load_all_models():
+    edit = load_editing_model()
+    order = load_ordering_model()
+    layout = load_layout_model()
+    nougat = load_nougat_model()
+    model_lst = [nougat, layout, order, edit]
+    return model_lst

marker/postprocessors/editor.py

@@ -0,0 +1,130 @@
+from collections import defaultdict, Counter
+from itertools import chain
+from typing import Optional
+import re
+
+from transformers import BloomForTokenClassification, AutoTokenizer, DataCollatorForTokenClassification
+from marker.settings import settings
+import torch
+
+tokenizer = AutoTokenizer.from_pretrained(settings.EDITOR_MODEL_NAME)
+
+
+def load_editing_model(disable_editor=False):
+    if disable_editor:
+        return None
+
+    if not settings.CUDA:
+        # Don't postprocess on CPU to save time
+        return None
+
+    model = BloomForTokenClassification.from_pretrained(
+        settings.EDITOR_MODEL_NAME,
+        load_in_4bit=True,
+        torch_dtype=torch.bfloat16,
+        bnb_4bit_compute_dtype=torch.bfloat16,
+        bnb_4bit_quant_type="nf4",
+        device_map="sequential"
+    )
+
+    model.config.label2id = {
+        "equal": 0,
+        "delete": 1,
+        "delete_trailing_newline": 2,
+        "delete_leading_space": 3,
+        "leading_space_to_newline": 4,
+        "newline-1": 5,
+        "space-1": 6,
+    }
+    model.config.id2label = {v: k for k, v in model.config.label2id.items()}
+    return model
+
+
+def edit_full_text(text: str, model: Optional[BloomForTokenClassification]):
+    if not model:
+        return text
+
+    tokenized = tokenizer(
+        text,
+        truncation=True,
+        max_length=settings.EDITOR_MAX_LENGTH,
+        return_overflowing_tokens=True,
+        padding="max_length",
+    )
+    input_ids = tokenized["input_ids"]
+
+    # Tokenize, and make sure reverse tokenization works
+    model_tokens = [tokenizer.convert_ids_to_tokens(t, skip_special_tokens=True) for t in input_ids]
+    model_str_tokens = [tokenizer.convert_tokens_to_string(t) for t in model_tokens]
+    full_text = "".join(model_str_tokens)
+    assert full_text == text
+
+    # Long list of all tokens
+    model_tokens = [tokenizer.convert_ids_to_tokens(t) for t in input_ids]
+    flat_tokens = list(chain.from_iterable(model_tokens))
+    flat_str_tokens = [tokenizer.convert_tokens_to_string([t]) for t in flat_tokens]
+
+    # Run model
+    token_masks = []
+    for i in range(0, len(input_ids), settings.EDITOR_BATCH_SIZE):
+        batch_input_ids = tokenized["input_ids"][i: i + settings.EDITOR_BATCH_SIZE]
+        batch_input_ids = torch.tensor(batch_input_ids, device=model.device)
+        batch_attention_mask = tokenized["attention_mask"][i: i + settings.EDITOR_BATCH_SIZE]
+        batch_attention_mask = torch.tensor(batch_attention_mask, device=model.device)
+        with torch.inference_mode():
+            predictions = model(batch_input_ids, attention_mask=batch_attention_mask)
+
+        logits = predictions.logits.cpu()
+
+        labels = logits.argmax(-1).squeeze().tolist()
+        labels = list(chain.from_iterable(labels))
+        token_masks.extend(labels)
+
+    assert len(token_masks) == len(flat_tokens) == len(flat_str_tokens)
+
+    edit_stats = defaultdict(int)
+    out_tokens = []
+    for i, (token, str_token, mask) in enumerate(zip(flat_tokens, flat_str_tokens, token_masks)):
+        label = model.config.id2label[mask]
+
+        match label:
+            case "equal":
+                out_tokens.append(str_token)
+                edit_stats[label] += 1
+            case "delete":
+                # If we delete whitespace, roll with it, otherwise ignore
+                if str_token.strip():
+                    out_tokens.append(str_token)
+                edit_stats[label] += 1
+            case "delete_trailing_newline":
+                if str_token.endswith("\n"):
+                    str_token = re.sub(r"\n+$", "", str_token)
+                    edit_stats[label] += 1
+                out_tokens.append(str_token)
+
+            case "delete_leading_space":
+                if str_token.startswith(" "):
+                    str_token = re.sub(r"^ +", "", str_token)
+                    edit_stats[label] += 1
+                out_tokens.append(str_token)
+            case "leading_space_to_newline":
+                if str_token.startswith(" "):
+                    str_token = "\n" + str_token[1:]
+                    edit_stats[label] += 1
+                out_tokens.append(str_token)
+            case "newline-1":
+                out_tokens.append("\n")
+                out_tokens.append(str_token)
+                edit_stats[label] += 1
+            case "space-1":
+                out_tokens.append(" ")
+                out_tokens.append(str_token)
+                edit_stats[label] += 1
+
+    return "".join(out_tokens), edit_stats
+
+
+
+
+
+

marker/settings.py

@@ -67,6 +67,11 @@ class Settings(BaseSettings):
     ORDERER_BATCH_SIZE: int = 16 # This can be high, because max token count is 128
     ORDERER_MODEL_NAME: str = "vikp/column_detector"
 
+    # Final editing model
+    EDITOR_BATCH_SIZE: int = 4
+    EDITOR_MAX_LENGTH: int = 1024
+    EDITOR_MODEL_NAME: str = "vikp/pdf_postprocessor"
+
     # Ray
     RAY_CACHE_PATH: Optional[str] = None # Where to save ray cache
     RAY_DASHBOARD_HOST: str = "127.0.0.1"

poetry.lock

@@ -361,6 +361,17 @@ soupsieve = ">1.2"
 html5lib = ["html5lib"]
 lxml = ["lxml"]
 
+[[package]]
+name = "bitsandbytes"
+version = "0.41.2.post2"
+description = "k-bit optimizers and matrix multiplication routines."
+optional = false
+python-versions = "*"
+files = [
+    {file = "bitsandbytes-0.41.2.post2-py3-none-any.whl", hash = "sha256:98e5e1979aea3d481ed06181c689f3a154d7f5dc1af770c5173485bc54cf7b72"},
+    {file = "bitsandbytes-0.41.2.post2.tar.gz", hash = "sha256:d374da4700651f36a285ed53e012ee527736109614e3f5c0249985d41027136d"},
+]
+
 [[package]]
 name = "bleach"
 version = "6.1.0"
@@ -5484,4 +5495,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p
 [metadata]
 lock-version = "2.0"
 python-versions = ">=3.9,<3.13"
-content-hash = "2fb598c6fe9ac11179d892a4c57a776a3274f2a26b24e4db308e44871ed2f8d3"
+content-hash = "867abbd491c21af26d74884792e63116aab25a1a362e1c719dfe145c6cc3c2bd"

pyproject.toml

@@ -28,6 +28,7 @@ pyspellchecker = "^0.7.2"
 ftfy = "^6.1.1"
 nltk = "^3.8.1"
 ocrmypdf = "^15.4.0"
+bitsandbytes = "^0.41.2.post2"
 
 
 [tool.poetry.group.dev.dependencies]