95 lines
4.0 KiB
Python
95 lines
4.0 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]+')
|
|
|
|
|
|
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_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) -> 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.
|
|
"""
|
|
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() == ':'
|
|
and previous.group() in _CONTENT_KEYS):
|
|
token = _repair_string(token)
|
|
output.extend((raw[end:match.start()], token))
|
|
end = match.end()
|
|
previous = match
|
|
output.append(raw[end:])
|
|
return ''.join(output)
|