"""LaTeX parsing utilities. Provides helpers to extract sections, chunk text, and convert LaTeX to plain readable prose by stripping markup that is not useful for NLP/embedding. """ from __future__ import annotations import re # --------------------------------------------------------------------------- # Internal helpers # --------------------------------------------------------------------------- _SECTION_RE = re.compile( r"\\(?:sub)*section\*?\s*\{([^}]*)\}", re.DOTALL, ) # \input{file} / \include{file} — arXiv papers routinely split the body across # files, so these must be inlined before parsing (R-C multi-file flattening). _INPUT_RE = re.compile(r"\\(?:input|include)\s*\{([^}]+)\}") # Patterns for LaTeX noise removal (applied in order). _COMMENT_RE = re.compile(r"%[^\n]*") _CITE_RE = re.compile(r"\\cite\*?\{[^}]*\}") _LABEL_RE = re.compile(r"\\label\{[^}]*\}") _REF_RE = re.compile(r"\\ref\{[^}]*\}") # Display math: $$...$$ (before single $ to avoid greedy mismatch) _DISPLAY_DOLLAR_RE = re.compile(r"\$\$.*?\$\$", re.DOTALL) # Inline math: $...$ _INLINE_MATH_RE = re.compile(r"\$[^$\n]*?\$", re.DOTALL) # \begin{equation}...\end{equation} _ENV_EQUATION_RE = re.compile( r"\\begin\{equation\*?\}.*?\\end\{equation\*?\}", re.DOTALL, ) # \begin{align}...\end{align} (covers align, align*, aligned, …) _ENV_ALIGN_RE = re.compile( r"\\begin\{align[^}]*\}.*?\\end\{align[^}]*\}", re.DOTALL, ) # Collapse excess whitespace _WHITESPACE_RE = re.compile(r"[ \t]+") _BLANK_LINES_RE = re.compile(r"\n{3,}") def _clean_latex(text: str) -> str: """Strip LaTeX markup and return readable prose.""" text = _COMMENT_RE.sub("", text) text = _ENV_EQUATION_RE.sub("", text) text = _ENV_ALIGN_RE.sub("", text) text = _DISPLAY_DOLLAR_RE.sub("", text) text = _INLINE_MATH_RE.sub("", text) text = _CITE_RE.sub("", text) text = _LABEL_RE.sub("", text) text = _REF_RE.sub("", text) text = _WHITESPACE_RE.sub(" ", text) text = _BLANK_LINES_RE.sub("\n\n", text) return text.strip() # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- def extract_sections(latex: str) -> list[tuple[str, str]]: """Split *latex* on \\section / \\subsection boundaries. Returns a list of ``(title, cleaned_text)`` tuples, one per section. Text between the document start and the first section command is discarded (preamble / abstract handling is out of scope). """ # Find all section positions matches = list(_SECTION_RE.finditer(latex)) if not matches: return [] sections: list[tuple[str, str]] = [] for idx, match in enumerate(matches): title = match.group(1).strip() body_start = match.end() body_end = matches[idx + 1].start() if idx + 1 < len(matches) else len(latex) body = latex[body_start:body_end] cleaned = _clean_latex(body) sections.append((title, cleaned)) return sections def chunk_text(text: str, size: int = 512, overlap: int = 64) -> list[str]: """Split *text* into overlapping word-based chunks. Parameters ---------- text: Plain text to chunk (not raw LaTeX). size: Target number of words per chunk. overlap: Number of words to carry over into the next chunk. Each chunk is at most ``size + overlap`` words. Boundaries are snapped to the nearest ``". "`` (sentence end) within ±20 words of the nominal boundary when possible. """ words = text.split() if not words: return [] snap_window = 20 chunks: list[str] = [] start = 0 while start < len(words): end = min(start + size, len(words)) # Try to snap *end* to a sentence boundary within ±snap_window words. if end < len(words): best = end # Build a small search window lo = max(start + 1, end - snap_window) hi = min(len(words), end + snap_window + 1) # Prefer the closest sentence-end ". " to *end* for offset in range(0, snap_window + 1): for candidate in (end - offset, end + offset): if ( lo <= candidate < hi and candidate > start and words[candidate - 1].endswith(".") ): # Sentence boundary: word at (candidate-1) ends with ".". best = candidate break if best != end: break end = best chunk_words = words[start:end] chunks.append(" ".join(chunk_words)) if end >= len(words): break # Next chunk starts *overlap* words before *end*. start = max(start + 1, end - overlap) return chunks def latex_to_text(latex: str) -> str: """Convert a LaTeX document to plain text. Extracts all sections, cleans each one, and joins them with a blank line. If no section commands are found, the whole document is cleaned and returned as a single block. """ sections = extract_sections(latex) if sections: return "\n\n".join(body for _, body in sections) # Fallback: no section structure — clean the whole string. return _clean_latex(latex) def _norm_texkey(name: str) -> str: """Normalize a tex file name / ``\\input`` argument to a comparable key. Strips a leading ``./``, any directory-irrelevant whitespace, and a trailing ``.tex`` so ``\\input{sections/intro}`` matches the archive member ``sections/intro.tex``. """ n = name.strip() if n.startswith("./"): # a single leading "./" only — don't eat "../" or ".hidden" n = n[2:] return n[:-4] if n.endswith(".tex") else n def flatten_inputs(primary: str, files: dict[str, str]) -> str: """Recursively inline ``\\input``/``\\include`` directives (R-C). *files* maps each source file's archive name to its contents (keys may carry a ``.tex`` suffix or a ``./`` prefix — they are normalized internally). arXiv multi-file projects leave the ``\\documentclass`` file as little more than a skeleton of ``\\input`` lines, so the body must be assembled before parsing. Unresolved references are dropped; recursion is capped at depth 10 to guard against include cycles. """ norm = {_norm_texkey(k): v for k, v in files.items()} def _expand(text: str, depth: int) -> str: if depth > 10: # Cap reached (likely an \input cycle): strip any still-unresolved # directives rather than leak a literal "\input{…}" into the parsed text. return _INPUT_RE.sub("", text) def _repl(match: re.Match[str]) -> str: content = norm.get(_norm_texkey(match.group(1))) return _expand(content, depth + 1) if content is not None else "" return _INPUT_RE.sub(_repl, text) return _expand(primary, 0) def chunk_sections(latex: str, size: int = 512, overlap: int = 64) -> list[tuple[str | None, str]]: """Chunk LaTeX *within* each section, pairing every chunk with its title. Unlike :func:`latex_to_text` (which discards titles and flattens the whole document into one block before word-window chunking), this keeps chunks from spanning section boundaries and remembers which ``\\section`` each came from. The caller can then label the stored ``section`` column from the real heading instead of mostly ``body`` (R-C / audit R-12). Falls back to title-less chunks of the whole cleaned document when there are no section commands. """ sections = extract_sections(latex) if not sections: return [(None, chunk) for chunk in chunk_text(_clean_latex(latex), size, overlap)] pairs: list[tuple[str | None, str]] = [] for title, body in sections: for chunk in chunk_text(body, size, overlap): pairs.append((title or None, chunk)) return pairs