render.py

#!/usr/bin/env python3
"""
Extract inner-most spans and their bounding boxes, and the mathML output, 
from rendered LaTeX equations using Playwright and KaTeX.
Caching is maintained via a SHA1-based hash stored as a JSON file.

Requirements:
    pip install playwright
    python -m playwright install chromium

    Place katex.min.css and katex.min.js in the same directory as this script
"""

import os
import re
import html
import hashlib
import pathlib
import json
import re
import shutil
from dataclasses import dataclass
from typing import List
import unittest
import html.entities
from lxml import etree

from playwright.sync_api import sync_playwright, Error as PlaywrightError

@dataclass
class BoundingBox:
    x: float
    y: float
    width: float
    height: float

@dataclass
class SpanInfo:
    text: str
    bounding_box: BoundingBox

@dataclass
class RenderedEquation:
    mathml: str
    spans: List[SpanInfo]

def get_equation_hash(equation, bg_color="white", text_color="black", font_size=24):
    """
    Calculate SHA1 hash of the equation string and rendering parameters.
    """
    params_str = f"{equation}|{bg_color}|{text_color}|{font_size}"
    return hashlib.sha1(params_str.encode('utf-8')).hexdigest()

def get_cache_dir():
    """
    Get the cache directory for equations, creating it if it doesn't exist.
    """
    cache_dir = pathlib.Path.home() / '.cache' / 'olmocr' / 'bench' / 'equations'
    cache_dir.mkdir(parents=True, exist_ok=True)
    return cache_dir


def clear_cache_dir():
    """
    Clear all files and subdirectories in the cache directory.
    """
    cache_dir = get_cache_dir()
    if cache_dir.exists() and cache_dir.is_dir():
        shutil.rmtree(cache_dir)
        cache_dir.mkdir(parents=True, exist_ok=True)  # Recreate the empty directory


def render_equation(
    equation, 
    bg_color="white",
    text_color="black",
    font_size=24,
    use_cache=True,
    debug_dom=False,
):
    """
    Render a LaTeX equation using Playwright and KaTeX, extract the inner-most span elements
    (those without child elements that contain non-whitespace text) along with their bounding boxes,
    and also extract the MathML output generated by KaTeX.

    Returns:
        RenderedEquation: A dataclass containing the mathml string and a list of SpanInfo dataclasses.
    """
    # Calculate hash for caching
    eq_hash = get_equation_hash(equation, bg_color, text_color, font_size)
    cache_dir = get_cache_dir()
    cache_file = cache_dir / f"{eq_hash}.json"
    cache_error_file = cache_dir / f"{eq_hash}_error"
    
    if use_cache:
        if cache_error_file.exists():
            return None
        if cache_file.exists():
            with open(cache_file, 'r') as f:
                data = json.load(f)
            spans = [
                SpanInfo(
                    text=s["text"],
                    bounding_box=BoundingBox(
                        x=s["boundingBox"]["x"],
                        y=s["boundingBox"]["y"],
                        width=s["boundingBox"]["width"],
                        height=s["boundingBox"]["height"],
                    )
                )
                for s in data["spans"]
            ]
            return RenderedEquation(mathml=data["mathml"], spans=spans)

    # Escape backslashes for JavaScript string
    escaped_equation = json.dumps(equation)
    
    # Get local paths for KaTeX files
    script_dir = os.path.dirname(os.path.abspath(__file__))
    katex_css_path = os.path.join(script_dir, "katex.min.css")
    katex_js_path = os.path.join(script_dir, "katex.min.js")
    
    if not os.path.exists(katex_css_path) or not os.path.exists(katex_js_path):
        raise FileNotFoundError(f"KaTeX files not found. Please ensure katex.min.css and katex.min.js are in {script_dir}")
    
    with sync_playwright() as p:
        browser = p.chromium.launch()
        page = browser.new_page(viewport={"width": 800, "height": 400})
        
        # Basic HTML structure
        page_html = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <style>
                body {{
                    display: flex;
                    justify-content: center;
                    align-items: center;
                    height: 100vh;
                    margin: 0;
                    background-color: {bg_color};
                    color: {text_color};
                }}
                #equation-container {{
                    padding: 0;
                    font-size: {font_size}px;
                }}
            </style>
        </head>
        <body>
            <div id="equation-container"></div>
        </body>
        </html>
        """
        page.set_content(page_html)
        page.add_style_tag(path=katex_css_path)
        page.add_script_tag(path=katex_js_path)
        page.wait_for_load_state("networkidle")
        
        katex_loaded = page.evaluate("typeof katex !== 'undefined'")
        if not katex_loaded:
            raise RuntimeError("KaTeX library failed to load. Check your katex.min.js file.")
        
        try:
            error_message = page.evaluate(f"""
            () => {{
                try {{
                    katex.render({escaped_equation}, document.getElementById("equation-container"), {{
                        displayMode: true,
                        throwOnError: true
                    }});
                    return null;
                }} catch (error) {{
                    console.error("KaTeX error:", error.message);
                    return error.message;
                }}
            }}
            """)
        except PlaywrightError as ex:
            print(escaped_equation)
            error_message = str(ex)
            raise 

        if error_message:
            print(f"Error rendering equation: '{equation}'")
            print(error_message)
            cache_error_file.touch()
            browser.close()
            return None
        
        page.wait_for_selector(".katex", state="attached")
        
        if debug_dom:
            katex_dom_html = page.evaluate("""
            () => {
                return document.getElementById("equation-container").innerHTML;
            }
            """)
            print("\n===== KaTeX DOM HTML =====")
            print(katex_dom_html)
        
        # Extract inner-most spans with non-whitespace text
        spans_info = page.evaluate("""
        () => {
            const spans = Array.from(document.querySelectorAll('span'));
            const list = [];
            spans.forEach(span => {
                // Check if this span has no child elements and contains non-whitespace text
                if (span.children.length === 0 && /\S/.test(span.textContent)) {
                    const rect = span.getBoundingClientRect();
                    list.push({
                        text: span.textContent.trim(),
                        boundingBox: {
                            x: rect.x,
                            y: rect.y,
                            width: rect.width,
                            height: rect.height
                        }
                    });
                }
            });
            return list;
        }
        """)
        
        if debug_dom:
            print("\n===== Extracted Span Information =====")
            print(spans_info)
        
        # Extract mathML output (if available) from the KaTeX output.
        # We try to get the <math> element within an element with class "katex-mathml".
        mathml = page.evaluate("""
        () => {
            const mathElem = document.querySelector('.katex-mathml math');
            return mathElem ? mathElem.outerHTML : "";
        }
        """)
        
        browser.close()
        
        # Build the result as a RenderedEquation dataclass
        rendered_eq = RenderedEquation(
            mathml=mathml,
            spans=[
                SpanInfo(
                    text=s["text"],
                    bounding_box=BoundingBox(
                        x=s["boundingBox"]["x"],
                        y=s["boundingBox"]["y"],
                        width=s["boundingBox"]["width"],
                        height=s["boundingBox"]["height"]
                    )
                )
                for s in spans_info
            ]
        )
        
        # Save to cache (convert dataclasses to a JSON-serializable dict)
        cache_data = {
            "mathml": rendered_eq.mathml,
            "spans": [
                {
                    "text": span.text,
                    "boundingBox": {
                        "x": span.bounding_box.x,
                        "y": span.bounding_box.y,
                        "width": span.bounding_box.width,
                        "height": span.bounding_box.height
                    }
                }
                for span in rendered_eq.spans
            ]
        }
        with open(cache_file, 'w') as f:
            json.dump(cache_data, f)
        return rendered_eq

def compare_rendered_equations(reference: RenderedEquation, hypothesis: RenderedEquation) -> bool:
    """
    First, try to determine whether the normalized MathML of the hypothesis is contained
    as a substring of the normalized MathML of the reference.
    
    If that fails, then perform a neighbor‐based matching:
    Each span in the hypothesis must be matched to a span in the reference (with identical text),
    and if a hypothesis span has an immediate neighbor (up, down, left, or right), then its candidate
    reference span must have a neighbor in the same direction that (if already matched) is the candidate
    for the hypothesis neighbor – otherwise, the candidate must have the same text as the hypothesis neighbor.
    The algorithm uses backtracking to explore all possible assignments.
    """
    from bs4 import BeautifulSoup

    def extract_inner(mathml: str) -> str:
        try:
            # Use the "xml" parser so that BeautifulSoup parses MathML correctly,
            # handling HTML entities along the way.
            soup = BeautifulSoup(mathml, "xml")
            semantics = soup.find("semantics")
            if semantics:
                # Concatenate the string representation of all children except <annotation>
                inner_parts = [
                    str(child)
                    for child in semantics.contents
                    if getattr(child, "name", None) != "annotation"
                ]
                return ''.join(inner_parts)
            else:
                return str(soup)
        except Exception as e:
            print("Error parsing MathML with BeautifulSoup:", e)
            print(mathml)
            return mathml

    def normalize(s: str) -> str:
        return re.sub(r'\s+', '', s)

    # First, try a fast mathML substring check.
    reference_inner = normalize(extract_inner(reference.mathml))
    hypothesis_inner = normalize(extract_inner(hypothesis.mathml))
    if reference_inner in hypothesis_inner:
        return True

    # Fallback: neighbor-based matching using the spans.
    # First, print out the original span lists.
    # print("Hypothesis spans:")
    # for s in hypothesis.spans:
    #     print(s)
    # print("---")
    # print("Reference spans:")
    # for s in reference.spans:
    #     print(s)
    # print("---")

    # We swap H and R so that we are effectively checking if the reference is contained in the hypothesis.
    H, R = reference.spans, hypothesis.spans

    H = [span for span in H if span.text != "\u200b"]
    R = [span for span in R if span.text != "\u200b"]

    # Build candidate map: for each span in H (reference), record indices in R with matching text.
    candidate_map = {}
    for i, hspan in enumerate(H):
        candidate_map[i] = [j for j, rsp in enumerate(R) if rsp.text == hspan.text]
        if not candidate_map[i]:
            return False  # no candidate for a given span, so we fail immediately
   
    # print("Candidate Map:")
    # print(candidate_map)

    # Function to compute neighbor mappings for a list of spans.
    def compute_neighbors(spans, tol=5):
        neighbors = {}
        for i, span in enumerate(spans):
            cx = span.bounding_box.x + span.bounding_box.width / 2
            cy = span.bounding_box.y + span.bounding_box.height / 2
            up = down = left = right = None
            up_dist = down_dist = left_dist = right_dist = None
            for j, other in enumerate(spans):
                if i == j:
                    continue
                ocx = other.bounding_box.x + other.bounding_box.width / 2
                ocy = other.bounding_box.y + other.bounding_box.height / 2
                # Up: candidate must be above (ocy < cy) and nearly aligned horizontally.
                if ocy < cy and abs(ocx - cx) <= tol:
                    dist = cy - ocy
                    if up is None or dist < up_dist:
                        up = j
                        up_dist = dist
                # Down: candidate below.
                if ocy > cy and abs(ocx - cx) <= tol:
                    dist = ocy - cy
                    if down is None or dist < down_dist:
                        down = j
                        down_dist = dist
                # Left: candidate left.
                if ocx < cx and abs(ocy - cy) <= tol:
                    dist = cx - ocx
                    if left is None or dist < left_dist:
                        left = j
                        left_dist = dist
                # Right: candidate right.
                if ocx > cx and abs(ocy - cy) <= tol:
                    dist = ocx - cx
                    if right is None or dist < right_dist:
                        right = j
                        right_dist = dist
            neighbors[i] = {"up": up, "down": down, "left": left, "right": right}
        return neighbors

    hyp_neighbors = compute_neighbors(H)
    ref_neighbors = compute_neighbors(R)
    # print("Neighbor Map for Reference spans (H):")
    # for i, nb in hyp_neighbors.items():
    #     print(f"Span {i}: {nb}")
    # print("Neighbor Map for Hypothesis spans (R):")
    # for i, nb in ref_neighbors.items():
    #     print(f"Span {i}: {nb}")

    # Backtracking search for an injection f: {0,...,n-1} -> {indices in R} that preserves neighbor relations.
    n = len(H)
    used = [False] * len(R)
    assignment = {}

    def backtrack(i):
        if i == n:
            return True
        for cand in candidate_map[i]:
            if used[cand]:
                continue
            # Tentatively assign hypothesis span i (from H) to reference span cand (in R).
            assignment[i] = cand
            used[cand] = True
            valid = True
            # Check neighbor constraints for all directions.
            for direction in ["up", "down", "left", "right"]:
                hyp_nb = hyp_neighbors[i].get(direction)
                ref_nb = ref_neighbors[cand].get(direction)
                if hyp_nb is not None:
                    expected_text = H[hyp_nb].text
                    # The candidate in R must have a neighbor in that direction.
                    if ref_nb is None:
                        valid = False
                        break
                    # If the neighbor in H is already assigned, then the candidate neighbor must match.
                    if hyp_nb in assignment:
                        if assignment[hyp_nb] != ref_nb:
                            valid = False
                            break
                    else:
                        # If not yet assigned, the neighbor text in R must match the expected text.
                        if R[ref_nb].text != expected_text:
                            valid = False
                            break
            if valid:
                if backtrack(i + 1):
                    return True
            # Backtrack this candidate assignment.
            used[cand] = False
            del assignment[i]
        return False

    result = backtrack(0)

    return result


class TestRenderedEquationComparison(unittest.TestCase):
    def test_exact_match(self):
        # Both calls with identical LaTeX should produce matching MathML output.
        eq1 = render_equation("a+b", use_cache=False)
        eq2 = render_equation("a+b", use_cache=False)
        self.assertTrue(compare_rendered_equations(eq1, eq2))
    
    def test_whitespace_difference(self):
        # Differences in whitespace in the LaTeX input should not affect the MathML output.
        eq1 = render_equation("a+b", use_cache=False)
        eq2 = render_equation("a + b", use_cache=False)
        self.assertTrue(compare_rendered_equations(eq1, eq2))
    
    def test_not_found(self):
        # Completely different equations should not match.
        eq1 = render_equation("c-d", use_cache=False)
        eq2 = render_equation("a+b", use_cache=False)
        self.assertFalse(compare_rendered_equations(eq1, eq2))
    
    def test_align_block_contains_needle(self):
        # The MathML output of the plain equation should be found within the align block output.
        eq_plain = render_equation("a+b", use_cache=False)
        eq_align = render_equation("\\begin{align*}a+b\\end{align*}", use_cache=False)
        self.assertTrue(compare_rendered_equations(eq_plain, eq_align))
    
    def test_align_block_needle_not_in(self):
        # An align block rendering a different equation should not contain the MathML of an unrelated equation.
        eq_align = render_equation("\\begin{align*}a+b\\end{align*}", use_cache=False)
        eq_diff = render_equation("c-d", use_cache=False)
        self.assertFalse(compare_rendered_equations(eq_diff, eq_align))

    def test_big(self):
        ref_rendered = render_equation("\\nabla \\cdot \\mathbf{E} = \\frac{\\rho}{\\varepsilon_0}", use_cache=False, debug_dom=False)
        align_rendered = render_equation("""\\begin{align*}\\nabla \\cdot \\mathbf{E} = \\frac{\\rho}{\\varepsilon_0}\\end{align*}""", use_cache=False, debug_dom=False)
        self.assertTrue(compare_rendered_equations(ref_rendered, align_rendered))

    def test_dot_end1(self):
        ref_rendered = render_equation("\\lambda_{g}=\\sum_{s \\in S} \\zeta_{n}^{\\psi(g s)}=\\sum_{i=1}^{k}\\left[\\sum_{s, R s=\\mathcal{I}_{i}} \\zeta_{n}^{\\varphi(g s)}\\right]")
        align_rendered = render_equation("\\lambda_{g}=\\sum_{s \\in S} \\zeta_{n}^{\\psi(g s)}=\\sum_{i=1}^{k}\\left[\\sum_{s, R s=\\mathcal{I}_{i}} \\zeta_{n}^{\\varphi(g s)}\\right].")
        self.assertTrue(compare_rendered_equations(ref_rendered, align_rendered))

    def test_dot_end2(self):
        ref_rendered = render_equation("\\lambda_{g}=\\sum_{s \\in S} \\zeta_{n}^{\\psi(g s)}=\\sum_{i=1}^{k}\\left[\\sum_{s, R s=\\mathcal{I}_{i}} \\zeta_{n}^{\\psi(g s)}\\right]")
        align_rendered = render_equation("\\lambda_g = \\sum_{s \\in S} \\zeta_n^{\\psi(gs)} = \\sum_{i=1}^{k} \\left[ \\sum_{s, Rs = \\mathcal{I}_i} \\zeta_n^{\\psi(gs)} \\right]")
        self.assertTrue(compare_rendered_equations(ref_rendered, align_rendered))

    def test_lambda(self):
        ref_rendered = render_equation("\\lambda_g = \\lambda_{g'}")
        align_rendered = render_equation("\\lambda_{g}=\\lambda_{g^{\\prime}}")
        self.assertTrue(compare_rendered_equations(ref_rendered, align_rendered))

    def test_gemini(self):
        ref_rendered = render_equation("u \\in (R/\\operatorname{Ann}_R(x_i))^{\\times}")
        align_rendered = render_equation("u \\in\\left(R / \\operatorname{Ann}_{R}\\left(x_{i}\\right)\\right)^{\\times}")
        self.assertTrue(compare_rendered_equations(ref_rendered, align_rendered))
        
if __name__ == "__main__":
    unittest.main()