marker-pdf 0.1.0__py3-none-any.whl

benchmark.py

@@ -0,0 +1,128 @@
+import argparse
+import tempfile
+import time
+from collections import defaultdict
+
+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
+from marker.benchmark.scoring import score_text
+from marker.extract_text import naive_get_text
+import json
+import os
+import subprocess
+import shutil
+import fitz as pymupdf
+from marker.settings import settings
+from tabulate import tabulate
+
+configure_logging()
+
+
+def nougat_prediction(pdf_filename, batch_size=1):
+    out_dir = tempfile.mkdtemp()
+    subprocess.run(["nougat", pdf_filename, "-o", out_dir, "--no-skipping", "--recompute", "--batchsize", str(batch_size)], check=True)
+    md_file = os.listdir(out_dir)[0]
+    with open(os.path.join(out_dir, md_file), "r") as f:
+        data = f.read()
+    shutil.rmtree(out_dir)
+    return data
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Benchmark PDF to MD conversion.  Needs source pdfs, and a refernece folder with the correct markdown.")
+    parser.add_argument("in_folder", help="Input PDF files")
+    parser.add_argument("reference_folder", help="Reference folder with reference markdown files")
+    parser.add_argument("out_file", help="Output filename")
+    parser.add_argument("--nougat", action="store_true", help="Run nougat and compare", default=False)
+    # Nougat batch size 1 uses about as much VRAM as default marker settings
+    parser.add_argument("--nougat_batch_size", type=int, default=1, help="Batch size to use for nougat when making predictions.")
+    parser.add_argument("--marker_parallel_factor", type=int, default=1, help="How much to multiply default parallel OCR workers and model batch sizes by.")
+    parser.add_argument("--md_out_path", type=str, default=None, help="Output path for generated markdown files")
+    args = parser.parse_args()
+
+    methods = ["naive", "marker"]
+    if args.nougat:
+        methods.append("nougat")
+
+    model_lst = load_all_models()
+
+    scores = defaultdict(dict)
+    benchmark_files = os.listdir(args.in_folder)
+    benchmark_files = [b for b in benchmark_files if b.endswith(".pdf")]
+    times = defaultdict(dict)
+    pages = defaultdict(int)
+
+    for fname in tqdm(benchmark_files):
+        md_filename = fname.rsplit(".", 1)[0] + ".md"
+
+        reference_filename = os.path.join(args.reference_folder, md_filename)
+        with open(reference_filename, "r") as f:
+            reference = f.read()
+
+        pdf_filename = os.path.join(args.in_folder, fname)
+        doc = pymupdf.open(pdf_filename)
+        pages[fname] = len(doc)
+
+        for method in methods:
+            start = time.time()
+            if method == "marker":
+                full_text, out_meta = convert_single_pdf(pdf_filename, model_lst, parallel_factor=args.marker_parallel_factor)
+            elif method == "nougat":
+                full_text = nougat_prediction(pdf_filename, batch_size=args.nougat_batch_size)
+            elif method == "naive":
+                full_text = naive_get_text(doc)
+            else:
+                raise ValueError(f"Unknown method {method}")
+
+            times[method][fname] = time.time() - start
+
+            score = score_text(full_text, reference)
+            scores[method][fname] = score
+
+            if args.md_out_path:
+                md_out_filename = f"{method}_{md_filename}"
+                with open(os.path.join(args.md_out_path, md_out_filename), "w+") as f:
+                    f.write(full_text)
+
+    total_pages = sum(pages.values())
+    with open(args.out_file, "w+") as f:
+        write_data = defaultdict(dict)
+        for method in methods:
+            total_time = sum(times[method].values())
+            file_stats = {
+                fname:
+                {
+                    "time": times[method][fname],
+                    "score": scores[method][fname],
+                    "pages": pages[fname]
+                }
+
+                for fname in benchmark_files
+            }
+            write_data[method] = {
+                "files": file_stats,
+                "avg_score": sum(scores[method].values()) / len(scores[method]),
+                "time_per_page": total_time / total_pages,
+                "time_per_doc": total_time / len(scores[method])
+            }
+
+        json.dump(write_data, f, indent=4)
+
+    summary_table = []
+    score_table = []
+    score_headers = benchmark_files
+    for method in methods:
+        summary_table.append([method, write_data[method]["avg_score"], write_data[method]["time_per_page"], write_data[method]["time_per_doc"]])
+        score_table.append([method, *[write_data[method]["files"][h]["score"] for h in score_headers]])
+
+    print(tabulate(summary_table, headers=["Method", "Average Score", "Time per page", "Time per document"]))
+    print("")
+    print("Scores by file")
+    print(tabulate(score_table, headers=["Method", *score_headers]))
+

chunk_convert.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+
+trap 'pkill -P $$' SIGINT
+
+# Check if NUM_DEVICES is set
+if [[ -z "$NUM_DEVICES" ]]; then
+    echo "Please set the NUM_DEVICES environment variable."
+    exit 1
+fi
+
+if [[ -z "$NUM_WORKERS" ]]; then
+    echo "Please set the NUM_WORKERS environment variable."
+    exit 1
+fi
+
+
+# Get input folder and output folder from args
+if [[ -z "$1" ]]; then
+    echo "Please provide an input folder."
+    exit 1
+fi
+
+if [[ -z "$2" ]]; then
+    echo "Please provide an output folder."
+    exit 1
+fi
+
+INPUT_FOLDER=$1
+OUTPUT_FOLDER=$2
+
+# Loop from 0 to NUM_DEVICES and run the Python script in parallel
+for (( i=0; i<$NUM_DEVICES; i++ )); do
+    DEVICE_NUM=$i
+    export DEVICE_NUM
+    export NUM_DEVICES
+    export NUM_WORKERS
+    echo "Running convert.py on GPU $DEVICE_NUM"
+    cmd="CUDA_VISIBLE_DEVICES=$DEVICE_NUM python convert.py $INPUT_FOLDER $OUTPUT_FOLDER --num_chunks $NUM_DEVICES --chunk_idx $DEVICE_NUM --workers $NUM_WORKERS"
+    [[ -n "$METADATA_FILE" ]] && cmd="$cmd --metadata_file $METADATA_FILE"
+    [[ -n "$MIN_LENGTH" ]] && cmd="$cmd --min_length $MIN_LENGTH"
+    eval $cmd &
+
+    sleep 5
+done
+
+# Wait for all background processes to finish
+wait

convert.py

@@ -0,0 +1,124 @@
+import argparse
+import os
+from typing import Dict, Optional
+
+import ray
+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.settings import settings
+from marker.logger import configure_logging
+import traceback
+import json
+
+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, model_refs, metadata: Optional[Dict] = None, min_length: Optional[int] = 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"
+    if os.path.exists(out_filename):
+        return
+    try:
+        # Skip trying to convert files that don't have a lot of embedded text
+        # This can indicate that they were scanned, and not OCRed properly
+        # Usually these files are not recent/high-quality
+        if min_length:
+            length = get_length_of_text(fname)
+            if length < min_length:
+                return
+
+        full_text, out_metadata = convert_single_pdf(fname, model_refs, metadata=metadata)
+        if len(full_text.strip()) > 0:
+            with open(out_filename, "w+", encoding='utf-8') as f:
+                f.write(full_text)
+            with open(out_meta_filename, "w+") as f:
+                f.write(json.dumps(out_metadata, indent=4))
+        else:
+            print(f"Empty file: {fname}.  Could not convert.")
+    except Exception as e:
+        print(f"Error converting {fname}: {e}")
+        print(traceback.format_exc())
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Convert multiple pdfs to markdown.")
+    parser.add_argument("in_folder", help="Input folder with pdfs.")
+    parser.add_argument("out_folder", help="Output folder")
+    parser.add_argument("--chunk_idx", type=int, default=0, help="Chunk index to convert")
+    parser.add_argument("--num_chunks", type=int, default=1, help="Number of chunks being processed in parallel")
+    parser.add_argument("--max", type=int, default=None, help="Maximum number of pdfs to convert")
+    parser.add_argument("--workers", type=int, default=5, help="Number of worker processes to use")
+    parser.add_argument("--metadata_file", type=str, default=None, help="Metadata json file to use for filtering")
+    parser.add_argument("--min_length", type=int, default=None, help="Minimum length of pdf to convert")
+
+    args = parser.parse_args()
+
+    in_folder = os.path.abspath(args.in_folder)
+    out_folder = os.path.abspath(args.out_folder)
+    files = [os.path.join(in_folder, f) for f in os.listdir(in_folder)]
+    os.makedirs(out_folder, exist_ok=True)
+
+    # Handle chunks if we're processing in parallel
+    # Ensure we get all files into a chunk
+    chunk_size = math.ceil(len(files) / args.num_chunks)
+    start_idx = args.chunk_idx * chunk_size
+    end_idx = start_idx + chunk_size
+    files_to_convert = files[start_idx:end_idx]
+
+    # Limit files converted if needed
+    if args.max:
+        files_to_convert = files_to_convert[:args.max]
+
+    metadata = {}
+    if args.metadata_file:
+        metadata_file = os.path.abspath(args.metadata_file)
+        with open(metadata_file, "r") as f:
+            metadata = json.load(f)
+
+    total_processes = min(len(files_to_convert), args.workers)
+
+    ray.init(
+        num_cpus=total_processes,
+        num_gpus=1 if settings.CUDA else 0,
+        storage=settings.RAY_CACHE_PATH,
+        _temp_dir=settings.RAY_CACHE_PATH,
+        dashboard_host=settings.RAY_DASHBOARD_HOST,
+        log_to_driver=settings.DEBUG
+    )
+
+    model_lst = load_all_models()
+    model_refs = ray.put(model_lst)
+
+    # Dynamically set GPU allocation per task based on GPU ram
+    gpu_frac = settings.VRAM_PER_TASK / settings.INFERENCE_RAM if settings.CUDA else 0
+
+    print(f"Converting {len(files_to_convert)} pdfs in chunk {args.chunk_idx + 1}/{args.num_chunks} with {total_processes} processes, and storing in {out_folder}")
+    futures = [
+        process_single_pdf.options(num_gpus=gpu_frac).remote(
+            filename,
+            out_folder,
+            model_refs,
+            metadata=metadata.get(os.path.basename(filename)),
+            min_length=args.min_length
+        ) for filename in files_to_convert
+    ]
+
+    # Run all ray conversion tasks
+    progress_bar = tqdm(total=len(futures))
+    while len(futures) > 0:
+        finished, futures = ray.wait(
+            futures, timeout=7.0
+        )
+        finished_lst = ray.get(finished)
+        if isinstance(finished_lst, list):
+            progress_bar.update(len(finished_lst))
+        else:
+            progress_bar.update(1)
+
+    # Shutdown ray to free resources
+    ray.shutdown()

convert_single.py

@@ -0,0 +1,29 @@
+import argparse
+
+from marker.convert import convert_single_pdf
+from marker.logger import configure_logging
+from marker.models import load_all_models
+from marker.settings import settings
+import json
+
+configure_logging()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("filename", help="PDF file to parse")
+    parser.add_argument("output", help="Output file name")
+    parser.add_argument("--max_pages", type=int, default=None, help="Maximum number of pages to parse")
+    parser.add_argument("--parallel_factor", type=int, default=1, help="How much to multiply default parallel OCR workers and model batch sizes by.")
+    args = parser.parse_args()
+
+    fname = args.filename
+    model_lst = load_all_models()
+    full_text, out_meta = convert_single_pdf(fname, model_lst, max_pages=args.max_pages, parallel_factor=args.parallel_factor)
+
+    with open(args.output, "w+", encoding='utf-8') as f:
+        f.write(full_text)
+
+    out_meta_filename = args.output.rsplit(".", 1)[0] + "_meta.json"
+    with open(out_meta_filename, "w+") as f:
+        f.write(json.dumps(out_meta, indent=4))

marker/bbox.py

@@ -0,0 +1,81 @@
+import fitz as pymupdf
+
+def should_merge_blocks(box1, box2, tol=5):
+    # Within tol y px, and to the right within tol px
+    merge = [
+        box2[0] > box1[0], # After in the x coordinate
+        abs(box2[1] - box1[1]) < tol, # Within tol y px
+        abs(box2[3] - box1[3]) < tol, # Within tol y px
+        abs(box2[0] - box1[2]) < tol, # Within tol x px
+    ]
+    return all(merge)
+
+
+def merge_boxes(box1, box2):
+    return (min(box1[0], box2[0]), min(box1[1], box2[1]), max(box2[2], box1[2]), max(box1[3], box2[3]))
+
+
+def boxes_intersect(box1, box2):
+    # Box1 intersects box2
+    return box1[0] < box2[2] and box1[2] > box2[0] and box1[1] < box2[3] and box1[3] > box2[1]
+
+
+def boxes_intersect_pct(box1, box2, pct=.9):
+    # determine the coordinates of the intersection rectangle
+    x_left = max(box1[0], box2[0])
+    y_top = max(box1[1], box2[1])
+    x_right = min(box1[2], box2[2])
+    y_bottom = min(box1[3], box2[3])
+
+    if x_right < x_left or y_bottom < y_top:
+        return 0.0
+
+    # The intersection of two axis-aligned bounding boxes is always an
+    # axis-aligned bounding box
+    intersection_area = (x_right - x_left) * (y_bottom - y_top)
+
+    # compute the area of both AABBs
+    bb1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
+    bb2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
+
+    iou = intersection_area / float(bb1_area + bb2_area - intersection_area)
+    return iou > pct
+
+
+def multiple_boxes_intersect(box1, boxes):
+    for box2 in boxes:
+        if boxes_intersect(box1, box2):
+            return True
+    return False
+
+
+def box_contained(box1, box2):
+    # Box1 inside box2
+    return box1[0] > box2[0] and box1[1] > box2[1] and box1[2] < box2[2] and box1[3] < box2[3]
+
+
+def unnormalize_box(bbox, width, height):
+    return [
+        width * (bbox[0] / 1000),
+        height * (bbox[1] / 1000),
+        width * (bbox[2] / 1000),
+        height * (bbox[3] / 1000),
+    ]
+
+
+def correct_rotation(bbox, page):
+    #bbox base is (x0, y0, x1, y1)
+    rotation = page.rotation
+    if rotation == 0:
+        return bbox
+
+    tl = pymupdf.Point(bbox[0], bbox[1]) * page.rotation_matrix
+    br = pymupdf.Point(bbox[2], bbox[3]) * page.rotation_matrix
+    if rotation == 90:
+        bbox = [br[0], tl[1], tl[0], br[1]]
+    elif rotation == 180:
+        bbox = [br[0], br[1], tl[0], tl[1]]
+    elif rotation == 270:
+        bbox = [tl[0], br[1], br[0], tl[1]]
+
+    return bbox

marker/benchmark/scoring.py

@@ -0,0 +1,51 @@
+import math
+
+from rapidfuzz import fuzz, distance
+import re
+
+CHUNK_MIN_CHARS = 25
+
+
+def tokenize(text):
+    # Combined pattern
+    pattern = r'([^\w\s\d\'])|([\w\']+)|(\d+)|(\n+)|( +)'
+    result = re.findall(pattern, text)
+    # Flatten the result and filter out empty strings
+    flattened_result = [item for sublist in result for item in sublist if item]
+    return flattened_result
+
+
+def chunk_text(text):
+    chunks = text.split("\n")
+    chunks = [c for c in chunks if c.strip() and len(c) > CHUNK_MIN_CHARS]
+    return chunks
+
+
+def overlap_score(hypothesis_chunks, reference_chunks):
+    length_modifier = len(hypothesis_chunks) / len(reference_chunks)
+    search_distance = max(len(reference_chunks) // 5, 10)
+    chunk_scores = []
+    chunk_weights = []
+    for i, hyp_chunk in enumerate(hypothesis_chunks):
+        max_score = 0
+        chunk_weight = 1
+        i_offset = int(i * length_modifier)
+        chunk_range = range(max(0, i_offset-search_distance), min(len(reference_chunks), i_offset+search_distance))
+        for j in chunk_range:
+            ref_chunk = reference_chunks[j]
+            score = fuzz.ratio(hyp_chunk, ref_chunk, score_cutoff=30) / 100
+            if score > max_score:
+                max_score = score
+                chunk_weight = math.sqrt(len(ref_chunk))
+        chunk_scores.append(max_score)
+        chunk_weights.append(chunk_weight)
+    chunk_scores = [chunk_scores[i] * chunk_weights[i] for i in range(len(chunk_scores))]
+    return chunk_scores, chunk_weights
+
+
+def score_text(hypothesis, reference):
+    # Returns a 0-1 alignment score
+    hypothesis_chunks = chunk_text(hypothesis)
+    reference_chunks = chunk_text(reference)
+    chunk_scores, chunk_weights = overlap_score(hypothesis_chunks, reference_chunks)
+    return sum(chunk_scores) / sum(chunk_weights)

marker/cleaners/bullets.py

@@ -0,0 +1,8 @@
+import re
+
+
+def replace_bullets(text):
+    # Replace bullet characters with a -
+    bullet_pattern = r"(^|[\n ])[•●○■▪▫–—]( )"
+    replaced_string = re.sub(bullet_pattern, r"\1-\2", text)
+    return replaced_string

marker/cleaners/code.py

@@ -0,0 +1,129 @@
+from marker.schema import Span, Line, Page
+import re
+from typing import List
+import fitz as pymupdf
+
+
+def is_code_linelen(lines, thresh=60):
+    # Decide based on chars per newline threshold
+    total_alnum_chars = sum(len(re.findall(r'\w', line.prelim_text)) for line in lines)
+    total_newlines = max(len(lines) - 1, 1)
+
+    if total_alnum_chars == 0:
+        return False
+
+    ratio = total_alnum_chars / total_newlines
+    return ratio < thresh
+
+
+def comment_count(lines):
+    pattern = re.compile(r"^(//|#|'|--|/\*|'''|\"\"\"|--\[\[|<!--|%|%{|\(\*)")
+    return sum([1 for line in lines if pattern.match(line)])
+
+
+def identify_code_blocks(blocks: List[Page]):
+    code_block_count = 0
+    font_info = None
+    for p in blocks:
+        stats = p.get_font_stats()
+        if font_info is None:
+            font_info = stats
+        else:
+            font_info += stats
+    try:
+        most_common_font = font_info.most_common(1)[0][0]
+    except IndexError:
+        print(f"Could not find most common font")
+        most_common_font = None
+
+    last_block = None
+    for page in blocks:
+        try:
+            min_start = page.get_min_line_start()
+        except IndexError:
+            continue
+
+        for block in page.blocks:
+            if block.most_common_block_type() != "Text":
+                last_block = block
+                continue
+
+            is_indent = []
+            line_fonts = []
+            for line in block.lines:
+                fonts = [span.font for span in line.spans]
+                line_fonts += fonts
+                line_start = line.bbox[0]
+                if line_start > min_start:
+                    is_indent.append(True)
+                else:
+                    is_indent.append(False)
+            comment_lines = comment_count([line.prelim_text for line in block.lines])
+            is_code = [
+                len(block.lines) > 3,
+                sum([f != most_common_font for f in line_fonts]) > len(line_fonts) * .8,  # At least 80% of the fonts are not the most common, since code usually uses a different font from the main body text
+                is_code_linelen(block.lines),
+                (
+                    sum(is_indent) > len(block.lines) * .2
+                    or
+                    comment_lines > len(block.lines) * .2
+                 ), # 20% lines indented or 20% of the lines are comments
+            ]
+
+            # Check if previous block is code, and this block is indented
+            is_code_prev = [
+                last_block and last_block.most_common_block_type() == "Code",
+                sum(is_indent) >= len(block.lines) * .8 # At least 80% indented
+            ]
+
+            if all(is_code) or all(is_code_prev):
+                code_block_count += 1
+                block.set_block_type("Code")
+
+            last_block = block
+    return code_block_count
+
+
+def indent_blocks(blocks: List[Page]):
+    span_counter = 0
+    for page in blocks:
+        for block in page.blocks:
+            block_types = [span.block_type for line in block.lines for span in line.spans]
+            if "Code" not in block_types:
+                continue
+
+            lines = []
+            min_left = 1000  # will contain x- coord of column 0
+            col_width = 0  # width of 1 char
+            for line in block.lines:
+                text = ""
+                min_left = min(line.bbox[0], min_left)
+                for span in line.spans:
+                    if col_width == 0 and len(span.text) > 0:
+                        col_width = (span.bbox[2] - span.bbox[0]) / len(span.text)
+                    text += span.text
+                lines.append((pymupdf.Rect(line.bbox), text))
+
+            block_text = ""
+            blank_line = False
+            for line in lines:
+                text = line[1]
+                prefix = " " * int((line[0].x0 - min_left) / col_width)
+                current_line_blank = len(text.strip()) == 0
+                if blank_line and current_line_blank:
+                    # Don't put multiple blank lines in a row
+                    continue
+
+                block_text += prefix + text + "\n"
+                blank_line = current_line_blank
+
+            new_span = Span(
+                text=block_text,
+                bbox=block.bbox,
+                color=block.lines[0].spans[0].color,
+                span_id=f"{span_counter}_fix_code",
+                font=block.lines[0].spans[0].font,
+                block_type="Code"
+            )
+            span_counter += 1
+            block.lines = [Line(spans=[new_span], bbox=block.bbox)]