Spaces:
Sleeping
Sleeping
File size: 2,143 Bytes
2987995 7e8815a 2987995 7e8815a 2987995 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | import ast
import pathlib
import re
import unittest
def _load_clean_output():
app_path = pathlib.Path(__file__).resolve().parents[1] / "app.py"
source = app_path.read_text(encoding="utf-8")
module = ast.parse(source, filename=str(app_path))
wanted = {
"_equation_text_key",
"_dedupe_repeated_math_blocks",
"_strip_malformed_grounding",
"clean_output",
}
fn_nodes = [n for n in module.body if isinstance(n, ast.FunctionDef) and n.name in wanted]
fn_nodes.sort(key=lambda n: n.lineno)
test_mod = ast.Module(body=fn_nodes, type_ignores=[])
code = compile(test_mod, filename=str(app_path), mode="exec")
scope = {
"re": re,
}
exec(code, scope)
return scope["clean_output"]
class CleanOutputTests(unittest.TestCase):
def test_removes_truncated_grounding_artifact_line(self):
clean_output = _load_clean_output()
raw = (
"\\[ \\frac{18x-34}{(2x-3)^2} \\]\n"
"<|ref|>equation<|/ref|><|det|>[[50, 0, 450, 100]]\n"
)
cleaned = clean_output(raw, include_images=True)
self.assertIn("\\[ \\frac{18x-34}{(2x-3)^2} \\]", cleaned)
self.assertNotIn("<|ref|>", cleaned)
self.assertNotIn("<|det|>", cleaned)
self.assertNotIn("[[50, 0, 450, 100]]", cleaned)
def test_replaces_full_image_reference(self):
clean_output = _load_clean_output()
raw = "prefix\n<|ref|>image<|/ref|><|det|>[[0,0,100,100]]<|/det|>\nsuffix"
cleaned = clean_output(raw, include_images=True)
self.assertIn("**[Figure 1]**", cleaned)
self.assertNotIn("<|ref|>", cleaned)
self.assertNotIn("<|det|>", cleaned)
def test_dedupes_equivalent_math_blocks(self):
clean_output = _load_clean_output()
raw = (
"\\[ \\frac{(18x-27-7)}{(2x-3)^2}=\\frac{18x-34}{(2x-3)^2} \\]\n"
"\\[ \\frac{(18x-27-7)}{(2x-3)^{2}}=\\frac{18x-34}{(2x-3)^{2}} \\]\n"
)
cleaned = clean_output(raw, include_images=True)
self.assertEqual(1, cleaned.count("\\["))
if __name__ == "__main__":
unittest.main()
|