Files
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
chore: import upstream snapshot with attribution
2026-07-13 13:27:09 +08:00

881 lines
35 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
import json
import os
import re
import unicodedata
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import asdict, dataclass
from enum import Enum
from typing import Dict, List, Optional, Tuple, Union
from fuzzysearch import find_near_matches
from rapidfuzz import fuzz
from tqdm import tqdm
from olmocr.repeatdetect import RepeatDetector
from .katex.render import compare_rendered_equations, render_equation
from .table_parsing import parse_html_tables, parse_markdown_tables
# Tell pytest these are not tests
__test__ = False
class TestType(str, Enum):
__test__ = False # Tell pytest this is not a test class
BASELINE = "baseline"
PRESENT = "present"
ABSENT = "absent"
ORDER = "order"
TABLE = "table"
MATH = "math"
FORMAT = "format"
FOOTNOTE = "footnote"
class TestChecked(str, Enum):
VERIFIED = "verified"
REJECTED = "rejected"
class ValidationError(Exception):
"""Exception raised for validation errors."""
pass
def normalize_text(md_content: str) -> str:
if md_content is None:
return None
# Normalize <br> and <br/> to newlines
md_content = re.sub(r"<br/?>", " ", md_content)
# Remove markdown bold formatting (** or __ for bold)
md_content = re.sub(r"\*\*(.*?)\*\*", r"\1", md_content)
md_content = re.sub(r"__(.*?)__", r"\1", md_content)
md_content = re.sub(r"</?b>", "", md_content) # Remove <b> tags if they exist
md_content = re.sub(r"</?i>", "", md_content) # Remove <i> tags if they exist
# Remove markdown italics formatting (* or _ for italics)
# Logic: The dot (.) in regex matches any character EXCEPT a newline.
# This automatically prevents matching **start \n\n end**.
# We use group \1 to ensure we match matching pairs (**...** or __...__).
md_content = re.sub(r"(\*\*|__)(.*?)\1", r"\2", md_content) # Bold
md_content = re.sub(r"(\*|_)(.*?)\1", r"\2", md_content) # Italics
# Normalize whitespace in the md_content
md_content = re.sub(r"\s+", " ", md_content)
# Convert down to a consistent unicode form, so é == e + accent, unicode forms
md_content = unicodedata.normalize("NFC", md_content)
# Dictionary of characters to replace: keys are fancy characters, values are ASCII equivalents, unicode micro with greek mu comes up often enough too
replacements = {"": "'", "": "'", "": "'", "“": '"', "”": '"', "„": '"', "_": "_", "": "-", "—": "-", "": "-", "": "-", "": "-", "\u00b5": "\u03bc"}
# Apply all replacements from the dictionary
for fancy_char, ascii_char in replacements.items():
md_content = md_content.replace(fancy_char, ascii_char)
return md_content
@dataclass(kw_only=True)
class BasePDFTest:
"""
Base class for all PDF test types.
Attributes:
pdf: The PDF filename.
page: The page number for the test.
id: Unique identifier for the test.
type: The type of test.
threshold: A float between 0 and 1 representing the threshold for fuzzy matching.
"""
pdf: str
page: int
id: str
type: str
max_diffs: int = 0
checked: Optional[TestChecked] = None
url: Optional[str] = None
def __post_init__(self):
if not self.pdf:
raise ValidationError("PDF filename cannot be empty")
if not self.id:
raise ValidationError("Test ID cannot be empty")
if not isinstance(self.max_diffs, int) or self.max_diffs < 0:
raise ValidationError("Max diffs must be positive number or 0")
if self.type not in {t.value for t in TestType}:
raise ValidationError(f"Invalid test type: {self.type}")
def run(self, md_content: str) -> Tuple[bool, str]:
"""
Run the test on the provided markdown content.
Args:
md_content: The content of the .md file.
Returns:
A tuple (passed, explanation) where 'passed' is True if the test passes,
and 'explanation' provides details when the test fails.
"""
raise NotImplementedError("Subclasses must implement the run method")
@dataclass
class TextPresenceTest(BasePDFTest):
"""
Test to verify the presence or absence of specific text in a PDF.
Attributes:
text: The text string to search for.
"""
text: str
case_sensitive: bool = True
first_n: Optional[int] = None
last_n: Optional[int] = None
def __post_init__(self):
super().__post_init__()
if self.type not in {TestType.PRESENT.value, TestType.ABSENT.value}:
raise ValidationError(f"Invalid type for TextPresenceTest: {self.type}")
self.text = normalize_text(self.text)
if not self.text.strip():
raise ValidationError("Text field cannot be empty")
def run(self, md_content: str) -> Tuple[bool, str]:
reference_query = self.text
# Normalize whitespace in the md_content
md_content = normalize_text(md_content)
if not self.case_sensitive:
reference_query = reference_query.lower()
md_content = md_content.lower()
if self.first_n and self.last_n:
md_content = md_content[: self.first_n] + md_content[-self.last_n :]
elif self.first_n:
md_content = md_content[: self.first_n]
elif self.last_n:
md_content = md_content[-self.last_n :]
# Threshold for fuzzy matching derived from max_diffs
threshold = 1.0 - (self.max_diffs / (len(reference_query) if len(reference_query) > 0 else 1))
best_ratio = fuzz.partial_ratio(reference_query, md_content) / 100.0
if self.type == TestType.PRESENT.value:
if best_ratio >= threshold:
return True, ""
else:
msg = f"Expected '{reference_query[:40]}...' with threshold {threshold} " f"but best match ratio was {best_ratio:.3f}"
return False, msg
else: # ABSENT
if best_ratio < threshold:
return True, ""
else:
msg = f"Expected absence of '{reference_query[:40]}...' with threshold {threshold} " f"but best match ratio was {best_ratio:.3f}"
return False, msg
@dataclass
class TextOrderTest(BasePDFTest):
"""
Test to verify that one text appears before another in a PDF.
Attributes:
before: The text expected to appear first.
after: The text expected to appear after the 'before' text.
"""
before: str
after: str
def __post_init__(self):
super().__post_init__()
if self.type != TestType.ORDER.value:
raise ValidationError(f"Invalid type for TextOrderTest: {self.type}")
self.before = normalize_text(self.before)
self.after = normalize_text(self.after)
if not self.before.strip():
raise ValidationError("Before field cannot be empty")
if not self.after.strip():
raise ValidationError("After field cannot be empty")
if self.max_diffs > len(self.before) // 2 or self.max_diffs > len(self.after) // 2:
raise ValidationError("Max diffs is too large for this test, greater than 50% of the search string")
def run(self, md_content: str) -> Tuple[bool, str]:
md_content = normalize_text(md_content)
before_matches = find_near_matches(self.before, md_content, max_l_dist=self.max_diffs)
after_matches = find_near_matches(self.after, md_content, max_l_dist=self.max_diffs)
if not before_matches:
return False, f"'before' text '{self.before[:40]}...' not found with max_l_dist {self.max_diffs}"
if not after_matches:
return False, f"'after' text '{self.after[:40]}...' not found with max_l_dist {self.max_diffs}"
for before_match in before_matches:
for after_match in after_matches:
if before_match.start < after_match.start:
return True, ""
return False, (f"Could not find a location where '{self.before[:40]}...' appears before " f"'{self.after[:40]}...'.")
@dataclass
class FormatTest(BasePDFTest):
"""
Test to verify that specific text appears with the correct formatting.
Attributes:
text: The text to search for.
format: The expected format ("heading", "bold", or "italic").
"""
text: str
format: str
case_sensitive: bool = True
def __post_init__(self):
super().__post_init__()
if self.type != TestType.FORMAT.value:
raise ValidationError(f"Invalid type for FormatTest: {self.type}")
self.text = normalize_text(self.text)
if not self.text.strip():
raise ValidationError("Text field cannot be empty")
if self.format not in {"heading", "bold", "italic"}:
raise ValidationError(f"Invalid format type: {self.format}. Must be 'heading', 'bold', or 'italic'")
def run(self, md_content: str) -> Tuple[bool, str]:
"""
Extract all text with the specified format and check if our text is among them.
"""
# Store the original content before any normalization for pattern matching
original_content = md_content
# Extract formatted text based on the format type
formatted_texts = []
if self.format == "heading":
# Markdown headings (# through ######)
heading_patterns = [
r"^#{1,6}\s+(.+?)$", # Standard markdown headings
]
for pattern in heading_patterns:
matches = re.findall(pattern, original_content, re.MULTILINE)
formatted_texts.extend(matches)
# HTML headings (<h1> through <h6>)
html_heading_pattern = r"<h[1-6][^>]*>(.*?)</h[1-6]>"
matches = re.findall(html_heading_pattern, original_content, re.IGNORECASE | re.DOTALL)
formatted_texts.extend(matches)
elif self.format == "bold":
# Markdown bold patterns
bold_patterns = [
r"\*\*(.*?)\*\*", # **text**
r"__(.*?)__", # __text__
]
for pattern in bold_patterns:
matches = re.findall(pattern, original_content, re.DOTALL)
formatted_texts.extend(matches)
# HTML bold patterns
html_bold_patterns = [r"<b[^>]*>(.*?)</b>", r"<strong[^>]*>(.*?)</strong>"] # <b>text</b> # <strong>text</strong>
for pattern in html_bold_patterns:
matches = re.findall(pattern, original_content, re.IGNORECASE | re.DOTALL)
formatted_texts.extend(matches)
elif self.format == "italic":
# Markdown italic patterns - be careful not to match bold
# We need to match single * or _ that are not part of ** or __
italic_patterns = [
r"(?<!\*)\*(?!\*)(.*?)(?<!\*)\*(?!\*)", # *text* but not **text**
r"(?<!_)_(?!_)(.*?)(?<!_)_(?!_)", # _text_ but not __text__
]
for pattern in italic_patterns:
matches = re.findall(pattern, original_content, re.DOTALL)
formatted_texts.extend(matches)
# HTML italic patterns
html_italic_patterns = [r"<i[^>]*>(.*?)</i>", r"<em[^>]*>(.*?)</em>"] # <i>text</i> # <em>text</em>
for pattern in html_italic_patterns:
matches = re.findall(pattern, original_content, re.IGNORECASE | re.DOTALL)
formatted_texts.extend(matches)
# Normalize all extracted formatted texts
normalized_formatted_texts = [normalize_text(text) for text in formatted_texts]
# Normalize the search text
search_text = self.text
if not self.case_sensitive:
search_text = search_text.lower()
normalized_formatted_texts = [text.lower() for text in normalized_formatted_texts]
# Check if the text appears in any of the formatted texts using fuzzy matching
threshold = 1.0 - (self.max_diffs / (len(search_text) if len(search_text) > 0 else 1))
for formatted_text in normalized_formatted_texts:
# Use partial_ratio for substring matching
ratio = fuzz.partial_ratio(search_text, formatted_text) / 100.0
if ratio >= threshold:
return True, ""
# If we didn't find the text with the specified format
found_formats = []
if len(normalized_formatted_texts) > 0:
# Show a sample of what we did find with this format
sample = normalized_formatted_texts[:3]
sample_str = ", ".join([f"'{t[:20]}...'" if len(t) > 20 else f"'{t}'" for t in sample])
found_formats.append(f"Found {self.format} text: {sample_str}")
else:
found_formats.append(f"No {self.format} formatted text found")
return False, f"Text '{self.text[:40]}...' not found with {self.format} formatting. {'; '.join(found_formats)}"
@dataclass
class TableTest(BasePDFTest):
"""
Test to verify certain properties of a table are held, namely that some cells appear relative to other cells correctly
"""
# This is the target cell, which must exist in at least one place in the table
cell: str
# These properties say that the cell immediately up/down/left/right of the target cell has the string specified
up: str = ""
down: str = ""
left: str = ""
right: str = ""
# These properties say that the cell all the way up, or all the way left of the target cell (ex. headings) has the string value specified
top_heading: str = ""
left_heading: str = ""
ignore_markdown_tables: bool = False
def __post_init__(self):
super().__post_init__()
if self.type != TestType.TABLE.value:
raise ValidationError(f"Invalid type for TableTest: {self.type}")
# Normalize the search text too
self.cell = normalize_text(self.cell)
self.up = normalize_text(self.up)
self.down = normalize_text(self.down)
self.left = normalize_text(self.left)
self.right = normalize_text(self.right)
self.top_heading = normalize_text(self.top_heading)
self.left_heading = normalize_text(self.left_heading)
def run(self, content: str) -> Tuple[bool, str]:
"""
Run the table test on provided content.
Finds all tables (markdown and/or HTML based on content_type) and checks if any cell
matches the target cell and satisfies the specified relationships.
Args:
content: The content containing tables (markdown or HTML)
Returns:
A tuple (passed, explanation) where 'passed' is True if the test passes,
and 'explanation' provides details when the test fails.
"""
# Initialize variables to track tables and results
tables_to_check = []
failed_reasons = []
# Threshold for fuzzy matching derived from max_diffs
threshold = 1.0 - (self.max_diffs / (len(self.cell) if len(self.cell) > 0 else 1))
threshold = max(0.5, threshold)
# Parse tables based on content_type
if not self.ignore_markdown_tables:
md_tables = parse_markdown_tables(content)
tables_to_check.extend(md_tables)
html_tables = parse_html_tables(content)
tables_to_check.extend(html_tables)
# If no tables found, return failure
if not tables_to_check:
return False, "No tables found in the content"
# Check each table
for table_data in tables_to_check:
# Find all cells that match the target cell using fuzzy matching
matches = []
for rowcol, cell_content in table_data.cell_text.items():
similarity = fuzz.ratio(self.cell, normalize_text(cell_content)) / 100.0
if similarity >= threshold:
matches.append(rowcol)
# If no matches found in this table, continue to the next table
if not matches:
continue
# Check the relationships for each matching cell
for rowcol in matches:
all_relationships_satisfied = True
current_failed_reasons = []
def _check_relationship(comparison_str: str, relation_func):
nonlocal all_relationships_satisfied
cur_relation_satisified = False
best_similarity = 0
best_similarity_text = None
for rowcol_up in relation_func(rowcol):
test_cell = normalize_text(table_data.cell_text[rowcol_up])
test_similarity = fuzz.ratio(comparison_str, test_cell) / 100.0
if test_similarity > best_similarity:
best_similarity = test_similarity
best_similarity_text = test_cell
if test_similarity >= max(0.5, 1.0 - (self.max_diffs / (len(comparison_str) if len(comparison_str) > 0 else 1))):
cur_relation_satisified = True
if not cur_relation_satisified:
all_relationships_satisfied = False
current_failed_reasons.append(
f"Cell compared to '{best_similarity_text}' doesn't match expected '{comparison_str}' (best similarity: {best_similarity:.2f})"
)
# Check up relationship
if self.up:
_check_relationship(self.up, lambda rowcol: table_data.up_relations[rowcol])
if self.down:
_check_relationship(self.down, lambda rowcol: table_data.down_relations[rowcol])
if self.left:
_check_relationship(self.left, lambda rowcol: table_data.left_relations[rowcol])
if self.right:
_check_relationship(self.right, lambda rowcol: table_data.right_relations[rowcol])
if self.left_heading:
_check_relationship(self.left_heading, lambda rowcol: table_data.left_heading_relations(*rowcol))
if self.top_heading:
_check_relationship(self.top_heading, lambda rowcol: table_data.top_heading_relations(*rowcol))
# If all relationships are satisfied for this cell, the test passes
if all_relationships_satisfied:
return True, ""
else:
failed_reasons.extend(current_failed_reasons)
# If we've gone through all tables and all matching cells and none satisfied all relationships
if not failed_reasons:
return False, f"No cell matching '{self.cell}' found in any table with threshold {threshold}"
else:
return False, f"Found cells matching '{self.cell}' but relationships were not satisfied: {'; '.join(failed_reasons)}"
@dataclass
class BaselineTest(BasePDFTest):
"""
This test makes sure that several baseline quality checks pass for the output generation.
Namely, the output is not blank, not endlessly repeating, and contains characters of the proper
character sets.
"""
max_length: Optional[int] = None # Used to implement blank page checks
max_length_skips_image_alt_tags: bool = False
max_repeats: int = 30
check_disallowed_characters: bool = True
def run(self, content: str) -> Tuple[bool, str]:
base_content_len = len("".join(c for c in content if c.isalnum()).strip())
# If this a blank page check, then it short circuits the rest of the checks
if self.max_length is not None:
if self.max_length_skips_image_alt_tags:
# Remove markdown image tags like ![alt text](image.png) from the text length count
content_for_length_check = re.sub(r"!\[.*?\]\(.*?\)", "", content)
base_content_len = len("".join(c for c in content_for_length_check if c.isalnum()).strip())
if base_content_len > self.max_length:
return False, f"{base_content_len} characters were output for a page we expected to be blank"
else:
return True, ""
if base_content_len == 0:
return False, "The text contains no alpha numeric characters"
# Makes sure that the content has no egregious repeated ngrams at the end, which indicate a degradation of quality
# Honestly, this test doesn't seem to catch anything at the moment, maybe it can be refactored to a "text-quality"
# test or something, that measures repetition, non-blanks, charsets, etc
d = RepeatDetector(max_ngram_size=5)
d.add_letters(content)
repeats = d.ngram_repeats()
for index, count in enumerate(repeats):
if count > self.max_repeats:
return False, f"Text ends with {count} repeating {index+1}-grams, invalid"
pattern = re.compile(
r"["
r"\u4e00-\u9FFF" # CJK Unified Ideographs (Chinese characters)
r"\u3040-\u309F" # Hiragana (Japanese)
r"\u30A0-\u30FF" # Katakana (Japanese)
r"\U0001F600-\U0001F64F" # Emoticons (Emoji)
r"\U0001F300-\U0001F5FF" # Miscellaneous Symbols and Pictographs (Emoji)
r"\U0001F680-\U0001F6FF" # Transport and Map Symbols (Emoji)
r"\U0001F1E0-\U0001F1FF" # Regional Indicator Symbols (flags, Emoji)
r"]",
flags=re.UNICODE,
)
matches = pattern.findall(content)
if self.check_disallowed_characters and matches:
return False, f"Text contains disallowed characters {matches}"
return True, ""
@dataclass
class MathTest(BasePDFTest):
math: str
ignore_dollar_delimited: bool = False
def __post_init__(self):
super().__post_init__()
if self.type != TestType.MATH.value:
raise ValidationError(f"Invalid type for MathTest: {self.type}")
if len(self.math.strip()) == 0:
raise ValidationError("Math test must have non-empty math expression")
self.reference_render = render_equation(self.math)
if self.reference_render is None:
raise ValidationError(f"Math equation {self.math} was not able to render")
def run(self, content: str) -> Tuple[bool, str]:
# Store both the search pattern and the full pattern to replace
patterns = [
(r"\\\((.+?)\\\)", r"\\\((.+?)\\\)"), # \(...\)
(r"\\\[(.+?)\\\]", r"\\\[(.+?)\\\]"), # \[...\]
]
if not self.ignore_dollar_delimited:
patterns.extend(
[
(r"\$\$(.+?)\$\$", r"\$\$(.+?)\$\$"), # $$...$$
(r"\$(.+?)\$", r"\$(.+?)\$"), # $...$])
]
)
equations = []
modified_content = content
for search_pattern, replace_pattern in patterns:
# Find all matches for the current pattern
matches = re.findall(search_pattern, modified_content, re.DOTALL)
equations.extend([e.strip() for e in matches])
# Replace all instances of this pattern with empty strings
modified_content = re.sub(replace_pattern, "", modified_content, flags=re.DOTALL)
# If an equation in the markdown exactly matches our math string, then that's good enough
# we don't have to do a more expensive comparison
if any(hyp == self.math for hyp in equations):
return True, ""
# If not, then let's render the math equation itself and now compare to each hypothesis
# But, to speed things up, since rendering equations is hard, we sort the equations on the page
# by fuzzy similarity to the hypothesis
equations.sort(key=lambda x: -fuzz.ratio(x, self.math))
for hypothesis in equations:
hypothesis_render = render_equation(hypothesis)
if not hypothesis_render:
continue
if compare_rendered_equations(self.reference_render, hypothesis_render):
return True, ""
# self.reference_render.save(f"maths/{self.id}_ref.png", format="PNG")
# best_match_render.save(f"maths/{self.id}_hyp.png", format="PNG")
return False, f"No match found for {self.math} anywhere in content"
@dataclass
class FootnoteTest(BasePDFTest):
"""
Test to verify that footnotes appear correctly on a page.
Attributes:
marker: The footnote marker (e.g., "1", "2"). Must appear as superscript or [^marker]. Required.
appears_before_marker: Optional text that should appear before the marker (ignoring whitespace/non-alpha).
appears_after_marker: Optional text that should appear after the marker (ignoring whitespace/non-alpha).
"""
marker: str
appears_before_marker: Optional[str] = None
appears_after_marker: Optional[str] = None
def __post_init__(self):
super().__post_init__()
if self.type != TestType.FOOTNOTE.value:
raise ValidationError(f"Invalid type for FootnoteTest: {self.type}")
# marker is required
if not self.marker:
raise ValidationError("marker field is required")
# Validate marker doesn't contain whitespace
if " " in self.marker:
raise ValidationError("Marker cannot contain whitespace")
# Normalize the optional text fields
if self.appears_before_marker:
self.appears_before_marker = normalize_text(self.appears_before_marker)
if not self.appears_before_marker.strip():
raise ValidationError("appears_before_marker field cannot be empty if provided")
if self.appears_after_marker:
self.appears_after_marker = normalize_text(self.appears_after_marker)
if not self.appears_after_marker.strip():
raise ValidationError("appears_after_marker field cannot be empty if provided")
def run(self, md_content: str) -> Tuple[bool, str]:
"""
Run the footnote test on provided markdown content.
Args:
md_content: The markdown content to test.
Returns:
A tuple (passed, explanation) where 'passed' is True if the test passes,
and 'explanation' provides details when the test fails.
"""
# Find all occurrences of the marker in various formats
marker_positions = []
# Check for markdown footnote reference [^marker] (but not definition [^marker]:)
markdown_pattern = rf"\[\^{re.escape(self.marker)}\](?!:)"
for match in re.finditer(markdown_pattern, md_content):
marker_positions.append({"start": match.start(), "end": match.end(), "type": "markdown"})
# Check for superscript HTML <sup>marker</sup>
html_sup_pattern = rf"<sup[^>]*>{re.escape(self.marker)}</sup>"
for match in re.finditer(html_sup_pattern, md_content, re.IGNORECASE):
marker_positions.append({"start": match.start(), "end": match.end(), "type": "html"})
# Check for Unicode superscript characters (for common digits)
superscript_map = {"0": "⁰", "1": "¹", "2": "²", "3": "³", "4": "⁴", "5": "⁵", "6": "⁶", "7": "⁷", "8": "⁸", "9": "⁹"}
# Convert marker to superscript if all characters are digits
if all(c in superscript_map for c in self.marker):
superscript_marker = "".join(superscript_map[c] for c in self.marker)
for match in re.finditer(re.escape(superscript_marker), md_content):
marker_positions.append({"start": match.start(), "end": match.end(), "type": "unicode"})
# If no markers found at all, fail
if not marker_positions:
return False, f"Footnote marker '{self.marker}' not found as [^{self.marker}], <sup>{self.marker}</sup>, or superscript"
# If no additional checks needed, pass
if not self.appears_before_marker and not self.appears_after_marker:
return True, ""
# Helper function to clean text for comparison (remove whitespace and non-alpha)
def clean_for_comparison(text):
# Remove all non-alphanumeric characters and normalize
return "".join(c for c in normalize_text(text) if c.isalnum()).lower()
# Check appears_before_marker if provided
before_found = False if self.appears_before_marker else True
if self.appears_before_marker:
clean_target_before = clean_for_comparison(self.appears_before_marker)
threshold = 1.0 - (self.max_diffs / (len(self.appears_before_marker) if len(self.appears_before_marker) > 0 else 1))
for pos in marker_positions:
# Get text before this marker position
start_pos = max(0, pos["start"] - 200) # Look back up to 200 chars
text_before = md_content[start_pos : pos["start"]]
# Clean the text before for comparison
clean_text_before = clean_for_comparison(text_before)
# Check if appears_before_marker is at the end of this text (using fuzzy matching)
if clean_text_before:
# Use partial_ratio to check if target appears at the end
# We'll check the last portion that's roughly the size of our target
check_length = min(len(clean_text_before), len(clean_target_before) * 2)
text_to_check = clean_text_before[-check_length:] if check_length > 0 else clean_text_before
similarity = fuzz.partial_ratio(clean_target_before, text_to_check) / 100.0
if similarity >= threshold:
before_found = True
break
# Check appears_after_marker if provided
after_found = False if self.appears_after_marker else True
if self.appears_after_marker:
clean_target_after = clean_for_comparison(self.appears_after_marker)
threshold = 1.0 - (self.max_diffs / (len(self.appears_after_marker) if len(self.appears_after_marker) > 0 else 1))
for pos in marker_positions:
# Get text after this marker position
end_pos = min(len(md_content), pos["end"] + 200) # Look ahead up to 200 chars
text_after = md_content[pos["end"] : end_pos]
# Clean the text after for comparison
clean_text_after = clean_for_comparison(text_after)
# Check if appears_after_marker is at the beginning of this text (using fuzzy matching)
if clean_text_after:
# Use partial_ratio to check if target appears at the beginning
# We'll check the first portion that's roughly the size of our target
check_length = min(len(clean_text_after), len(clean_target_after) * 2)
text_to_check = clean_text_after[:check_length] if check_length > 0 else clean_text_after
similarity = fuzz.partial_ratio(clean_target_after, text_to_check) / 100.0
if similarity >= threshold:
after_found = True
break
# Build failure message if needed
failures = []
if self.appears_before_marker and not before_found:
failures.append(f"Text '{self.appears_before_marker[:40]}...' not found before any occurrence of marker '{self.marker}'")
if self.appears_after_marker and not after_found:
failures.append(f"Text '{self.appears_after_marker[:40]}...' not found after any occurrence of marker '{self.marker}'")
if failures:
return False, "; ".join(failures)
else:
return True, ""
def load_single_test(data: Union[str, Dict]) -> BasePDFTest:
"""
Load a single test from a JSON line string or JSON object.
Args:
data: Either a JSON string to parse or a dictionary containing test data.
Returns:
A test object of the appropriate type.
Raises:
ValidationError: If the test type is unknown or data is invalid.
json.JSONDecodeError: If the string cannot be parsed as JSON.
"""
# Handle JSON string input
if isinstance(data, str):
data = data.strip()
if not data:
raise ValueError("Empty string provided")
data = json.loads(data)
# Process the test data
test_type = data.get("type")
if test_type in {TestType.PRESENT.value, TestType.ABSENT.value}:
test = TextPresenceTest(**data)
elif test_type == TestType.ORDER.value:
test = TextOrderTest(**data)
elif test_type == TestType.TABLE.value:
test = TableTest(**data)
elif test_type == TestType.MATH.value:
test = MathTest(**data)
elif test_type == TestType.BASELINE.value:
test = BaselineTest(**data)
elif test_type == TestType.FORMAT.value:
test = FormatTest(**data)
elif test_type == TestType.FOOTNOTE.value:
test = FootnoteTest(**data)
else:
raise ValidationError(f"Unknown test type: {test_type}")
return test
def load_tests(jsonl_file: str) -> List[BasePDFTest]:
"""
Load tests from a JSONL file using parallel processing with a ThreadPoolExecutor.
Args:
jsonl_file: Path to the JSONL file containing test definitions.
Returns:
A list of test objects.
"""
def process_line_with_number(line_tuple: Tuple[int, str]) -> Optional[Tuple[int, BasePDFTest]]:
"""
Process a single line from the JSONL file and return a tuple of (line_number, test object).
Returns None for empty lines.
"""
line_number, line = line_tuple
line = line.strip()
if not line:
return None
try:
test = load_single_test(line)
return (line_number, test)
except json.JSONDecodeError as e:
print(f"Error parsing JSON on line {line_number}: {e}")
raise
except (ValidationError, KeyError) as e:
print(f"Error on line {line_number}: {e}")
raise
except Exception as e:
print(f"Unexpected error on line {line_number}: {e}")
raise
tests = []
# Read all lines along with their line numbers.
with open(jsonl_file, "r") as f:
lines = list(enumerate(f, start=1))
# Use a ThreadPoolExecutor to process each line in parallel.
with ThreadPoolExecutor(max_workers=min(os.cpu_count() or 1, 64)) as executor:
# Submit all tasks concurrently.
futures = {executor.submit(process_line_with_number, item): item[0] for item in lines}
# Use tqdm to show progress as futures complete.
for future in tqdm(as_completed(futures), total=len(futures), desc="Loading tests"):
result = future.result()
if result is not None:
_, test = result
tests.append(test)
# Check for duplicate test IDs after parallel processing.
unique_ids = set()
for test in tests:
if test.id in unique_ids:
raise ValidationError(f"Test with duplicate id {test.id} found, error loading tests.")
unique_ids.add(test.id)
return tests
def save_tests(tests: List[BasePDFTest], jsonl_file: str) -> None:
"""
Save tests to a JSONL file using asdict for conversion.
Args:
tests: A list of test objects.
jsonl_file: Path to the output JSONL file.
"""
with open(jsonl_file, "w") as file:
for test in tests:
file.write(json.dumps(asdict(test)) + "\n")