yujunjun/olmocr
1
0
Fork 0
mirror of https://github.com/allenai/olmocr.git synced 2025-06-27 04:00:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

Code cleanup, version bump, remove unused permutation test

Browse Source
This commit is contained in:
Jake Poznanski 2025-05-16 21:25:32 +00:00
parent 5de52e7d13
commit 63aee2c1e5
14 changed files with 72 additions and 266 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

167
olmocr/bench/benchmark.py
View File

@ -28,7 +28,7 @@ from tqdm import tqdm
from .report import generate_html_report
from .tests import BaselineTest, BasePDFTest, load_tests, save_tests
from .utils import calculate_bootstrap_ci, perform_permutation_test
from .utils import calculate_bootstrap_ci
def evaluate_candidate(
@ -186,16 +186,6 @@ def main():
default=0.95,
help="Confidence level for interval calculation (default: 0.95 for 95% CI).",
)
parser.add_argument(
"--permutation_tests",
nargs="?",
const="default",
help=(
"Run permutation testing. If provided without candidate names, run default tests. "
"If provided with a comma-separated list of candidate names (e.g. --permutation_tests asdf,qwe,ert), "
"run permutation tests on all pairs of the specified candidates."
),
)
# New arguments
parser.add_argument("--sample", type=int, default=None, help="Randomly sample N tests to run instead of all tests.")
parser.add_argument("--test_report", type=str, default=None, help="Generate an HTML report of test results. Provide a filename (e.g., results.html).")
@ -418,161 +408,6 @@ def main():
print(f" {jsonl_file:30s}: {pass_rate:0.1f}% ({results['passed']}/{results['total']} tests)")
print("")
if args.permutation_tests is not None:
print("\n" + "=" * 60)
print("Pairwise Permutation Tests:")
valid_candidates = [c for c in summary if not c[3]]
if args.permutation_tests == "default":
olmocr_candidates = sorted([c for c in valid_candidates if "olmocr" in c[0].lower()], key=lambda x: x[1], reverse=True)
non_olmocr_candidates = sorted([c for c in valid_candidates if "olmocr" not in c[0].lower()], key=lambda x: x[1], reverse=True)
top_olmocr = olmocr_candidates[0] if olmocr_candidates else None
top_non_olmocr = non_olmocr_candidates[0] if non_olmocr_candidates else None
top_two_olmocr = olmocr_candidates[:2]
if top_olmocr and top_non_olmocr:
olmocr_name, olmocr_score = top_olmocr[0], top_olmocr[1]
non_olmocr_name, non_olmocr_score = top_non_olmocr[0], top_non_olmocr[1]
# Extract file sizes and scores for both candidates
olmocr_jsonl_sizes = []
non_olmocr_jsonl_sizes = []
# Extract jsonl file sizes for each candidate
for test in all_tests:
jsonl_file = test_to_jsonl.get(test.id, "unknown")
# Process for top_olmocr
if not top_olmocr[3] and hasattr(test, "pdf") and hasattr(test, "page"):
pdf_name = test.pdf
page = test.page
if pdf_name in test_results_by_candidate.get(top_olmocr[0], {}) and page in test_results_by_candidate[top_olmocr[0]].get(pdf_name, {}):
for t, _, _ in test_results_by_candidate[top_olmocr[0]][pdf_name][page]:
if t.id == test.id:
if jsonl_file not in olmocr_jsonl_sizes:
olmocr_jsonl_sizes.append(len([t for t in all_tests if test_to_jsonl.get(t.id, "") == jsonl_file]))
break
# Process for top_non_olmocr
if not top_non_olmocr[3] and hasattr(test, "pdf") and hasattr(test, "page"):
pdf_name = test.pdf
page = test.page
if pdf_name in test_results_by_candidate.get(top_non_olmocr[0], {}) and page in test_results_by_candidate[top_non_olmocr[0]].get(pdf_name, {}):
for t, _, _ in test_results_by_candidate[top_non_olmocr[0]][pdf_name][page]:
if t.id == test.id:
if jsonl_file not in non_olmocr_jsonl_sizes:
non_olmocr_jsonl_sizes.append(len([t for t in all_tests if test_to_jsonl.get(t.id, "") == jsonl_file]))
break
diff, p_value = perform_permutation_test(
top_olmocr[7], top_non_olmocr[7],
splits_a=olmocr_jsonl_sizes if olmocr_jsonl_sizes else None,
splits_b=non_olmocr_jsonl_sizes if non_olmocr_jsonl_sizes else None
)
print("\nComparison 1: Top olmocr vs Top non-olmocr candidate")
print(f" {olmocr_name} ({olmocr_score*100:.1f}%) vs {non_olmocr_name} ({non_olmocr_score*100:.1f}%)")
print(f" Difference: {diff*100:.2f}% (positive means {olmocr_name} is better)")
print(f" p-value: {p_value:.4f}")
if p_value < 0.05:
print(" Result: Statistically significant difference (p < 0.05)")
else:
print(" Result: No statistically significant difference (p ≥ 0.05)")
else:
print("\nCannot perform olmocr vs non-olmocr comparison: Missing candidates")
if len(top_two_olmocr) >= 2:
# Extract file sizes for each candidate
olmocr1_jsonl_sizes = []
olmocr2_jsonl_sizes = []
# Extract jsonl file sizes for each candidate
for test in all_tests:
jsonl_file = test_to_jsonl.get(test.id, "unknown")
# Process for first olmocr candidate
if not top_two_olmocr[0][3] and hasattr(test, "pdf") and hasattr(test, "page"):
pdf_name = test.pdf
page = test.page
if pdf_name in test_results_by_candidate.get(top_two_olmocr[0][0], {}) and page in test_results_by_candidate[top_two_olmocr[0][0]].get(pdf_name, {}):
for t, _, _ in test_results_by_candidate[top_two_olmocr[0][0]][pdf_name][page]:
if t.id == test.id:
if jsonl_file not in olmocr1_jsonl_sizes:
olmocr1_jsonl_sizes.append(len([t for t in all_tests if test_to_jsonl.get(t.id, "") == jsonl_file]))
break
# Process for second olmocr candidate
if not top_two_olmocr[1][3] and hasattr(test, "pdf") and hasattr(test, "page"):
pdf_name = test.pdf
page = test.page
if pdf_name in test_results_by_candidate.get(top_two_olmocr[1][0], {}) and page in test_results_by_candidate[top_two_olmocr[1][0]].get(pdf_name, {}):
for t, _, _ in test_results_by_candidate[top_two_olmocr[1][0]][pdf_name][page]:
if t.id == test.id:
if jsonl_file not in olmocr2_jsonl_sizes:
olmocr2_jsonl_sizes.append(len([t for t in all_tests if test_to_jsonl.get(t.id, "") == jsonl_file]))
break
diff, p_value = perform_permutation_test(
top_two_olmocr[0][7], top_two_olmocr[1][7],
splits_a=olmocr1_jsonl_sizes if olmocr1_jsonl_sizes else None,
splits_b=olmocr2_jsonl_sizes if olmocr2_jsonl_sizes else None
)
print("\nComparison 2: Top two olmocr candidates")
print(f" {top_two_olmocr[0][0]} ({top_two_olmocr[0][1]*100:.1f}%) vs {top_two_olmocr[1][0]} ({top_two_olmocr[1][1]*100:.1f}%)")
print(f" Difference: {diff*100:.2f}% (positive means {top_two_olmocr[0][0]} is better)")
print(f" p-value: {p_value:.4f}")
if p_value < 0.05:
print(" Result: Statistically significant difference (p < 0.05)")
else:
print(" Result: No statistically significant difference (p ≥ 0.05)")
else:
print("\nCannot perform top two olmocr comparison: Not enough olmocr candidates")
else:
candidate_names = [name.strip() for name in args.permutation_tests.split(",")]
selected_candidates = [c for c in valid_candidates if c[0] in candidate_names]
if len(selected_candidates) < 2:
print("\nNot enough valid candidates among the selected for permutation tests.")
else:
for cand1, cand2 in combinations(selected_candidates, 2):
# Extract file sizes for each candidate
cand1_jsonl_sizes = []
cand2_jsonl_sizes = []
# Extract jsonl file sizes for each candidate
for test in all_tests:
jsonl_file = test_to_jsonl.get(test.id, "unknown")
# Process for first candidate
if not cand1[3] and hasattr(test, "pdf") and hasattr(test, "page"):
pdf_name = test.pdf
page = test.page
if pdf_name in test_results_by_candidate.get(cand1[0], {}) and page in test_results_by_candidate[cand1[0]].get(pdf_name, {}):
for t, _, _ in test_results_by_candidate[cand1[0]][pdf_name][page]:
if t.id == test.id:
if jsonl_file not in cand1_jsonl_sizes:
cand1_jsonl_sizes.append(len([t for t in all_tests if test_to_jsonl.get(t.id, "") == jsonl_file]))
break
# Process for second candidate
if not cand2[3] and hasattr(test, "pdf") and hasattr(test, "page"):
pdf_name = test.pdf
page = test.page
if pdf_name in test_results_by_candidate.get(cand2[0], {}) and page in test_results_by_candidate[cand2[0]].get(pdf_name, {}):
for t, _, _ in test_results_by_candidate[cand2[0]][pdf_name][page]:
if t.id == test.id:
if jsonl_file not in cand2_jsonl_sizes:
cand2_jsonl_sizes.append(len([t for t in all_tests if test_to_jsonl.get(t.id, "") == jsonl_file]))
break
diff, p_value = perform_permutation_test(
cand1[7], cand2[7],
splits_a=cand1_jsonl_sizes if cand1_jsonl_sizes else None,
splits_b=cand2_jsonl_sizes if cand2_jsonl_sizes else None
)
print(f"\nComparison: {cand1[0]} vs {cand2[0]}")
print(f" {cand1[0]} ({cand1[1]*100:.1f}%) vs {cand2[0]} ({cand2[1]*100:.1f}%)")
print(f" Difference: {diff*100:.2f}% (positive means {cand1[0]} is better)")
print(f" p-value: {p_value:.4f}")
if p_value < 0.05:
print(" Result: Statistically significant difference (p < 0.05)")
else:
print(" Result: No statistically significant difference (p ≥ 0.05)")
print("=" * 60)
# Generate HTML report if requested
if args.test_report:
generate_html_report(test_results_by_candidate, pdf_folder, args.test_report)

5
olmocr/bench/runners/run_claude.py
View File

@ -2,10 +2,7 @@ import json
import os
from anthropic import Anthropic
from prompts import (
build_openai_silver_data_prompt,
claude_response_format_schema,
)
from prompts import build_openai_silver_data_prompt, claude_response_format_schema
from olmocr.data.renderpdf import render_pdf_to_base64png
from olmocr.prompts.anchor import get_anchor_text

6
olmocr/bench/synth/mine_html_templates.py
View File

@ -17,11 +17,7 @@ from playwright.async_api import async_playwright
from syntok.segmenter import process
from tqdm import tqdm
from olmocr.bench.tests import (
TableTest,
TestType,
parse_html_tables,
)
from olmocr.bench.tests import TableTest, TestType, parse_html_tables
from olmocr.data.renderpdf import (
get_png_dimensions_from_base64,
render_pdf_to_base64png,

5
olmocr/bench/utils.py
View File

@ -67,8 +67,9 @@ def calculate_bootstrap_ci(test_scores: List[float], n_bootstrap: int = 1000, ci
return (lower_bound, upper_bound)
def perform_permutation_test(scores_a: List[float], scores_b: List[float], n_permutations: int = 10000,
splits_a: List[int] = None, splits_b: List[int] = None) -> Tuple[float, float]:
def perform_permutation_test(
scores_a: List[float], scores_b: List[float], n_permutations: int = 10000, splits_a: List[int] = None, splits_b: List[int] = None
) -> Tuple[float, float]:
"""
Perform a permutation test to determine if there's a significant difference
between two sets of test scores.

1
olmocr/eval/dolma_refine/metrics.py
View File

@ -208,7 +208,6 @@ class ParagraphEditSimilarity(DocumentEditSimilarity):
self.sent_window = sent_window
def segment(self, seq_a_tokens: list[str], seq_b_tokens: list[str]) -> list[tuple[list[str], list[str]]]:
all_spans = []
for seq_tokens in (seq_a_tokens, seq_b_tokens):

9
olmocr/eval/dolma_refine/registry.py
View File

@ -91,15 +91,18 @@ class BaseRegistry(Generic[T]):
@overload
@classmethod
def get(cls, name: str) -> T: ...
def get(cls, name: str) -> T:
...
@overload
@classmethod
def get(cls, name: str, raise_on_missing: Literal[True]) -> T: ...
def get(cls, name: str, raise_on_missing: Literal[True]) -> T:
...
@overload
@classmethod
def get(cls, name: str, raise_on_missing: Literal[False]) -> Optional[T]: ...
def get(cls, name: str, raise_on_missing: Literal[False]) -> Optional[T]:
...
@classmethod
def get(cls, name: str, raise_on_missing: bool = True) -> Optional[T]:

3
olmocr/train/core/errors.py
View File

@ -1 +1,2 @@
class DolmaRefineError(RuntimeError): ...
class DolmaRefineError(RuntimeError):
...

2
olmocr/version.py
View File

@ -2,7 +2,7 @@ _MAJOR = "0"
_MINOR = "1"
# On main and in a nightly release the patch should be one ahead of the last
# released build.
_PATCH = "67"
_PATCH = "68"
# This is mainly for nightly builds which have the suffix ".dev$DATE". See
# https://semver.org/#is-v123-a-semantic-version for the semantics.
_SUFFIX = ""

1
scripts/jsonl_to_markdown.py
View File

@ -46,7 +46,6 @@ def jsonl_to_markdown(input_file, output_dir):
# It takes two arguments: the input JSONL file and the output directory.
# The script will create the output directory if it does not exist.
if __name__ == "__main__":
if len(sys.argv) != 3:
print("Usage: python jsonl_to_markdown.py <input_file> <output_dir>")
sys.exit(1)

63
scripts/pareto_plot.py
View File

@ -71,7 +71,7 @@ data = {
"Qwen 2 VL (A100)",
"Qwen 2 VL (H100,L40S)",
"Qwen 2.5 VL (A100)",
"Qwen 2.5 VL (H100,L40S)"
"Qwen 2.5 VL (H100,L40S)",
],
COST_COLUMN_NAME: [
12480,
@ -86,7 +86,7 @@ data = {
270, # Same cost as Ours
190, # Same cost as Ours
270, # Same cost as Ours
190 # Same cost as Ours
190, # Same cost as Ours
],
PERF_COLUMN_NAME: [
69.9, # GPT-4o (Anchored)
@ -101,8 +101,8 @@ data = {
31.5, # Qwen2VL
31.5, # Qwen2VL
65.5, # Qwen2.5VL
65.5 # Qwen2.5VL
]
65.5, # Qwen2.5VL
],
}
df = pd.DataFrame(data)
@ -121,41 +121,23 @@ model_categories = {
"Qwen 2 VL (A100)": "Open VLM",
"Qwen 2 VL (H100,L40S)": "Open VLM",
"Qwen 2.5 VL (A100)": "Open VLM",
"Qwen 2.5 VL (H100,L40S)": "Open VLM"
"Qwen 2.5 VL (H100,L40S)": "Open VLM",
}
df[CATEGORY_COLUMN_NAME] = df[MODEL_COLUMN_NAME].map(model_categories)
# Category colors
category_colors = {
"Commercial API Tool": DARK_BLUE,
"Commercial VLM": DARK_GREEN,
"Open Source Tool": LIGHT_GREEN,
"Ours": DARK_PINK,
"Open VLM": PURPLE
}
category_colors = {"Commercial API Tool": DARK_BLUE, "Commercial VLM": DARK_GREEN, "Open Source Tool": LIGHT_GREEN, "Ours": DARK_PINK, "Open VLM": PURPLE}
df[COLOR_COLUMN_NAME] = df[CATEGORY_COLUMN_NAME].map(category_colors)
# Define marker types
category_markers = {
"Commercial API Tool": "o",
"Commercial VLM": "D",
"Open Source Tool": "s",
"Ours": "*",
"Open VLM": "^"
}
category_markers = {"Commercial API Tool": "o", "Commercial VLM": "D", "Open Source Tool": "s", "Ours": "*", "Open VLM": "^"}
df[MARKER_COLUMN_NAME] = df[CATEGORY_COLUMN_NAME].map(category_markers)
# Define marker sizes - increased sizes
category_marker_sizes = {
"Commercial API Tool": 120,
"Commercial VLM": 120,
"Open Source Tool": 140,
"Ours": 300,
"Open VLM": 140
}
category_marker_sizes = {"Commercial API Tool": 120, "Commercial VLM": 120, "Open Source Tool": 140, "Ours": 300, "Open VLM": 140}
# Define text colors
category_text_colors = {
@ -163,7 +145,7 @@ category_text_colors = {
"Commercial VLM": DARK_GREEN,
"Open Source Tool": DARK_TEAL,
"Ours": "#a51c5c", # darker pink
"Open VLM": "#6f1188" # darker purple
"Open VLM": "#6f1188", # darker purple
}
# Label offsets for better readability
@ -180,7 +162,7 @@ model_label_offsets = {
"Qwen 2 VL (A100)": [-20, 10],
"Qwen 2 VL (H100,L40S)": [-60, 25],
"Qwen 2.5 VL (A100)": [-20, 10],
"Qwen 2.5 VL (H100,L40S)": [-60, 25]
"Qwen 2.5 VL (H100,L40S)": [-60, 25],
}
df[OFFSET_COLUMN_NAME] = df[MODEL_COLUMN_NAME].map(model_label_offsets)
@ -220,16 +202,22 @@ for idx, row in df.iterrows():
# Set up axes
plt.ylim(25, 85) # Set y-axis limits from 25 to 85 to include Qwen2VL
plt.xlim(100, 15000)
plt.xscale('log') # Use log scale for cost
plt.xscale("log") # Use log scale for cost
plt.grid(True, which="both", ls=":", color=TEAL, alpha=0.2)
# Format y-axis to show percentages without scientific notation
def percent_formatter(y, pos):
return f'{y:.1f}%'
return f"{y:.1f}%"
plt.gca().yaxis.set_major_formatter(ticker.FuncFormatter(percent_formatter))
# Format x-axis to show dollar amounts
def dollar_formatter(x, pos):
return f'${x:,.0f}'
return f"${x:,.0f}"
# Set specific x-axis ticks with increased font size
plt.gca().xaxis.set_major_formatter(ticker.FuncFormatter(dollar_formatter))
@ -243,8 +231,8 @@ plt.ylabel("Overall Performance (Pass Rate %)", fontsize=16, weight="medium")
# plt.title("OCR Engines: Performance vs. Cost", fontsize=12, weight="medium")
# Remove spines
plt.gca().spines['top'].set_visible(False)
plt.gca().spines['right'].set_visible(False)
plt.gca().spines["top"].set_visible(False)
plt.gca().spines["right"].set_visible(False)
# Add the legend with custom ordering and increased font size
handles, labels = plt.gca().get_legend_handles_labels()
@ -254,14 +242,7 @@ ordered_handles = [label_to_handle[label] for label in desired_order if label in
ordered_labels = [label for label in desired_order if label in labels]
plt.legend(
ordered_handles,
ordered_labels,
loc="lower right",
fontsize=12, # Increased from 10
frameon=True,
framealpha=0.9,
edgecolor=TEAL,
facecolor="white"
ordered_handles, ordered_labels, loc="lower right", fontsize=12, frameon=True, framealpha=0.9, edgecolor=TEAL, facecolor="white" # Increased from 10
)
# Adjust layout

4
scripts/rich_tagging_pipeline.py
View File

@ -26,9 +26,7 @@ import zstandard as zstd
from huggingface_hub import snapshot_download
from pydantic import BaseModel, Field, ValidationError
from olmocr.check import (
check_torch_gpu_available,
)
from olmocr.check import check_torch_gpu_available
from olmocr.metrics import MetricsKeeper
from olmocr.s3_utils import (
download_directory,

4
scripts/tagging_pipeline.py
View File

@ -27,9 +27,7 @@ import zstandard as zstd
from huggingface_hub import snapshot_download
from pydantic import BaseModel, Field, ValidationError
from olmocr.check import (
check_torch_gpu_available,
)
from olmocr.check import check_torch_gpu_available
from olmocr.metrics import MetricsKeeper
from olmocr.s3_utils import (
download_directory,

4
scripts/tagging_pipeline_v2.py
View File

@ -26,9 +26,7 @@ import zstandard as zstd
from huggingface_hub import snapshot_download
from pydantic import BaseModel, Field, ValidationError
from olmocr.check import (
check_torch_gpu_available,
)
from olmocr.check import check_torch_gpu_available
from olmocr.metrics import MetricsKeeper
from olmocr.s3_utils import (
download_directory,

2
tests/test_integration.py
View File

@ -19,7 +19,7 @@ class TestPipelineIntegration(unittest.TestCase):
print(self.data[-1])
def test_edgar(self) -> None:
self.assertTrue(any("King of England" in line["text"] for line in self.data))
self.assertTrue(any("King of the English" in line["text"] for line in self.data))
def test_ambig(self) -> None:
self.assertTrue(any("Apples and Bananas" in line["text"] for line in self.data))