Fix LaTeX display and add quality-of-life improvements

- Fix display math mis-parsing: convert \[...\] to block $$\ncontent\n$$
form so arithmatex correctly identifies display math instead of treating
it as inline math surrounded by literal $ signs
- Fix MathJax typeset in Gradio shadow DOM: walk shadow roots to find
.math-preview instead of relying on document.querySelector which fails
inside Gradio's shadow DOM isolation; add 2s fallback for GPU cold-starts
- Add MathJax mathtools extension so \coloneqq / \eqqcolon render natively
without requiring text substitutions
- Text tab: convert \[...\] → $$...$$ and \(...\) → $...$ for readability
- Add Download Markdown button that appears after each extraction

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

app.py

@@ -110,14 +110,14 @@ def clean_output(text, include_images=False):
         else:
             text = re.sub(rf'(?m)^[^\n]*{re.escape(match[0])}[^\n]*\n?', '', text)
     
-    text = text.replace('\\coloneqq', ':=').replace('\\eqqcolon', '=:')
-    
     return text.strip()
 
 MATHJAX_HEAD = """
 <script>
 window.MathJax = {
+  loader: {load: ['[tex]/mathtools']},
   tex: {
+    packages:    {'[+]': ['mathtools']},
     inlineMath:  [['\\\\(', '\\\\)']],
     displayMath: [['\\\\[', '\\\\]']],
     processEscapes: true,
@@ -169,8 +169,16 @@ def to_math_html(text):
     # Pre-convert \[...\] and \(...\) to $$...$$ and $...$
     # Markdown strips backslashes before arithmatex can protect them,
     # so convert to $-delimiters first (arithmatex recognises those).
-    text = re.sub(r'\\\[(.+?)\\\]', r'$$\1$$', text, flags=re.DOTALL)
-    text = re.sub(r'\\\((.+?)\\\)', r'$\1$', text)
+    #
+    # IMPORTANT: the model outputs " \[ content \] " with surrounding spaces.
+    # Naively replacing to "$$ content $$" (space after $$) causes arithmatex
+    # to mis-parse as inline math surrounded by literal $ signs, which is why
+    # equations appeared as "$ ... $" instead of rendered display math.
+    # Fix: use the multi-line block form ($$\ncontent\n$$) and strip whitespace.
+    text = re.sub(r'\\\[(.+?)\\\]',
+                  lambda m: f'\n\n$$\n{m.group(1).strip()}\n$$\n\n',
+                  text, flags=re.DOTALL)
+    text = re.sub(r'\\\((.+?)\\\)', lambda m: f'${m.group(1).strip()}$', text)
     html = md_lib.markdown(text, extensions=[
         'pymdownx.arithmatex',
         'tables',
@@ -354,6 +362,7 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
                     gallery = gr.Gallery(show_label=False, columns=3, height=400)
                 with gr.Tab("Raw Text", id="tab_raw"):
                     raw_out = gr.Textbox(lines=20, buttons=["copy"], show_label=False)
+            download_btn = gr.DownloadButton("Download Markdown", visible=False, variant="secondary")
     
     with gr.Accordion("Image Examples", open=True):
         gr.Examples(
@@ -406,21 +415,48 @@ with gr.Blocks(title="DeepSeek-OCR-2") as demo:
         elif image is not None:
             cleaned, markdown, raw, img_out, crops = process_image(image, task, custom_prompt)
         else:
-            return "Error: Upload a file or image", "", "", None, []
-        return cleaned, to_math_html(markdown), raw, img_out, crops
+            return "Error: Upload a file or image", "", "", None, [], gr.DownloadButton(visible=False)
+
+        # Text tab: convert \[...\] → $$...$$ and \(...\) → $...$ for readability
+        text_display = re.sub(r'\\\[(.+?)\\\]',
+                              lambda m: f'\n$$\n{m.group(1).strip()}\n$$\n',
+                              cleaned, flags=re.DOTALL)
+        text_display = re.sub(r'\\\((.+?)\\\)', lambda m: f'${m.group(1).strip()}$', text_display)
+
+        # Download file: write cleaned markdown to a temp .md file
+        dl_tmp = tempfile.NamedTemporaryFile(delete=False, suffix='.md', mode='w', encoding='utf-8')
+        dl_tmp.write(cleaned)
+        dl_tmp.close()
+
+        return (text_display, to_math_html(markdown), raw, img_out, crops,
+                gr.DownloadButton(value=dl_tmp.name, visible=True))
 
     submit_event = btn.click(run, [input_img, file_in, task, prompt, page_selector],
-                             [text_out, md_out, raw_out, img_out, gallery])
+                             [text_out, md_out, raw_out, img_out, gallery, download_btn])
     submit_event.then(select_boxes, [task], [tabs])
     submit_event.then(fn=None, js="""() => {
       const tryTypeset = () => {
         if (!window.MathJax || !MathJax.typesetPromise) { setTimeout(tryTypeset, 100); return; }
-        const el = document.querySelector('.math-preview');
-        if (!el) return;
+        // Gradio renders gr.HTML inside a shadow DOM, so document.querySelector
+        // won't find .math-preview. Walk all shadow roots to locate it.
+        const findInShadows = (root) => {
+          const el = root.querySelector('.math-preview');
+          if (el) return el;
+          for (const node of root.querySelectorAll('*')) {
+            if (node.shadowRoot) {
+              const found = findInShadows(node.shadowRoot);
+              if (found) return found;
+            }
+          }
+          return null;
+        };
+        const el = findInShadows(document);
+        if (!el) { MathJax.typesetPromise(); return; }
         MathJax.typesetClear([el]);
         MathJax.typesetPromise([el]);
       };
-      setTimeout(tryTypeset, 100);
+      setTimeout(tryTypeset, 200);   // primary: catches most cases
+      setTimeout(tryTypeset, 2000);  // fallback: GPU cold-start can delay DOM update
     }""")
 
 if __name__ == "__main__":