131 lines
5.3 KiB
Python
131 lines
5.3 KiB
Python
"""Conservative repair of LaTeX escaping in raw model-generated JSON.
|
|
|
|
Call before JSON decoding, including for valid JSON: ``\\frac`` and ``\\not``
|
|
can otherwise be silently decoded as control characters. This is heuristic;
|
|
unknown commands starting with a valid JSON escape are left alone.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
|
|
# Only these content fields carry TeX. Never rewrite labels or structural keys.
|
|
_CONTENT_FIELDS = {'question_content', 'solution_content', 'context_content', 'rubric_content'}
|
|
_CONTENT_KEYS = {f'"{field}"' for field in _CONTENT_FIELDS}
|
|
_STRING = re.compile(r'"(?:\\.|[^"\\])*"', re.DOTALL)
|
|
_COMMAND = re.compile(r'[a-zA-Z]+')
|
|
_MATH = re.compile(r"(\$\$.*?\$\$|\$.*?\$|\\\(.*?\\\)|\\\[.*?\\\])", re.DOTALL)
|
|
_DECODED_COMMAND = re.compile(r"[\b\f\n\r\t]([A-Za-z]+)")
|
|
_JSON_CONTROL_PREFIX = {"\b": "b", "\f": "f", "\n": "n", "\r": "r", "\t": "t"}
|
|
|
|
|
|
def normalize_literal_newlines(text: str) -> str:
|
|
r"""Decode stray \n tokens in extracted TeX, preserving commands and \\.
|
|
|
|
An odd backslash run ends in an escaped newline; any preceding pairs
|
|
remain LaTeX row breaks. Even runs are already legitimate LaTeX.
|
|
"""
|
|
def replace(match: re.Match[str]) -> str:
|
|
slashes = match.group(1)
|
|
return slashes[:-1] + "\n" if len(slashes) % 2 else match.group(0)
|
|
|
|
return re.sub(r"(\\+)n(?![A-Za-z])", replace, text)
|
|
# Commands whose first letter is also a valid JSON escape. Other commands do
|
|
# not need a dictionary: their single backslash is invalid JSON already.
|
|
_AMBIGUOUS_COMMANDS = set('''
|
|
begin beta bar backslash binom boxed boldsymbol big bigl bigr biggl biggr
|
|
bigskip blacktriangle blacksquare bmod breve bullet bra braket
|
|
frac fbox forall footnote footnotesize framebox
|
|
not neq neg nabla newline newcommand newenvironment noindent nonumber
|
|
nexists notin nmid nparallel nsubseteq nsupseteq normalfont normalsize
|
|
rightarrow right leftarrow rangle rbrace rceil rfloor rightarrowtail
|
|
Rightarrow rightharpoonup rightharpoondown rightleftharpoons rightleftarrows
|
|
rm rule ref rho rightskip raisebox require renewcommand
|
|
text textbf textit textrm textsf texttt textnormal textcolor textsc textsl
|
|
textup textwidth theta tau tan tanh times to top triangle triangleleft
|
|
triangleright tilde tfrac today tag tiny thickspace thinspace tbinom
|
|
underbrace underline underset uparrow updownarrow upsilon usepackage
|
|
'''.split())
|
|
|
|
|
|
def repair_decoded_latex_controls(text: str) -> str:
|
|
r"""Restore known TeX commands after JSON decoded their first letter.
|
|
|
|
Only math spans and whitelisted command names are changed. This avoids
|
|
treating ordinary paragraph newlines or tabs as damaged LaTeX.
|
|
"""
|
|
def repair_math(math_match: re.Match[str]) -> str:
|
|
def repair_command(command_match: re.Match[str]) -> str:
|
|
candidate = (
|
|
_JSON_CONTROL_PREFIX[command_match.group(0)[0]]
|
|
+ command_match.group(1)
|
|
)
|
|
return (
|
|
f"\\{candidate}"
|
|
if candidate in _AMBIGUOUS_COMMANDS
|
|
else command_match.group(0)
|
|
)
|
|
|
|
return _DECODED_COMMAND.sub(repair_command, math_match.group(0))
|
|
|
|
return _MATH.sub(repair_math, text)
|
|
|
|
|
|
def _repair_string(token: str) -> str:
|
|
output = []
|
|
index = 0
|
|
while index < len(token):
|
|
char = token[index]
|
|
if char != '\\' or index + 1 == len(token):
|
|
output.append(char)
|
|
index += 1
|
|
continue
|
|
following = token[index + 1]
|
|
# Consume already escaped pairs together, never double-escape them.
|
|
if following in '\\"/':
|
|
output.append(token[index:index + 2])
|
|
index += 2
|
|
continue
|
|
command_match = _COMMAND.match(token, index + 1)
|
|
command = command_match.group() if command_match else ''
|
|
unicode_escape = following == 'u' and re.match(r'[0-9a-fA-F]{4}', token[index + 2:])
|
|
if not unicode_escape and (following not in 'bfnrtu' or command in _AMBIGUOUS_COMMANDS
|
|
or following == 'u'):
|
|
output.append('\\\\')
|
|
else:
|
|
output.append('\\')
|
|
index += 1
|
|
return ''.join(output)
|
|
|
|
|
|
def repair_latex_json(
|
|
raw: str, *, additional_content_fields: set[str] | None = None
|
|
) -> str:
|
|
"""Repair single TeX backslashes without decoding or rewriting other JSON.
|
|
|
|
Preserve standard escapes, Unicode escapes, and already escaped TeX.
|
|
A literal newline followed by e.g. 'ot' has the same wire representation
|
|
as a broken ``\\not`` command; known TeX command names take precedence.
|
|
Raw responses/caches should be retained to allow reviewing such cases.
|
|
"""
|
|
additional_content_keys = {
|
|
f'"{field}"' for field in (additional_content_fields or set())
|
|
}
|
|
output = []
|
|
end = 0
|
|
previous = None
|
|
for match in _STRING.finditer(raw):
|
|
token = match.group()
|
|
if previous is not None and raw[previous.end():match.start()].strip() == ':':
|
|
if previous.group() in _CONTENT_KEYS:
|
|
token = _repair_string(token)
|
|
elif previous.group() in additional_content_keys:
|
|
token = _MATH.sub(
|
|
lambda math_match: _repair_string(math_match.group(0)),
|
|
token,
|
|
)
|
|
output.extend((raw[end:match.start()], token))
|
|
end = match.end()
|
|
previous = match
|
|
output.append(raw[end:])
|
|
return ''.join(output)
|