From edec3bf3b0dadaf78342479659e0c7711eab347b Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Wed, 13 Mar 2024 08:14:34 -0700
Subject: [PATCH] UI: avoid caching convert_to_markdown calls during streaming

---
 modules/html_generator.py | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/modules/html_generator.py b/modules/html_generator.py
index 86771dcf..902ffcc7 100644
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -101,6 +101,17 @@ def convert_to_markdown(string):
     return html_output
 
 
+def convert_to_markdown_wrapped(string, use_cache=True):
+    '''
+    Used to avoid caching convert_to_markdown calls during streaming.
+    '''
+
+    if use_cache:
+        return convert_to_markdown(string)
+
+    return convert_to_markdown.__wrapped__(string)
+
+
 def generate_basic_html(string):
     string = convert_to_markdown(string)
     string = f'<style>{readable_css}</style><div class="readable-container">{string}</div>'
@@ -196,7 +207,7 @@ def get_image_cache(path):
 def generate_instruct_html(history):
     output = f'<style>{instruct_css}</style><div class="chat" id="chat"><div class="messages">'
     for i, _row in enumerate(history):
-        row = [convert_to_markdown(entry) for entry in _row]
+        row = [convert_to_markdown_wrapped(entry, use_cache=i != len(history) - 1) for entry in _row]
 
         if row[0]:  # don't display empty user messages
             output += f"""
@@ -232,7 +243,7 @@ def generate_cai_chat_html(history, name1, name2, style, character, reset_cache=
     img_me = f'<img src="file/cache/pfp_me.png?{time.time() if reset_cache else ""}">' if Path("cache/pfp_me.png").exists() else ''
 
     for i, _row in enumerate(history):
-        row = [convert_to_markdown(entry) for entry in _row]
+        row = [convert_to_markdown_wrapped(entry, use_cache=i != len(history) - 1) for entry in _row]
 
         if row[0]:  # don't display empty user messages
             output += f"""
@@ -275,7 +286,7 @@ def generate_chat_html(history, name1, name2, reset_cache=False):
     output = f'<style>{chat_styles["wpp"]}</style><div class="chat" id="chat"><div class="messages">'
 
     for i, _row in enumerate(history):
-        row = [convert_to_markdown(entry) for entry in _row]
+        row = [convert_to_markdown_wrapped(entry, use_cache=i != len(history) - 1) for entry in _row]
 
         if row[0]:  # don't display empty user messages
             output += f"""