From a74dd9003f0f06395cceaf1565618c6f6f70bf61 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Sun, 20 Aug 2023 21:40:22 -0700
Subject: [PATCH] Fix HTML escaping for perplexity_colors extension

---
 modules/html_generator.py  | 1 -
 modules/text_generation.py | 8 ++++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/modules/html_generator.py b/modules/html_generator.py
index 0026ba4e..86f31f3a 100644
--- a/modules/html_generator.py
+++ b/modules/html_generator.py
@@ -89,7 +89,6 @@ def convert_to_markdown(string):
 
 
 def generate_basic_html(string):
-    string = html.escape(string)
     string = convert_to_markdown(string)
     string = f'<style>{readable_css}</style><div class="container">{string}</div>'
     return string
diff --git a/modules/text_generation.py b/modules/text_generation.py
index 30e81355..efafd67a 100644
--- a/modules/text_generation.py
+++ b/modules/text_generation.py
@@ -1,5 +1,6 @@
 import ast
 import copy
+import html
 import random
 import re
 import time
@@ -31,7 +32,7 @@ def generate_reply(*args, **kwargs):
         shared.generation_lock.release()
 
 
-def _generate_reply(question, state, stopping_strings=None, is_chat=False):
+def _generate_reply(question, state, stopping_strings=None, is_chat=False, escape_html=False):
 
     # Find the appropriate generation function
     generate_func = apply_extensions('custom_generate_reply')
@@ -73,6 +74,9 @@ def _generate_reply(question, state, stopping_strings=None, is_chat=False):
 
     # Generate
     for reply in generate_func(question, original_question, seed, state, stopping_strings, is_chat=is_chat):
+        if escape_html:
+            reply = html.escape(reply)
+
         reply, stop_found = apply_stopping_strings(reply, all_stop_strings)
         if is_stream:
             cur_time = time.time()
@@ -138,7 +142,7 @@ def generate_reply_wrapper(question, state, stopping_strings=None):
     reply = question if not shared.is_seq2seq else ''
     yield formatted_outputs(reply, shared.model_name)
 
-    for reply in generate_reply(question, state, stopping_strings, is_chat=False):
+    for reply in generate_reply(question, state, stopping_strings, is_chat=False, escape_html=True):
         if not shared.is_seq2seq:
             reply = question + reply