wehub-resource-sync
917eedffcf
Main / Python 3.11 - Docs (push) Has been cancelled
Main / Python 3.11 - Build (push) Has been cancelled
Main / Python 3.11 - Lint (push) Has been cancelled
Main / Python 3.11 - Style (push) Has been cancelled
Main / Python 3.11 - Test (push) Has been cancelled
Main / GPU CI (push) Has been cancelled
Main / Release (push) Has been cancelled
Main / Build and Push Docker Images (push) Has been cancelled
1790 lines
66 KiB
Python
1790 lines
66 KiB
Python
import argparse
|
||
import base64
|
||
import json
|
||
import logging
|
||
import re
|
||
import shutil
|
||
from abc import ABC, abstractmethod
|
||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||
from dataclasses import dataclass, replace
|
||
from html.parser import HTMLParser
|
||
from io import BytesIO
|
||
from os import PathLike
|
||
from pathlib import Path
|
||
from typing import (
|
||
Any,
|
||
Dict,
|
||
List,
|
||
Optional,
|
||
Tuple,
|
||
)
|
||
|
||
import numpy as np
|
||
import torch
|
||
from PIL import Image
|
||
from pypdf import PdfReader
|
||
from torch.utils.data import Dataset
|
||
from tqdm import tqdm
|
||
|
||
from olmocr.data.renderpdf import render_pdf_to_base64png
|
||
from olmocr.prompts.anchor import get_anchor_text
|
||
from olmocr.prompts.prompts import (
|
||
PageResponse,
|
||
build_finetuning_prompt,
|
||
build_no_anchoring_v4_yaml_prompt,
|
||
)
|
||
from olmocr.train.front_matter import FrontMatterParser, Sample
|
||
|
||
# Configure logging
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def validate_pdf_pair(md_path: Path) -> Tuple[Optional[Dict[str, Path]], Optional[Tuple[Path, str]]]:
|
||
"""Validate a single markdown-PDF pair.
|
||
|
||
Args:
|
||
md_path: Path to the markdown file
|
||
|
||
Returns:
|
||
Tuple of (valid_sample, invalid_pdf_info)
|
||
- valid_sample: Dict with markdown_path and pdf_path if valid, None otherwise
|
||
- invalid_pdf_info: Tuple of (pdf_path, reason) if invalid, None otherwise
|
||
"""
|
||
# Look for PDF with same stem (filename without extension)
|
||
pdf_path = md_path.with_suffix(".pdf")
|
||
|
||
if pdf_path.exists() or pdf_path.is_symlink():
|
||
# Resolve symlink if it is one
|
||
if pdf_path.is_symlink():
|
||
pdf_path = pdf_path.resolve()
|
||
|
||
# Verify the resolved path exists
|
||
if pdf_path.exists():
|
||
# Validate PDF - check it loads and has exactly one page and that you can get document-anchoring from it
|
||
try:
|
||
reader = PdfReader(str(pdf_path))
|
||
num_pages = len(reader.pages)
|
||
|
||
if num_pages != 1:
|
||
return None, (pdf_path, f"Expected 1 page, found {num_pages}")
|
||
|
||
# Test that document anchoring works
|
||
from olmocr.prompts.anchor import get_anchor_text
|
||
|
||
get_anchor_text(pdf_path, page=1, pdf_engine="pdfreport", target_length=100)
|
||
|
||
return {"markdown_path": md_path, "pdf_path": pdf_path}, None
|
||
|
||
except Exception as e:
|
||
return None, (pdf_path, f"Failed to load: {str(e)}")
|
||
|
||
return None, None
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class PipelineStep(ABC):
|
||
"""Abstract base class for pipeline steps."""
|
||
|
||
@abstractmethod
|
||
def __call__(self, sample: Sample) -> Optional[Sample]:
|
||
"""Process a sample and return the modified sample, or None to skip this sample."""
|
||
...
|
||
|
||
|
||
class BaseMarkdownPDFDataset(Dataset):
|
||
"""Base dataset class that loads and verifies markdown-PDF pairs."""
|
||
|
||
def __init__(self, root_dir: str | PathLike, pipeline_steps: Optional[List[PipelineStep]] = None):
|
||
"""
|
||
Initialize the dataset by finding all markdown files with corresponding PDFs.
|
||
|
||
Args:
|
||
root_dir: Path to the root folder containing processed markdown and PDF files
|
||
pipeline_steps: Optional list of pipeline steps to apply to each sample
|
||
"""
|
||
self.root_dir = Path(root_dir)
|
||
self.pipeline_steps = pipeline_steps or []
|
||
self.samples = []
|
||
|
||
# Find all markdown files recursively
|
||
logger.info(f"Scanning for markdown files in {self.root_dir}...")
|
||
md_files = list(self.root_dir.rglob("*.md"))
|
||
|
||
# Verify each markdown file has a corresponding PDF using ProcessPoolExecutor
|
||
valid_count = 0
|
||
invalid_pdfs = []
|
||
|
||
logger.info(f"Validating {len(md_files)} markdown-PDF pairs using ProcessPoolExecutor...")
|
||
|
||
# Use ProcessPoolExecutor for parallel validation
|
||
with ProcessPoolExecutor(max_workers=8) as executor:
|
||
# Submit all validation tasks
|
||
future_to_md = {executor.submit(validate_pdf_pair, md_path): md_path for md_path in md_files}
|
||
|
||
# Process results as they complete
|
||
with tqdm(total=len(md_files), desc="Validating PDFs") as pbar:
|
||
for future in as_completed(future_to_md):
|
||
md_path = future_to_md[future]
|
||
try:
|
||
valid_sample, invalid_pdf_info = future.result()
|
||
|
||
if valid_sample:
|
||
self.samples.append(valid_sample)
|
||
valid_count += 1
|
||
elif invalid_pdf_info:
|
||
invalid_pdfs.append(invalid_pdf_info)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Error processing {md_path}: {str(e)}")
|
||
invalid_pdfs.append((md_path.with_suffix(".pdf"), f"Processing error: {str(e)}"))
|
||
|
||
pbar.update(1)
|
||
|
||
# Sort samples by markdown path for consistent ordering across runs
|
||
self.samples.sort(key=lambda x: x["markdown_path"])
|
||
|
||
logger.info(f"Found {valid_count} valid markdown-PDF pairs")
|
||
|
||
if invalid_pdfs:
|
||
logger.warning(f"{len(invalid_pdfs)} invalid PDFs found:")
|
||
for pdf_path, reason in invalid_pdfs[:5]: # Show first 5
|
||
logger.warning(f" - {pdf_path.name}: {reason}")
|
||
if len(invalid_pdfs) > 5:
|
||
logger.warning(f" ... and {len(invalid_pdfs) - 5} more")
|
||
|
||
def __len__(self) -> int:
|
||
return len(self.samples)
|
||
|
||
def __getitem__(self, idx: int) -> Optional[Dict[str, Any]]:
|
||
"""
|
||
Get a single sample from the dataset.
|
||
|
||
Returns:
|
||
dict containing at minimum:
|
||
- 'markdown_path': Path to the markdown file
|
||
- 'pdf_path': Path to the PDF file
|
||
|
||
Additional fields will be added by pipeline steps.
|
||
Returns None if any pipeline step returns None.
|
||
"""
|
||
# Start with basic sample info
|
||
sample = self.samples[idx].copy()
|
||
|
||
# Apply pipeline steps, returning None if any step returns None
|
||
for step in self.pipeline_steps:
|
||
sample = step(sample)
|
||
if sample is None:
|
||
return None
|
||
|
||
return sample
|
||
|
||
|
||
# FrontMatterParser is imported from olmocr.train.front_matter
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class PDFRenderer(PipelineStep):
|
||
"""Pipeline step that renders PDF to image."""
|
||
|
||
target_longest_image_dim: int
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
"""Render PDF to image."""
|
||
# Render PDF to image
|
||
base64_png = render_pdf_to_base64png(str(sample["pdf_path"]), page_num=1, target_longest_image_dim=self.target_longest_image_dim)
|
||
png_bytes = base64.b64decode(base64_png)
|
||
image = Image.open(BytesIO(png_bytes))
|
||
|
||
# Update sample
|
||
sample["image"] = image
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class StaticLengthDocumentAnchoring(PipelineStep):
|
||
target_anchor_text_len: int
|
||
|
||
"""Pipeline step that runs document anchoring on the PDF and puts in the data to be used by later prompting stages"""
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
anchor_text = get_anchor_text(sample["pdf_path"], page=1, pdf_engine="pdfreport", target_length=self.target_anchor_text_len)
|
||
sample["anchor_text"] = anchor_text
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class FinetuningPrompt(PipelineStep):
|
||
"""Applies the standard fine tuning prompt"""
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
sample["instruction_prompt"] = build_finetuning_prompt(sample["anchor_text"])
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class NewYamlFinetuningPromptWithAnchoring(PipelineStep):
|
||
"""Applies the standard fine tuning prompt"""
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
sample["instruction_prompt"] = (
|
||
f"Attached is one page of a document, as well as some raw textual content that was previously extracted for it. "
|
||
f"Just return the plain text representation of this document as if you were reading it naturally. Convert equations to LateX and tables to markdown.\n"
|
||
f"RAW_TEXT_START\n{sample['anchor_text']}\nRAW_TEXT_END\n"
|
||
f"Return your output as markdown, with a front matter section on top specifying values for the primary_language, is_rotation_valid, rotation_correction, is_table, and is_diagram parameters."
|
||
)
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class NewYamlFinetuningPromptWithNoAnchoring(PipelineStep):
|
||
"""Applies the standard fine tuning prompt"""
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
sample["instruction_prompt"] = build_no_anchoring_v4_yaml_prompt()
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class FrontMatterOutputFormat(PipelineStep):
|
||
"""Takes the output and applies the standard yaml formatting to it"""
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
page_data = sample["page_data"]
|
||
assert type(page_data) is PageResponse
|
||
|
||
sample["response"] = f"""---
|
||
primary_language: {page_data.primary_language}
|
||
is_rotation_valid: {page_data.is_rotation_valid}
|
||
rotation_correction: {page_data.rotation_correction}
|
||
is_table: {page_data.is_table}
|
||
is_diagram: {page_data.is_diagram}
|
||
---
|
||
{page_data.natural_text if page_data.natural_text is not None and len(page_data.natural_text.strip()) > 0 else ""}
|
||
""".strip()
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class JSONOutputFormat(PipelineStep):
|
||
"""Takes the output and applies the standard yaml formatting to it"""
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
page_data = sample["page_data"]
|
||
assert type(page_data) is PageResponse
|
||
|
||
sample["response"] = json.dumps(
|
||
{
|
||
"primary_language": page_data.primary_language,
|
||
"is_rotation_valid": page_data.is_rotation_valid,
|
||
"rotation_correction": page_data.rotation_correction,
|
||
"is_table": page_data.is_table,
|
||
"is_diagram": page_data.is_diagram,
|
||
"natural_text": page_data.natural_text,
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class LatexBracketNormalizer(PipelineStep):
|
||
"""Normalizes LaTeX brackets in natural text field."""
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
"""Normalize LaTeX brackets in the natural text field."""
|
||
# Get the page_data object
|
||
if "page_data" not in sample:
|
||
return sample
|
||
|
||
page_data = sample["page_data"]
|
||
if not hasattr(page_data, "natural_text") or not page_data.natural_text:
|
||
return sample
|
||
|
||
text = page_data.natural_text
|
||
|
||
# Define patterns for LaTeX normalization
|
||
# Order matters: process display math first, then inline
|
||
patterns = [
|
||
(r"\$\$(.+?)\$\$", r"\[\1\]"), # $$...$$ to \[...\]
|
||
(r"\$(.+?)\$", r"\(\1\)"), # $...$ to \(...\)
|
||
]
|
||
|
||
# Apply replacements
|
||
for pattern, replacement in patterns:
|
||
text = re.sub(pattern, replacement, text, flags=re.DOTALL)
|
||
|
||
# Update the page_data with normalized text
|
||
# Since PageResponse is frozen, we need to create a new instance
|
||
new_page_data = PageResponse(
|
||
primary_language=page_data.primary_language,
|
||
is_rotation_valid=page_data.is_rotation_valid,
|
||
rotation_correction=page_data.rotation_correction,
|
||
is_table=page_data.is_table,
|
||
is_diagram=page_data.is_diagram,
|
||
natural_text=text,
|
||
)
|
||
|
||
sample["page_data"] = new_page_data
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class RotationAugmentation(PipelineStep):
|
||
"""Pipeline step that randomly rotates images for augmentation."""
|
||
|
||
probability: float = 0.5 # Probability of applying rotation
|
||
|
||
def __call__(self, sample: Sample) -> Optional[Sample]:
|
||
"""Randomly rotate image and update rotation metadata."""
|
||
# Only proceed with given probability
|
||
if np.random.random() > self.probability:
|
||
return sample
|
||
|
||
# Check if image exists
|
||
if "image" not in sample:
|
||
return sample
|
||
|
||
# Check if page_data exists (we need to update it)
|
||
if "page_data" not in sample:
|
||
return sample
|
||
|
||
# Randomly choose a rotation (90, 180, or 270 degrees)
|
||
rotation_degrees = np.random.choice([90, 180, 270])
|
||
|
||
# Apply rotation to image
|
||
image = sample["image"]
|
||
if rotation_degrees == 90:
|
||
transpose = Image.Transpose.ROTATE_90
|
||
elif rotation_degrees == 180:
|
||
transpose = Image.Transpose.ROTATE_180
|
||
else: # 270
|
||
transpose = Image.Transpose.ROTATE_270
|
||
|
||
rotated_image = image.transpose(transpose)
|
||
sample["image"] = rotated_image
|
||
|
||
# Update page_data
|
||
page_data = sample["page_data"]
|
||
|
||
# Create new PageResponse with updated rotation info
|
||
# The rotation_correction should be the inverse of what we applied
|
||
# If we rotated 90 clockwise, we need 270 counter-clockwise to correct it
|
||
if rotation_degrees == 90:
|
||
correction = 270
|
||
elif rotation_degrees == 180:
|
||
correction = 180
|
||
else: # 270
|
||
correction = 90
|
||
|
||
new_page_data = PageResponse(
|
||
primary_language=page_data.primary_language,
|
||
is_rotation_valid=False, # Mark as invalid since we rotated it
|
||
rotation_correction=correction, # The correction needed to fix it
|
||
is_table=page_data.is_table,
|
||
is_diagram=page_data.is_diagram,
|
||
natural_text=page_data.natural_text,
|
||
)
|
||
|
||
sample["page_data"] = new_page_data
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class FilterOutRotatedDocuments(PipelineStep):
|
||
"""Pipeline step that filters out documents with rotation issues."""
|
||
|
||
def __call__(self, sample: Sample) -> Optional[Sample]:
|
||
"""Filter out samples where rotation is invalid or rotation correction is needed."""
|
||
# Check if page_data exists
|
||
if "page_data" not in sample:
|
||
return sample
|
||
|
||
page_data = sample["page_data"]
|
||
|
||
# Check if page_data has the required attributes
|
||
if not hasattr(page_data, "is_rotation_valid") or not hasattr(page_data, "rotation_correction"):
|
||
return sample
|
||
|
||
# Filter out if rotation is invalid or rotation correction is not 0
|
||
if page_data.is_rotation_valid is False or page_data.rotation_correction != 0:
|
||
return None
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class DatasetTextRuleFilter(PipelineStep):
|
||
"""Pipeline step that filters samples based on text content rules.
|
||
|
||
Filters out samples that:
|
||
- Contain markdown tables
|
||
- Contain malformed HTML tables
|
||
- Contain math equations that fail to render
|
||
- Contain mathematical symbols (∈, ∉, ⊂, ⊃, ⊆, ⊇, ∅, ∪, ∩, ∀, ∃, ¬) outside of table cells
|
||
- Contain LaTeX formatting commands (\\textit, \\textbf, \\texttt, etc.) outside of math equations
|
||
- Contain LaTeX table environments (\begin{table}, \begin{tabular}, etc.)
|
||
"""
|
||
|
||
def _contains_markdown_table(self, text: str) -> bool:
|
||
"""Check if text contains markdown tables."""
|
||
# Look for pipe-separated table patterns
|
||
# Markdown tables have lines like: | col1 | col2 | col3 |
|
||
# And separator lines like: |------|------|------|
|
||
lines = text.split("\n")
|
||
for i, line in enumerate(lines):
|
||
line = line.strip()
|
||
# Check if line looks like a table row
|
||
if line.startswith("|") and line.endswith("|") and line.count("|") >= 3:
|
||
# Check if next line is a separator (for header rows)
|
||
if i + 1 < len(lines):
|
||
next_line = lines[i + 1].strip()
|
||
if next_line.startswith("|") and "-" in next_line:
|
||
return True
|
||
# Check if previous line is a separator (for data rows)
|
||
if i > 0:
|
||
prev_line = lines[i - 1].strip()
|
||
if prev_line.startswith("|") and "-" in prev_line:
|
||
return True
|
||
return False
|
||
|
||
def _contains_math_symbols(self, text: str) -> bool:
|
||
"""Check if text contains specific mathematical symbols outside of table cells.
|
||
|
||
Returns:
|
||
True if text contains any of the specified math symbols outside tables
|
||
False otherwise
|
||
"""
|
||
# List of mathematical symbols to check for
|
||
math_symbols = [
|
||
# Set theory and logic
|
||
"∈",
|
||
"∉",
|
||
"⊂",
|
||
"⊃",
|
||
"⊆",
|
||
"⊇",
|
||
"∅",
|
||
"∪",
|
||
"∩",
|
||
"∀",
|
||
"∃",
|
||
"¬",
|
||
# Common mathematical operators
|
||
"⊕",
|
||
"⊗",
|
||
"⊙",
|
||
# Calculus and analysis
|
||
"∂",
|
||
"∇",
|
||
"∆",
|
||
"∫",
|
||
"∬",
|
||
"∭",
|
||
"∮",
|
||
"∏",
|
||
"∑",
|
||
"√",
|
||
"∛",
|
||
"∜",
|
||
# Arrows and relations
|
||
"⊥",
|
||
# Other common math symbols
|
||
"∠",
|
||
"∡",
|
||
"⊤",
|
||
"⊢",
|
||
"⊣",
|
||
"∴",
|
||
"∵",
|
||
"∶",
|
||
"∷",
|
||
"∝",
|
||
"≅",
|
||
"≆",
|
||
"≇",
|
||
"≊",
|
||
"≋",
|
||
# Matrix and vector notation
|
||
"⊕",
|
||
"⊖",
|
||
"⊗",
|
||
"⊘",
|
||
"⊙",
|
||
"⊚",
|
||
"⊛",
|
||
"⊜",
|
||
"⊝",
|
||
]
|
||
|
||
# First, remove all HTML tables from the text
|
||
text_without_tables = text
|
||
|
||
# Remove HTML tables
|
||
table_pattern = re.compile(r"<table\b[^>]*>.*?</table>", re.IGNORECASE | re.DOTALL)
|
||
text_without_tables = table_pattern.sub("", text_without_tables)
|
||
|
||
# Now check if any of these symbols appear in the text without tables
|
||
for symbol in math_symbols:
|
||
if symbol in text_without_tables:
|
||
return True
|
||
|
||
return False
|
||
|
||
def _contains_latex_tables(self, text: str) -> bool:
|
||
"""Check if text contains LaTeX table environments.
|
||
|
||
Returns:
|
||
True if text contains LaTeX tables (\\begin{table}, \\begin{tabular}, etc.)
|
||
False otherwise
|
||
"""
|
||
|
||
# Check for various LaTeX table environments
|
||
latex_table_patterns = [
|
||
r"\\begin\{table\}",
|
||
r"\\begin\{tabular\}",
|
||
]
|
||
|
||
# Check if any LaTeX table pattern exists in the text
|
||
for pattern in latex_table_patterns:
|
||
if re.search(pattern, text, re.IGNORECASE):
|
||
return True
|
||
|
||
return False
|
||
|
||
def _contains_latex_formatting_outside_math(self, text: str) -> bool:
|
||
"""Check if text contains LaTeX formatting commands outside of math equations.
|
||
|
||
Returns:
|
||
True if text contains LaTeX formatting commands outside math equations
|
||
False otherwise
|
||
"""
|
||
|
||
# List of common LaTeX formatting commands to check for
|
||
latex_commands = [
|
||
# Lists & basic content
|
||
r"\begin{itemize}",
|
||
r"\begin{enumerate}",
|
||
r"\item",
|
||
# Figures, tables, and captions
|
||
r"\begin{figure}",
|
||
r"\includegraphics",
|
||
r"\caption",
|
||
r"\label",
|
||
r"\ref",
|
||
r"\eqref",
|
||
r"\begin{table}",
|
||
r"\begin{tabular}",
|
||
# Formatting,
|
||
r"\textit",
|
||
r"\textbb",
|
||
# Math (strong signals)
|
||
r"\begin{equation}",
|
||
r"\begin{align}",
|
||
r"\frac",
|
||
r"\sum",
|
||
r"\int",
|
||
r"\sqrt",
|
||
r"\prod",
|
||
r"\lim",
|
||
r"\binom",
|
||
r"\mathbb",
|
||
r"\mathcal",
|
||
r"\to",
|
||
r"\varphi",
|
||
r"\cdot",
|
||
r"\langle",
|
||
r"\rangle",
|
||
# Citations (bibliography stacks)
|
||
r"\cite",
|
||
]
|
||
|
||
# First, remove all math equations from the text
|
||
text_without_math = text
|
||
|
||
# Patterns for math equations
|
||
math_patterns = [
|
||
r"\$\$(.+?)\$\$", # $$...$$
|
||
r"\\\((.+?)\\\)", # \(...\)
|
||
r"\\\[(.+?)\\\]", # \[...\]
|
||
]
|
||
|
||
# Remove all math equations
|
||
for pattern in math_patterns:
|
||
text_without_math = re.sub(pattern, "", text_without_math, flags=re.DOTALL)
|
||
|
||
# Check if any LaTeX commands appear in the remaining text
|
||
for command in latex_commands:
|
||
if command in text_without_math:
|
||
return True
|
||
|
||
return False
|
||
|
||
def _validate_math_equations(self, text: str) -> bool:
|
||
"""Check if all math equations in the text can render without errors.
|
||
|
||
Returns:
|
||
True if all equations render successfully or no equations exist
|
||
False if any equation fails to render
|
||
"""
|
||
|
||
# Patterns to find math equations (same as in MathTest)
|
||
patterns = [
|
||
r"\$\$(.+?)\$\$", # $$...$$
|
||
r"\\\((.+?)\\\)", # \(...\)
|
||
r"\\\[(.+?)\\\]", # \[...\]
|
||
]
|
||
|
||
equations = []
|
||
for pattern in patterns:
|
||
# Find all matches for the current pattern
|
||
matches = re.findall(pattern, text, re.DOTALL)
|
||
equations.extend([eq.strip() for eq in matches])
|
||
|
||
# If no equations found, that's fine
|
||
if not equations:
|
||
return True
|
||
|
||
# Try to render each equation
|
||
try:
|
||
from olmocr.bench.katex.render import render_equation
|
||
|
||
for equation in equations:
|
||
# Skip empty or whitespace-only equations
|
||
if not equation or not equation.strip():
|
||
continue
|
||
|
||
# Try to render the equation
|
||
rendered = render_equation(equation)
|
||
|
||
# Check if there was an error
|
||
if rendered is None or (hasattr(rendered, "error") and rendered.error):
|
||
# Equation failed to render
|
||
logger.warning(f"Could not render equation '{repr(equation)}', skipping sample")
|
||
return False
|
||
|
||
# All equations rendered successfully
|
||
return True
|
||
except Exception as e:
|
||
# If any unexpected error occurs during validation, be conservative and filter out
|
||
print(f"Error validating math equations: {e}")
|
||
return False
|
||
|
||
def _contains_br_in_table_cells(self, text: str) -> bool:
|
||
"""Check if text contains <br> tags within HTML table cells.
|
||
|
||
Returns:
|
||
True if any table cell contains <br> tags
|
||
False otherwise
|
||
"""
|
||
|
||
# Check if there are any tables in the text
|
||
if "<table" not in text.lower() or "<br" not in text.lower():
|
||
return False # No tables or no <br> tags at all
|
||
|
||
# Pattern to find HTML tables (case-insensitive)
|
||
table_pattern = re.compile(r"<table\b[^>]*>.*?</table>", re.IGNORECASE | re.DOTALL)
|
||
tables = table_pattern.findall(text)
|
||
|
||
# Check each table for <br> tags in cells
|
||
for table_html in tables:
|
||
# Pattern to find table cells (td and th tags)
|
||
cell_pattern = re.compile(r"<(td|th)\b[^>]*>(.*?)</\1>", re.IGNORECASE | re.DOTALL)
|
||
cells = cell_pattern.findall(table_html)
|
||
|
||
for tag_type, cell_content in cells:
|
||
# Check if cell content contains <br> tags (any variation)
|
||
if re.search(r"<br\s*/?>", cell_content, re.IGNORECASE):
|
||
return True
|
||
|
||
return False
|
||
|
||
def _extract_and_validate_html_tables(self, text: str) -> bool:
|
||
"""Extract HTML tables and validate they parse correctly.
|
||
|
||
Returns:
|
||
True if all HTML tables are valid or no tables exist
|
||
False if any HTML table is malformed
|
||
"""
|
||
# Find all HTML table blocks
|
||
|
||
# Check if there are any <table> tags at all
|
||
if "<table" not in text.lower():
|
||
return True # No tables, that's fine
|
||
|
||
# Pattern to find HTML tables (case-insensitive)
|
||
# Note: This pattern might not catch malformed tables where </table> is missing
|
||
table_pattern = re.compile(r"<table\b[^>]*>.*?</table>", re.IGNORECASE | re.DOTALL)
|
||
tables = table_pattern.findall(text)
|
||
|
||
# Also check for unclosed table tags
|
||
table_open_count = len(re.findall(r"<table\b[^>]*>", text, re.IGNORECASE))
|
||
table_close_count = len(re.findall(r"</table>", text, re.IGNORECASE))
|
||
|
||
if table_open_count != table_close_count:
|
||
return False # Mismatched table tags
|
||
|
||
if not tables and table_open_count > 0:
|
||
# Found table tags but couldn't extract complete tables
|
||
return False
|
||
|
||
# Try to parse each table
|
||
|
||
class TableValidator(HTMLParser):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.tag_stack = []
|
||
self.is_valid = True
|
||
self.error_msg = None
|
||
|
||
def handle_starttag(self, tag, attrs):
|
||
self.tag_stack.append(tag.lower())
|
||
|
||
def handle_endtag(self, tag):
|
||
tag = tag.lower()
|
||
if not self.tag_stack:
|
||
self.is_valid = False
|
||
self.error_msg = f"Unexpected closing tag: {tag}"
|
||
return
|
||
|
||
# Check if the closing tag matches the most recent opening tag
|
||
if self.tag_stack[-1] == tag:
|
||
self.tag_stack.pop()
|
||
else:
|
||
# For HTML, some tags can be implicitly closed (like td, tr)
|
||
# But we should still detect truly malformed tables
|
||
if tag in self.tag_stack:
|
||
# Pop until we find the matching tag
|
||
while self.tag_stack and self.tag_stack[-1] != tag:
|
||
self.tag_stack.pop()
|
||
if self.tag_stack:
|
||
self.tag_stack.pop()
|
||
else:
|
||
self.is_valid = False
|
||
self.error_msg = f"Mismatched tag: expected {self.tag_stack[-1]}, got {tag}"
|
||
|
||
def error(self, message):
|
||
self.is_valid = False
|
||
self.error_msg = message
|
||
|
||
# Validate each table
|
||
for table_html in tables:
|
||
parser = TableValidator()
|
||
try:
|
||
parser.feed(table_html)
|
||
# Check if all tags were closed
|
||
if parser.tag_stack:
|
||
return False # Unclosed tags
|
||
if not parser.is_valid:
|
||
return False # Parser found an error
|
||
except Exception:
|
||
# Any parsing exception means the table is malformed
|
||
return False
|
||
|
||
return True
|
||
|
||
def __call__(self, sample: Sample) -> Optional[Sample]:
|
||
"""Filter samples based on text content rules."""
|
||
# Get the natural text from page_data if it exists
|
||
text = None
|
||
|
||
if "page_data" in sample:
|
||
page_data = sample["page_data"]
|
||
if hasattr(page_data, "natural_text") and page_data.natural_text:
|
||
text = page_data.natural_text
|
||
|
||
# If no text to check, pass the sample through
|
||
if text is None:
|
||
return sample
|
||
|
||
# Check for markdown tables
|
||
if self._contains_markdown_table(text):
|
||
return None # Filter out samples with markdown tables
|
||
|
||
# Check for HTML tables and validate them
|
||
if not self._extract_and_validate_html_tables(text):
|
||
return None # Filter out samples with malformed HTML tables
|
||
|
||
# We had a check for <br> tags in table cells
|
||
# Note, this was maybe removing too much stuff
|
||
|
||
# Check if all math equations can render without errors
|
||
if not self._validate_math_equations(text):
|
||
return None # Filter out samples with invalid math equations
|
||
|
||
# Check for mathematical symbols
|
||
if self._contains_math_symbols(text):
|
||
return None # Filter out samples with mathematical symbols
|
||
|
||
# Check for LaTeX formatting outside math equations
|
||
if self._contains_latex_formatting_outside_math(text):
|
||
return None # Filter out samples with \textit or \textbf outside math
|
||
|
||
# Check for LaTeX tables
|
||
if self._contains_latex_tables(text):
|
||
return None # Filter out samples with LaTeX tables
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class ReformatLatexBoldItalic(PipelineStep):
|
||
"""Pipeline step that converts LaTeX formatting commands to markdown equivalents.
|
||
|
||
Converts:
|
||
- \\textit{...} to *...* (italic)
|
||
- \\textbf{...} to **...** (bold)
|
||
|
||
These conversions only happen outside of math equations.
|
||
"""
|
||
|
||
def __call__(self, sample: Sample) -> Optional[Sample]:
|
||
"""Convert LaTeX formatting to markdown in the sample text."""
|
||
# Get the natural text from page_data if it exists
|
||
if "page_data" not in sample:
|
||
return sample
|
||
|
||
page_data = sample["page_data"]
|
||
if not hasattr(page_data, "natural_text") or not page_data.natural_text:
|
||
return sample
|
||
|
||
text = page_data.natural_text
|
||
|
||
# Math equation patterns to preserve
|
||
math_patterns = [
|
||
r"\$\$(.+?)\$\$", # $$...$$
|
||
r"\\\((.+?)\\\)", # \(...\)
|
||
r"\\\[(.+?)\\\]", # \[...\]
|
||
]
|
||
|
||
# Store math equations with placeholders
|
||
math_placeholders = []
|
||
preserved_text = text
|
||
|
||
# Replace math equations with placeholders
|
||
for i, pattern in enumerate(math_patterns):
|
||
matches = re.finditer(pattern, preserved_text, re.DOTALL)
|
||
for j, match in enumerate(matches):
|
||
placeholder = f"__MATH_PLACEHOLDER_{i}_{j}__"
|
||
math_placeholders.append((placeholder, match.group(0)))
|
||
preserved_text = preserved_text.replace(match.group(0), placeholder, 1)
|
||
|
||
# Now convert LaTeX formatting to markdown
|
||
# We need to handle nested braces properly
|
||
# Use a function to find matching braces
|
||
def replace_latex_command(text, command, markdown):
|
||
"""Replace LaTeX command with markdown, handling nested braces."""
|
||
pattern = r"\\" + command + r"\{"
|
||
result = []
|
||
i = 0
|
||
|
||
while i < len(text):
|
||
match = re.search(pattern, text[i:])
|
||
if not match:
|
||
result.append(text[i:])
|
||
break
|
||
|
||
# Add text before the match
|
||
result.append(text[i : i + match.start()])
|
||
|
||
# Find the matching closing brace
|
||
start_pos = i + match.end()
|
||
brace_count = 1
|
||
j = start_pos
|
||
|
||
while j < len(text) and brace_count > 0:
|
||
if text[j] == "{":
|
||
brace_count += 1
|
||
elif text[j] == "}":
|
||
brace_count -= 1
|
||
j += 1
|
||
|
||
if brace_count == 0:
|
||
# Extract the content between braces
|
||
content = text[start_pos : j - 1]
|
||
result.append(markdown + content + markdown)
|
||
i = j
|
||
else:
|
||
# Unmatched braces, keep original
|
||
result.append(text[i + match.start() : i + match.end()])
|
||
i = i + match.end()
|
||
|
||
return "".join(result)
|
||
|
||
# Handle \textbf{...} -> **...**
|
||
preserved_text = replace_latex_command(preserved_text, "textbf", "**")
|
||
|
||
# Handle \textit{...} -> *...*
|
||
preserved_text = replace_latex_command(preserved_text, "textit", "*")
|
||
|
||
# Restore math equations
|
||
for placeholder, original in math_placeholders:
|
||
preserved_text = preserved_text.replace(placeholder, original)
|
||
|
||
# Create a new PageResponse with the updated text (since it's frozen)
|
||
|
||
updated_page_data = replace(page_data, natural_text=preserved_text)
|
||
sample["page_data"] = updated_page_data
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class TableTransformation(PipelineStep):
|
||
"""Pipeline step that applies transformations to HTML tables in the natural text.
|
||
|
||
Supported transformations:
|
||
- "annotate_dims": Adds data-totalrows and data-totalcols attributes to each table
|
||
showing the total number of rows and columns.
|
||
- "firstrowpreview": Adds an HTML comment after the opening <table> tag showing
|
||
a preview of the first row that has the maximum number of columns.
|
||
"""
|
||
|
||
transformation: str = "annotate_dims" # The transformation to apply
|
||
|
||
def _firstrowpreview(self, text: str) -> str:
|
||
"""Add an HTML comment showing a preview of the first data row."""
|
||
from olmocr.bench.table_parsing import parse_html_tables
|
||
|
||
# Find all HTML tables
|
||
table_pattern = re.compile(r"<table\b[^>]*>.*?</table>", re.IGNORECASE | re.DOTALL)
|
||
tables = table_pattern.findall(text)
|
||
|
||
if not tables:
|
||
return text
|
||
|
||
result = text
|
||
for table_html in tables:
|
||
# Parse the table to get its structure
|
||
parsed_tables = parse_html_tables(table_html)
|
||
|
||
if not parsed_tables:
|
||
continue
|
||
|
||
table_data = parsed_tables[0]
|
||
|
||
if not table_data.cell_text:
|
||
continue
|
||
|
||
# Get max columns
|
||
max_col = max(col for _, col in table_data.cell_text.keys()) + 1
|
||
|
||
# Group cells by row
|
||
rows_data: dict[int, dict[int, str]] = {}
|
||
for (row, col), cell_text in table_data.cell_text.items():
|
||
if row not in rows_data:
|
||
rows_data[row] = {}
|
||
rows_data[row][col] = cell_text
|
||
|
||
# Find first row with max_col columns
|
||
preview_row_idx = None
|
||
preview_row_data = None
|
||
for row_idx in sorted(rows_data.keys()):
|
||
if len(rows_data[row_idx]) == max_col:
|
||
preview_row_idx = row_idx
|
||
preview_row_data = rows_data[row_idx]
|
||
break
|
||
|
||
if preview_row_idx is None or preview_row_data is None:
|
||
continue
|
||
|
||
# Build the comment string
|
||
col_descriptions = []
|
||
for col_idx in sorted(preview_row_data.keys()):
|
||
cell_value = preview_row_data[col_idx].strip()
|
||
# Truncate long values
|
||
if len(cell_value) > 50:
|
||
cell_value = cell_value[:47] + "..."
|
||
col_descriptions.append(f"Column {col_idx + 1}: {cell_value}")
|
||
|
||
comment = f"<!-- Sample row ({preview_row_idx + 1}): {', '.join(col_descriptions)} -->"
|
||
|
||
# Insert the comment after the opening <table> tag
|
||
table_open_match = re.match(r"<table\b[^>]*>", table_html, re.IGNORECASE)
|
||
if table_open_match:
|
||
table_open_tag = table_open_match.group(0)
|
||
new_table_html = table_html.replace(table_open_tag, table_open_tag + comment, 1)
|
||
result = result.replace(table_html, new_table_html, 1)
|
||
|
||
return result
|
||
|
||
def _annotate_dims(self, text: str) -> str:
|
||
"""Add data-totalrows and data-totalcols attributes to HTML tables."""
|
||
from olmocr.bench.table_parsing import parse_html_tables
|
||
|
||
# Find all HTML tables
|
||
table_pattern = re.compile(r"<table\b[^>]*>.*?</table>", re.IGNORECASE | re.DOTALL)
|
||
tables = table_pattern.findall(text)
|
||
|
||
if not tables:
|
||
return text
|
||
|
||
result = text
|
||
for table_html in tables:
|
||
# Parse the table to get its structure
|
||
parsed_tables = parse_html_tables(table_html)
|
||
|
||
if not parsed_tables:
|
||
continue
|
||
|
||
table_data = parsed_tables[0]
|
||
|
||
# Get the max row and col from cell_text keys
|
||
if not table_data.cell_text:
|
||
continue
|
||
|
||
max_row = max(row for row, col in table_data.cell_text.keys()) + 1
|
||
max_col = max(col for row, col in table_data.cell_text.keys()) + 1
|
||
|
||
# Find the opening <table> tag and add the attributes
|
||
table_open_match = re.match(r"<table\b([^>]*)>", table_html, re.IGNORECASE)
|
||
if table_open_match:
|
||
existing_attrs = table_open_match.group(1)
|
||
new_attrs = f' data-totalrows="{max_row}" data-totalcols="{max_col}"'
|
||
|
||
# Check if attributes already exist
|
||
if "data-totalrows" not in existing_attrs.lower():
|
||
new_table_open = f"<table{existing_attrs}{new_attrs}>"
|
||
new_table_html = table_html.replace(table_open_match.group(0), new_table_open, 1)
|
||
result = result.replace(table_html, new_table_html, 1)
|
||
|
||
return result
|
||
|
||
def __call__(self, sample: Sample) -> Optional[Sample]:
|
||
"""Apply the specified transformation to HTML tables in the sample text."""
|
||
# Get the natural text from page_data if it exists
|
||
if "page_data" not in sample:
|
||
return sample
|
||
|
||
page_data = sample["page_data"]
|
||
if not hasattr(page_data, "natural_text") or not page_data.natural_text:
|
||
return sample
|
||
|
||
text = page_data.natural_text
|
||
|
||
# Apply the specified transformation
|
||
if self.transformation == "annotate_dims":
|
||
text = self._annotate_dims(text)
|
||
elif self.transformation == "firstrowpreview":
|
||
text = self._firstrowpreview(text)
|
||
else:
|
||
raise ValueError(f"Unknown table transformation: {self.transformation}")
|
||
|
||
# Create a new PageResponse with the updated text (since it's frozen)
|
||
updated_page_data = replace(page_data, natural_text=text)
|
||
sample["page_data"] = updated_page_data
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class AugraphyBasicAugmentations(PipelineStep):
|
||
"""Pipeline step that applies a decent selection of augraphy augmentations to the data"""
|
||
|
||
probability: float = 0.5 # Overall probability of applying any augmentation
|
||
|
||
def __call__(self, sample: Sample) -> Optional[Sample]:
|
||
"""Apply augraphy augmentations to the image in the sample."""
|
||
# Check that the image data exists
|
||
if "image" not in sample:
|
||
return sample
|
||
|
||
# Import opencv only here
|
||
import cv2
|
||
|
||
image = sample["image"]
|
||
|
||
# Skip all augmentations based on overall probability
|
||
if np.random.random() > self.probability:
|
||
return sample
|
||
|
||
# Convert from PIL to BGR for OpenCV/Augraphy
|
||
image_numpy = np.array(image)
|
||
if len(image_numpy.shape) < 3:
|
||
image_bgr = cv2.cvtColor(image_numpy, cv2.COLOR_GRAY2BGR)
|
||
else:
|
||
image_bgr = cv2.cvtColor(image_numpy, cv2.COLOR_RGB2BGR)
|
||
|
||
# Apply a basic augraphy pipeline
|
||
from augraphy import (
|
||
AugraphyPipeline,
|
||
Brightness,
|
||
InkBleed,
|
||
InkMottling,
|
||
InkShifter,
|
||
Jpeg,
|
||
LowInkPeriodicLines,
|
||
LowInkRandomLines,
|
||
OneOf,
|
||
)
|
||
|
||
# Apply geometric transformations first, maintaing scale
|
||
if np.random.random() < 0.50:
|
||
# Get dimensions
|
||
height, width = image_bgr.shape[:2]
|
||
|
||
# Random parameters for geometric transformations
|
||
angle = max(min(np.random.standard_normal(), 3), -3) # Small rotation range
|
||
scale = np.random.uniform(0.95, 1.05) # Small scale range
|
||
tx = np.random.uniform(-0.02, 0.02) * width # Translation as fraction of width
|
||
ty = np.random.uniform(-0.02, 0.02) * height # Translation as fraction of height
|
||
|
||
# Calculate center point
|
||
center = (width / 2, height / 2)
|
||
|
||
# Create transformation matrix
|
||
M = cv2.getRotationMatrix2D(center, angle, scale)
|
||
|
||
# Add translation
|
||
M[0, 2] += tx
|
||
M[1, 2] += ty
|
||
|
||
# Apply transformation
|
||
image_bgr = cv2.warpAffine(
|
||
image_bgr,
|
||
M,
|
||
(width, height),
|
||
flags=cv2.INTER_LINEAR,
|
||
borderMode=cv2.BORDER_CONSTANT,
|
||
borderValue=(255, 255, 255), # White background for documents
|
||
)
|
||
|
||
ink_phase = [
|
||
OneOf([InkBleed(p=1), LowInkRandomLines(p=1), LowInkPeriodicLines(p=1), InkMottling(p=1), InkShifter(p=1, text_shift_scale_range=(10, 15))], p=0.2),
|
||
]
|
||
|
||
paper_phase = [OneOf([Brightness(p=0.2), Jpeg(p=1)])]
|
||
|
||
post_phase = [
|
||
# Empty on purpose or else augmentations are too strong
|
||
]
|
||
|
||
augmentation_pipeline = AugraphyPipeline(ink_phase=ink_phase, paper_phase=paper_phase, post_phase=post_phase)
|
||
|
||
# Apply augmentations
|
||
augmented_image_bgr = augmentation_pipeline(image_bgr)
|
||
|
||
# Convert back to RGB and then to PIL format
|
||
augmented_image_rgb = cv2.cvtColor(augmented_image_bgr, cv2.COLOR_BGR2RGB)
|
||
augmented_image_pil = Image.fromarray(augmented_image_rgb)
|
||
|
||
# Update the sample with the augmented image
|
||
sample["image"] = augmented_image_pil
|
||
|
||
# Double-check PIL image size matches original
|
||
assert augmented_image_pil.size == image.size, f"PIL image size changed during augmentation: {image.size} -> {augmented_image_pil.size}"
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class InstructUserMessages(PipelineStep):
|
||
"""Creates instruction-following messages format for training."""
|
||
|
||
prompt_first: bool = False
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
# Prepare messages
|
||
if self.prompt_first:
|
||
messages = {
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "text", "text": sample["instruction_prompt"]},
|
||
{"type": "image", "image": sample["image"]},
|
||
],
|
||
}
|
||
else:
|
||
messages = {
|
||
"role": "user",
|
||
"content": [
|
||
{"type": "image", "image": sample["image"]},
|
||
{"type": "text", "text": sample["instruction_prompt"]},
|
||
],
|
||
}
|
||
|
||
sample["user_messages"] = messages
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class Tokenizer(PipelineStep):
|
||
"""Tokenizes messages and creates training labels with proper masking."""
|
||
|
||
processor: Any # The model processor (e.g., AutoProcessor)
|
||
masking_index: int = -100
|
||
end_of_message_token: str = "<|im_end|>" # Configurable, defaults to Qwen format
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
"""Tokenize messages and create labels for training."""
|
||
if torch is None:
|
||
raise ImportError("torch is required for Tokenizer step")
|
||
|
||
# Extract user message and response
|
||
user_messages = sample["user_messages"]
|
||
response = sample["response"]
|
||
|
||
# Apply chat template to user message only with generation prompt
|
||
# user_messages is a single dict, so wrap it in a list
|
||
text = self.processor.apply_chat_template([user_messages], tokenize=False, add_generation_prompt=True)
|
||
|
||
main_image = None
|
||
for usg_msg in user_messages["content"]:
|
||
if "image" in usg_msg:
|
||
main_image = usg_msg["image"]
|
||
break
|
||
|
||
assert main_image is not None
|
||
|
||
# Process inputs using processor
|
||
inputs = self.processor(
|
||
text=[text],
|
||
images=[main_image],
|
||
padding=True,
|
||
return_tensors="pt",
|
||
)
|
||
|
||
# Get labels by tokenizing the output text
|
||
labels = self.processor(text=[response], padding=True, return_tensors="pt")
|
||
|
||
# Append end-of-message token to the labels
|
||
end_tokens = self.processor.tokenizer(self.end_of_message_token, add_special_tokens=False)["input_ids"]
|
||
end_tokens = torch.tensor(end_tokens, dtype=inputs.input_ids.dtype)
|
||
|
||
# Handle the case where labels['input_ids'] is empty
|
||
if labels["input_ids"].shape[1] == 0:
|
||
labels_input_ids_0 = torch.tensor([], dtype=inputs.input_ids.dtype)
|
||
else:
|
||
labels_input_ids_0 = labels["input_ids"][0].to(inputs.input_ids.dtype)
|
||
|
||
labels["input_ids"] = torch.cat([labels_input_ids_0, end_tokens])
|
||
labels["input_ids"] = labels["input_ids"].unsqueeze(0)
|
||
|
||
# Concatenate input_ids and labels
|
||
input_ids = torch.cat([inputs.input_ids[0], labels.input_ids[0]], dim=0)
|
||
|
||
# All columns will participate in attention fully
|
||
attention_mask = torch.ones_like(input_ids)
|
||
|
||
# Create labels, masking the input portion with -100
|
||
labels_full = torch.full_like(input_ids, fill_value=self.masking_index)
|
||
labels_full[len(inputs.input_ids[0]) :] = labels.input_ids[0]
|
||
|
||
# Return as dict, including pixel_values
|
||
sample["input_ids"] = input_ids
|
||
sample["attention_mask"] = attention_mask
|
||
sample["labels"] = labels_full
|
||
sample["pixel_values"] = inputs.pixel_values
|
||
|
||
if hasattr(inputs, "image_grid_thw"):
|
||
sample["image_grid_thw"] = inputs.image_grid_thw[0]
|
||
|
||
return sample
|
||
|
||
|
||
@dataclass(frozen=True, slots=True)
|
||
class RandomTokenFlipper(PipelineStep):
|
||
"""Randomly flips tokens in the output (non-masked) portion and masks their labels."""
|
||
|
||
valid_token_ids: List[int] # List of valid token IDs to substitute with
|
||
token_flip_rate: float = 1e-4
|
||
masking_index: int = -100
|
||
|
||
def __call__(self, sample: Sample) -> Sample:
|
||
"""Randomly flip tokens in the non-masked portion of labels."""
|
||
if "labels" not in sample or "input_ids" not in sample:
|
||
return sample
|
||
|
||
# Work with clones to avoid modifying original tensors
|
||
labels = sample["labels"].clone() if torch.is_tensor(sample["labels"]) else torch.tensor(sample["labels"])
|
||
input_ids = sample["input_ids"].clone() if torch.is_tensor(sample["input_ids"]) else torch.tensor(sample["input_ids"])
|
||
|
||
# Find indices where labels are not masked (i.e., output tokens)
|
||
non_masked_indices = torch.where(labels != self.masking_index)[0]
|
||
|
||
if len(non_masked_indices) == 0:
|
||
return sample
|
||
|
||
# For each non-masked token, independently decide whether to flip
|
||
for idx in non_masked_indices:
|
||
if torch.rand(1).item() < self.token_flip_rate:
|
||
# Pick a random token from the valid tokens list
|
||
random_token = self.valid_token_ids[torch.randint(len(self.valid_token_ids), (1,)).item()]
|
||
input_ids[idx] = random_token
|
||
labels[idx] = self.masking_index
|
||
|
||
# Update sample with modified tensors
|
||
sample["input_ids"] = input_ids
|
||
sample["labels"] = labels
|
||
|
||
return sample
|
||
|
||
|
||
class MarkdownPDFDocumentDataset(BaseMarkdownPDFDataset):
|
||
"""Dataset that includes front matter parsing and PDF rendering by default."""
|
||
|
||
def __init__(self, root_dir: str | PathLike, target_longest_image_dim: int, front_matter_class=None):
|
||
"""
|
||
Initialize the dataset with default pipeline steps.
|
||
|
||
Args:
|
||
root_dir: Path to the root folder containing processed markdown and PDF files
|
||
target_longest_image_dim: Target dimension for the longest side of the image
|
||
front_matter_class: Optional dataclass type to validate front matter against
|
||
"""
|
||
# Create default pipeline steps
|
||
pipeline_steps = [
|
||
FrontMatterParser(front_matter_class),
|
||
PDFRenderer(target_longest_image_dim),
|
||
StaticLengthDocumentAnchoring(target_anchor_text_len=6000),
|
||
FinetuningPrompt(),
|
||
FrontMatterOutputFormat(),
|
||
InstructUserMessages(),
|
||
]
|
||
|
||
# Initialize base class with pipeline
|
||
super().__init__(root_dir, pipeline_steps)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
import argparse
|
||
from pathlib import Path
|
||
|
||
# Set up logging for testing
|
||
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
||
|
||
parser = argparse.ArgumentParser(description="Test MarkdownPDFDocumentDataset with YAML configuration")
|
||
parser.add_argument(
|
||
"--config",
|
||
type=str,
|
||
required=True,
|
||
help="Path to YAML configuration file",
|
||
)
|
||
parser.add_argument(
|
||
"--dataset-type",
|
||
type=str,
|
||
choices=["train", "eval"],
|
||
default="train",
|
||
help="Which dataset subset to display (train or eval)",
|
||
)
|
||
parser.add_argument(
|
||
"--dataset-index",
|
||
type=int,
|
||
default=0,
|
||
help="Index of dataset to use from the train/eval list",
|
||
)
|
||
parser.add_argument(
|
||
"--sample-index",
|
||
type=int,
|
||
default=0,
|
||
help="Index of sample to display in detail",
|
||
)
|
||
parser.add_argument(
|
||
"--sample-md",
|
||
type=str,
|
||
default=None,
|
||
help="Substring of markdown path to search for and display",
|
||
)
|
||
parser.add_argument(
|
||
"--analyze-tokens",
|
||
action="store_true",
|
||
help="Analyze token length distribution across entire dataset",
|
||
)
|
||
parser.add_argument(
|
||
"--save-image",
|
||
type=str,
|
||
help="Save the processed image to the specified file path (e.g., output.png)",
|
||
)
|
||
parser.add_argument(
|
||
"--save-filtered",
|
||
type=str,
|
||
help="Directory to save .md and .pdf files of filtered samples (samples that return None from pipeline)",
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
# Import config module
|
||
from olmocr.train.config import Config
|
||
|
||
# Load configuration
|
||
print(f"\n=== Loading configuration from {args.config} ===")
|
||
config = Config.from_yaml(args.config)
|
||
|
||
# Validate configuration
|
||
try:
|
||
config.validate()
|
||
except ValueError as e:
|
||
print(f"Configuration validation failed: {e}")
|
||
exit(1)
|
||
|
||
# Load processor for tokenization
|
||
print(f"\nLoading processor: {config.model.name}")
|
||
from transformers import AutoProcessor
|
||
|
||
processor = AutoProcessor.from_pretrained(config.model.name)
|
||
|
||
# Select dataset based on type
|
||
if args.dataset_type == "train":
|
||
dataset_configs = config.dataset.train
|
||
dataset_name = "train"
|
||
else:
|
||
dataset_configs = config.dataset.eval
|
||
dataset_name = "eval"
|
||
|
||
if args.dataset_index >= len(dataset_configs):
|
||
print(f"Error: Dataset index {args.dataset_index} out of range. Only {len(dataset_configs)} {dataset_name} datasets available.")
|
||
exit(1)
|
||
|
||
dataset_cfg = dataset_configs[args.dataset_index]
|
||
root_dir = dataset_cfg["root_dir"]
|
||
pipeline_steps = config.get_pipeline_steps(dataset_cfg["pipeline"], processor)
|
||
|
||
print(f"\n=== Testing {dataset_name} dataset {args.dataset_index} ===")
|
||
print(f"Root directory: {root_dir}")
|
||
print(f"Pipeline steps: {[step.__class__.__name__ for step in pipeline_steps]}")
|
||
|
||
# Create dataset
|
||
dataset = BaseMarkdownPDFDataset(root_dir, pipeline_steps)
|
||
|
||
print(f"Dataset length: {len(dataset)}")
|
||
|
||
# Handle --save-filtered option
|
||
if args.save_filtered:
|
||
import shutil
|
||
from pathlib import Path
|
||
|
||
save_dir = Path(args.save_filtered)
|
||
|
||
# Clear and create directory
|
||
if save_dir.exists():
|
||
shutil.rmtree(save_dir)
|
||
save_dir.mkdir(parents=True, exist_ok=True)
|
||
|
||
print(f"\n=== Checking for filtered samples ===")
|
||
print(f"Will save filtered samples to: {save_dir}")
|
||
|
||
# Function to process and copy a single sample
|
||
def process_and_copy_sample(idx, dataset_samples, save_dir_str):
|
||
"""Process a sample and return info if it's filtered.
|
||
|
||
Note: This function needs to be picklable for ProcessPoolExecutor,
|
||
so it takes simple arguments rather than complex objects.
|
||
"""
|
||
import shutil
|
||
from pathlib import Path
|
||
|
||
# Recreate dataset with same parameters
|
||
# This is needed because dataset objects can't be pickled
|
||
temp_dataset = BaseMarkdownPDFDataset.__new__(BaseMarkdownPDFDataset)
|
||
temp_dataset.samples = dataset_samples
|
||
temp_dataset.pipeline_steps = pipeline_steps
|
||
|
||
try:
|
||
sample = temp_dataset[idx]
|
||
if sample is None:
|
||
# This sample was filtered out - get the original paths
|
||
original_sample = dataset_samples[idx]
|
||
md_path = original_sample["markdown_path"]
|
||
pdf_path = original_sample["pdf_path"]
|
||
|
||
save_dir = Path(save_dir_str)
|
||
|
||
# Create subdirectory to preserve some structure
|
||
# Use the parent directory name and file name
|
||
rel_path = md_path.parent.name
|
||
target_subdir = save_dir / rel_path
|
||
target_subdir.mkdir(parents=True, exist_ok=True)
|
||
|
||
# Copy markdown file
|
||
target_md = target_subdir / md_path.name
|
||
shutil.copy2(md_path, target_md)
|
||
|
||
# Copy PDF file
|
||
target_pdf = target_subdir / pdf_path.name
|
||
shutil.copy2(pdf_path, target_pdf)
|
||
|
||
return {"index": idx, "markdown_path": str(md_path), "pdf_path": str(pdf_path)}
|
||
return None
|
||
except Exception as e:
|
||
print(f"Error processing sample {idx}: {e}")
|
||
return None
|
||
|
||
# Process all samples in parallel
|
||
filtered_samples = []
|
||
print(f"Processing {len(dataset)} samples to find and copy filtered ones...")
|
||
|
||
with ProcessPoolExecutor(max_workers=8) as executor:
|
||
# Submit all tasks
|
||
futures = {executor.submit(process_and_copy_sample, idx, dataset.samples, str(save_dir)): idx for idx in range(len(dataset))}
|
||
|
||
# Process results with progress bar
|
||
with tqdm(total=len(dataset), desc="Processing samples") as pbar:
|
||
for future in as_completed(futures):
|
||
result = future.result()
|
||
if result is not None:
|
||
filtered_samples.append(result)
|
||
pbar.update(1)
|
||
|
||
# Sort filtered samples by index for consistent output
|
||
filtered_samples.sort(key=lambda x: x["index"])
|
||
|
||
print(f"\nFound and copied {len(filtered_samples)} filtered samples to: {save_dir}")
|
||
|
||
if filtered_samples:
|
||
print(f"First 10 filtered samples:")
|
||
for i, sample_info in enumerate(filtered_samples[:10]):
|
||
md_name = Path(sample_info["markdown_path"]).name
|
||
print(f" Sample {sample_info['index']}: {md_name}")
|
||
if len(filtered_samples) > 10:
|
||
print(f" ... and {len(filtered_samples) - 10} more")
|
||
|
||
# Exit early if --save-filtered is used (don't continue with other analyses)
|
||
print("\nCompleted saving filtered samples. Exiting.")
|
||
exit(0)
|
||
|
||
if len(dataset) > 0:
|
||
# Show first few samples
|
||
print("\nFirst 5 samples:")
|
||
for i in range(min(5, len(dataset))):
|
||
sample = dataset.samples[i]
|
||
print(f" {i}: MD: {sample['markdown_path'].name}, PDF: {sample['pdf_path'].name}")
|
||
|
||
# Determine which sample to display
|
||
sample_idx = args.sample_index
|
||
|
||
# If --sample-md is provided, search for matching sample
|
||
if args.sample_md:
|
||
matching_indices = []
|
||
for i, s in enumerate(dataset.samples):
|
||
if args.sample_md in str(s["markdown_path"]):
|
||
matching_indices.append(i)
|
||
|
||
if len(matching_indices) == 0:
|
||
print(f"\nError: No samples found containing '{args.sample_md}' in markdown path.")
|
||
exit(1)
|
||
elif len(matching_indices) > 1:
|
||
print(f"\nError: Multiple samples found containing '{args.sample_md}':")
|
||
for idx in matching_indices[:10]: # Show first 10 matches
|
||
print(f" {idx}: {dataset.samples[idx]['markdown_path']}")
|
||
if len(matching_indices) > 10:
|
||
print(f" ... and {len(matching_indices) - 10} more")
|
||
print("\nPlease use a more specific substring.")
|
||
exit(1)
|
||
else:
|
||
sample_idx = matching_indices[0]
|
||
print(f"\nFound sample at index {sample_idx}: {dataset.samples[sample_idx]['markdown_path']}")
|
||
|
||
# Check if sample index is valid
|
||
if sample_idx >= len(dataset):
|
||
print(f"\nError: Sample index {sample_idx} out of range. Only {len(dataset)} samples available.")
|
||
exit(1)
|
||
|
||
# Get the requested sample
|
||
print(f"\n=== Displaying sample {sample_idx} ===")
|
||
sample = dataset[sample_idx]
|
||
|
||
# Display sample information based on pipeline output
|
||
print("\nSample keys:", list(sample.keys()))
|
||
|
||
# If it's raw data (no tokenization)
|
||
if "markdown_path" in sample:
|
||
print(f"\nMarkdown file: {sample['markdown_path']}")
|
||
if "pdf_path" in sample:
|
||
print(f"PDF file: {sample['pdf_path']}")
|
||
if "image" in sample and hasattr(sample["image"], "size"):
|
||
print(f"Image size: {sample['image'].size}")
|
||
|
||
# Save image if requested
|
||
if args.save_image:
|
||
sample["image"].save(args.save_image)
|
||
print(f"Saved image to: {args.save_image}")
|
||
|
||
if "page_data" in sample:
|
||
print(f"\nPage data: {sample['page_data']}")
|
||
if "messages" in sample:
|
||
print(f"\n=== Messages ===")
|
||
for i, msg in enumerate(sample["messages"]):
|
||
print(f"\nMessage {i}:")
|
||
print(f" Role: {msg['role']}")
|
||
print(f" Content preview: {str(msg['content'])[:200]}...")
|
||
|
||
# If it's tokenized data
|
||
if "input_ids" in sample:
|
||
print(f"\n=== Tokenized Output ===")
|
||
print(f" Keys: {list(sample.keys())}")
|
||
print(f" Input IDs shape: {sample['input_ids'].shape}")
|
||
print(f" Labels shape: {sample['labels'].shape}")
|
||
print(f" Attention mask shape: {sample['attention_mask'].shape}")
|
||
|
||
if "pixel_values" in sample:
|
||
print(f" Pixel values shape: {sample['pixel_values'].shape}")
|
||
if "image_grid_thw" in sample:
|
||
print(f" Image grid THW: {sample['image_grid_thw']}")
|
||
|
||
# Show label masking
|
||
print(f"\nLabel masking analysis:")
|
||
labels = sample["labels"]
|
||
# Handle both numpy arrays and torch tensors
|
||
if torch.is_tensor(labels):
|
||
masked_count = (labels == -100).sum().item()
|
||
total_count = labels.numel()
|
||
labels_array = labels.cpu().numpy() if labels.is_cuda else labels.numpy()
|
||
else:
|
||
masked_count = np.sum(labels == -100)
|
||
total_count = len(labels)
|
||
labels_array = labels
|
||
print(f" Total tokens: {total_count}")
|
||
print(f" Masked tokens: {masked_count} ({masked_count/total_count*100:.1f}%)")
|
||
print(f" Unmasked tokens: {total_count - masked_count} ({(total_count - masked_count)/total_count*100:.1f}%)")
|
||
|
||
# Find the transition point
|
||
transition_idx = None
|
||
for i in range(len(labels_array) - 1):
|
||
if labels_array[i] == -100 and labels_array[i + 1] != -100:
|
||
transition_idx = i + 1
|
||
break
|
||
|
||
if transition_idx:
|
||
print(f" Transition from masked to unmasked at position: {transition_idx}")
|
||
|
||
# Print all tokens
|
||
input_ids = sample["input_ids"]
|
||
# Handle both numpy arrays and torch tensors
|
||
if torch.is_tensor(input_ids):
|
||
input_ids_array = input_ids.cpu().numpy() if input_ids.is_cuda else input_ids.numpy()
|
||
else:
|
||
input_ids_array = input_ids
|
||
|
||
print(f"\nAll tokens ({len(input_ids_array)} total):")
|
||
print("Format: [index] Token (repr) | Label | Token ID")
|
||
print("-" * 80)
|
||
|
||
for i in range(len(input_ids_array)):
|
||
token = processor.tokenizer.decode([int(input_ids_array[i])])
|
||
token_repr = repr(token)
|
||
label = labels_array[i] if i < len(labels_array) else "N/A"
|
||
token_id = int(input_ids_array[i])
|
||
|
||
# Mark special positions
|
||
marker = ""
|
||
if transition_idx and i == transition_idx:
|
||
marker = " <-- TRANSITION (first unmasked)"
|
||
elif i == 0:
|
||
marker = " <-- START"
|
||
elif label != -100 and i > 0 and labels_array[i - 1] == -100:
|
||
marker = " <-- response begins"
|
||
|
||
print(f"[{i:4d}] {token_repr:20s} | {str(label):6s} | {token_id:6d}{marker}")
|
||
|
||
# Calculate and show token statistics after the table
|
||
print(f"\nToken statistics:")
|
||
|
||
# Count consecutive high-value tokens that represent the image
|
||
# Qwen uses tokens like 151859, 151860, etc. for image patches
|
||
image_token_threshold = 151000 # Typical threshold for Qwen image tokens
|
||
image_token_count = np.sum(input_ids_array > image_token_threshold)
|
||
|
||
# Calculate prompt tokens (everything masked)
|
||
prompt_token_count = masked_count
|
||
|
||
# Calculate output tokens (everything not masked)
|
||
output_token_count = total_count - masked_count
|
||
|
||
# Calculate non-image prompt tokens
|
||
non_image_prompt_tokens = prompt_token_count - image_token_count
|
||
|
||
print(f" Image tokens: {image_token_count}")
|
||
print(f" Prompt tokens (total): {prompt_token_count}")
|
||
print(f" Prompt tokens (non-image): {non_image_prompt_tokens}")
|
||
print(f" Output tokens: {output_token_count}")
|
||
print(f" Total sequence length: {total_count}")
|
||
|
||
# Analyze token length distribution across entire dataset
|
||
if args.analyze_tokens and "input_ids" in sample:
|
||
print(f"\n\n=== Analyzing token length distribution across entire dataset ===")
|
||
print(f"Processing {len(dataset)} samples...")
|
||
|
||
# Process samples sequentially with progress bar
|
||
# (ProcessPoolExecutor doesn't work well here because the dataset
|
||
# and pipeline steps can't be easily pickled for multiprocessing)
|
||
sequence_lengths = []
|
||
max_sequence_length = 0
|
||
max_sequence_sample_idx = 0
|
||
errors = []
|
||
|
||
for idx in tqdm(range(len(dataset)), desc="Analyzing samples"):
|
||
try:
|
||
current_sample = dataset[idx]
|
||
if current_sample is None:
|
||
continue
|
||
if "labels" in current_sample:
|
||
# Count total sequence length (all tokens, prompt + completion)
|
||
labels = current_sample["labels"]
|
||
if torch.is_tensor(labels):
|
||
total_length = labels.numel()
|
||
else:
|
||
total_length = len(labels)
|
||
sequence_lengths.append(total_length)
|
||
if total_length > max_sequence_length:
|
||
max_sequence_length = total_length
|
||
max_sequence_sample_idx = idx
|
||
else:
|
||
errors.append((idx, "No labels in sample"))
|
||
except Exception as e:
|
||
errors.append((idx, str(e)))
|
||
|
||
if errors:
|
||
print(f"\nEncountered {len(errors)} errors during processing")
|
||
if len(errors) <= 5:
|
||
for idx, error in errors:
|
||
print(f" Sample {idx}: {error}")
|
||
|
||
if sequence_lengths:
|
||
sequence_lengths = np.array(sequence_lengths)
|
||
|
||
print(f"\nTotal sequence length statistics (prompt + completion):")
|
||
print(f" Total samples analyzed: {len(sequence_lengths)}")
|
||
print(f" Max sequence length: {max_sequence_length} tokens (sample index: {max_sequence_sample_idx})")
|
||
print(f" Min sequence length: {np.min(sequence_lengths)} tokens")
|
||
print(f" Mean sequence length: {np.mean(sequence_lengths):.1f} tokens")
|
||
print(f" Median sequence length: {np.median(sequence_lengths):.1f} tokens")
|
||
print(f" Std dev: {np.std(sequence_lengths):.1f} tokens")
|
||
|
||
# Create histogram with 100-token buckets
|
||
print(f"\nSequence length histogram (100-token buckets):")
|
||
|
||
# Define buckets
|
||
bucket_size = 100
|
||
max_bucket = ((max_sequence_length // bucket_size) + 1) * bucket_size
|
||
buckets = list(range(0, max_bucket + bucket_size, bucket_size))
|
||
|
||
# Count samples in each bucket
|
||
hist, _ = np.histogram(sequence_lengths, bins=buckets)
|
||
|
||
# Find max count for scaling
|
||
max_count = max(hist)
|
||
bar_width = 50 # Width of histogram bars
|
||
|
||
print(f"\n{'Range':>15} | {'Count':>6} | Distribution")
|
||
print("-" * 80)
|
||
|
||
for i in range(len(hist)):
|
||
start = buckets[i]
|
||
end = buckets[i + 1] - 1
|
||
count = hist[i]
|
||
|
||
# Create bar
|
||
if max_count > 0:
|
||
bar_length = int((count / max_count) * bar_width)
|
||
bar = "█" * bar_length
|
||
else:
|
||
bar = ""
|
||
|
||
range_str = f"{start:>5}-{end:>5}"
|
||
print(f"{range_str:>15} | {count:>6} | {bar}")
|
||
|
||
else:
|
||
raise AssertionError("Expected some data to be created at this point")
|