improve latex parsing

app.py

@@ -133,34 +133,43 @@ def update_inputs(task):
         ]
     
 def parse_latex_output(res):
-    lines = res.split('\n')
+    # Split the input, preserving newlines and empty lines
+    lines = re.split(r'(\$\$.*?\$\$)', res, flags=re.DOTALL)
     parsed_lines = []
     in_latex = False
-    
+    latex_buffer = []
+
     for line in lines:
-        line = line.strip()
-        if not line:
+        if line == '\n':
+            if in_latex:
+                latex_buffer.append(line)
+            else:
+                parsed_lines.append(line)
             continue
+
+        line = line.strip()
         
-        latex_patterns = [r'\{', r'\}', r'\[', r'\]', r'\\', r'\$', r'_', r'^']
+        latex_patterns = [r'\{', r'\}', r'\[', r'\]', r'\\', r'\$', r'_', r'^', r'"']
         contains_latex = any(re.search(pattern, line) for pattern in latex_patterns)
         
-        is_key_value = ':' in line and not line.startswith('{') and not line.endswith('}')
-        
-        if contains_latex or is_key_value:
+        if contains_latex:
             if not in_latex:
-                parsed_lines.append('$$')
                 in_latex = True
-            parsed_lines.append(line)
+                latex_buffer = ['$$']
+            latex_buffer.append(line)
         else:
             if in_latex:
-                parsed_lines.append('$$')
+                latex_buffer.append('$$')
+                parsed_lines.extend(latex_buffer)
                 in_latex = False
+                latex_buffer = []
             parsed_lines.append(line)
+
     if in_latex:
-        parsed_lines.append('$$')
-    
-    return '\n'.join(parsed_lines)
+        latex_buffer.append('$$')
+        parsed_lines.extend(latex_buffer)
+
+    return '$$\n$$'.join(parsed_lines)
                          
 def ocr_demo(image, task, ocr_type, ocr_box, ocr_color):
     res, html_content, unique_id = process_image(image, task, ocr_type, ocr_box, ocr_color)