160 lines
5.1 KiB
Python
160 lines
5.1 KiB
Python
"""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,
|
|
)
|
|
|
|
# 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)
|