From a0492ce325b951a9c000fa3cad45806adc8d8926 Mon Sep 17 00:00:00 2001
From: oobabooga <oobabooga4@gmail.com>
Date: Sat, 11 Jan 2025 21:14:10 -0300
Subject: [PATCH] Optimize syntax highlighting during chat streaming (#6655)

---
 js/global_scope_js.js |  2 +-
 js/main.js            | 55 +++++++++++++++----------------------------
 modules/ui_chat.py    | 40 ++++++++++++++++++-------------
 3 files changed, 44 insertions(+), 53 deletions(-)

diff --git a/js/global_scope_js.js b/js/global_scope_js.js
index 79b673d7..983d60f1 100644
--- a/js/global_scope_js.js
+++ b/js/global_scope_js.js
@@ -19,5 +19,5 @@ function copyToClipboard(element) {
 }
 
 function regenerateClick() {
-    document.getElementById("Regenerate").click();
+  document.getElementById("Regenerate").click();
 }
diff --git a/js/main.js b/js/main.js
index ab2499d4..c5c47d04 100644
--- a/js/main.js
+++ b/js/main.js
@@ -177,47 +177,30 @@ function isElementVisibleOnScreen(element) {
   );
 }
 
-function getVisibleMessagesIndexes() {
-  const elements = document.querySelectorAll(".message-body");
-  const visibleIndexes = [];
-
-  elements.forEach((element, index) => {
-    if (isElementVisibleOnScreen(element) && !element.hasAttribute("data-highlighted")) {
-      visibleIndexes.push(index);
-    }
-  });
-
-  return visibleIndexes;
-}
-
 function doSyntaxHighlighting() {
-  const indexes = getVisibleMessagesIndexes();
-  const elements = document.querySelectorAll(".message-body");
+  const messageBodies = document.querySelectorAll(".message-body");
 
-  if (indexes.length > 0) {
+  if (messageBodies.length > 0) {
     observer.disconnect();
 
-    indexes.forEach((index) => {
-      const element = elements[index];
+    messageBodies.forEach((messageBody) => {
+      if (isElementVisibleOnScreen(messageBody)) {
+        // Handle both code and math in a single pass through each message
+        const codeBlocks = messageBody.querySelectorAll("pre code:not([data-highlighted])");
+        codeBlocks.forEach((codeBlock) => {
+          hljs.highlightElement(codeBlock);
+          codeBlock.setAttribute("data-highlighted", "true");
+        });
 
-      // Tag this element to prevent it from being highlighted twice
-      element.setAttribute("data-highlighted", "true");
-
-      // Perform syntax highlighting
-      const codeBlocks = element.querySelectorAll("pre code");
-
-      codeBlocks.forEach((codeBlock) => {
-        hljs.highlightElement(codeBlock);
-      });
-
-      renderMathInElement(element, {
-        delimiters: [
-          { left: "$$", right: "$$", display: true },
-          { left: "$", right: "$", display: false },
-          { left: "\\(", right: "\\)", display: false },
-          { left: "\\[", right: "\\]", display: true },
-        ],
-      });
+        renderMathInElement(messageBody, {
+          delimiters: [
+            { left: "$$", right: "$$", display: true },
+            { left: "$", right: "$", display: false },
+            { left: "\\(", right: "\\)", display: false },
+            { left: "\\[", right: "\\]", display: true },
+          ],
+        });
+      }
     });
 
     observer.observe(targetElement, config);
diff --git a/modules/ui_chat.py b/modules/ui_chat.py
index 61be17e3..8497f7df 100644
--- a/modules/ui_chat.py
+++ b/modules/ui_chat.py
@@ -182,23 +182,31 @@ def create_event_handlers():
 
     # Morph HTML updates instead of updating everything
     shared.gradio['display'].change(None, gradio('display'), None,
-        js="""
-            (text) => {
-                morphdom(
-                    document.getElementById('chat').parentNode,
-                    '<div class="prose svelte-1ybaih5">' + text + '</div>',
-                    {
-                        onBeforeElUpdated: function(fromEl, toEl) {
-                            if (fromEl.isEqualNode(toEl)) {
-                                return false; // Skip identical nodes
-                            }
-                            return true; // Update only if nodes differ
-                        }
-                    }
-                );
+      js="""
+      (text) => {
+        morphdom(
+          document.getElementById('chat').parentNode,
+          '<div class="prose svelte-1ybaih5">' + text + '</div>',
+          {
+            onBeforeElUpdated: function(fromEl, toEl) {
+              if (fromEl.tagName === 'PRE' && fromEl.querySelector('code[data-highlighted]')) {
+                const fromCode = fromEl.querySelector('code');
+                const toCode = toEl.querySelector('code');
+
+                if (fromCode && toCode && fromCode.textContent === toCode.textContent) {
+                  // If the <code> content is the same, preserve the entire <pre> element
+                  toEl.className = fromEl.className;
+                  toEl.innerHTML = fromEl.innerHTML;
+                  return false; // Skip updating the <pre> element
+                }
+              }
+              return !fromEl.isEqualNode(toEl); // Update only if nodes differ
             }
-        """
-    )
+          }
+        );
+      }
+      """
+    );
 
     shared.gradio['Generate'].click(
         ui.gather_interface_values, gradio(shared.input_elements), gradio('interface_state')).then(