text-generation-webui/modules/html_generator.py

274 lines
8.1 KiB
Python
Raw Normal View History

import os
2023-01-07 03:14:08 +01:00
import re
import time
from pathlib import Path
2023-01-07 03:14:08 +01:00
2023-03-15 16:33:26 +01:00
import markdown
2023-04-05 03:52:15 +02:00
from PIL import Image, ImageOps
from modules.utils import get_available_chat_styles
# This is to store the paths to the thumbnails of the profile pictures
image_cache = {}
2023-03-15 18:19:28 +01:00
with open(Path(__file__).resolve().parent / '../css/html_readable_style.css', 'r') as f:
readable_css = f.read()
with open(Path(__file__).resolve().parent / '../css/html_4chan_style.css', 'r') as css_f:
_4chan_css = css_f.read()
2023-04-05 16:49:59 +02:00
with open(Path(__file__).resolve().parent / '../css/html_instruct_style.css', 'r') as f:
instruct_css = f.read()
2023-03-15 16:33:26 +01:00
# Custom chat styles
chat_styles = {}
for k in get_available_chat_styles():
chat_styles[k] = open(Path(f'css/chat_style-{k}.css'), 'r').read()
2023-03-17 20:06:11 +01:00
def fix_newlines(string):
string = string.replace('\n', '\n\n')
string = re.sub(r"\n{3,}", "\n\n", string)
string = string.strip()
return string
2023-04-17 02:26:19 +02:00
2023-04-16 23:00:12 +02:00
def replace_blockquote(m):
return m.group().replace('\n', '\n> ').replace('\\begin{blockquote}', '').replace('\\end{blockquote}', '')
2023-04-17 02:26:19 +02:00
2023-03-17 20:06:11 +01:00
def convert_to_markdown(string):
2023-04-16 23:00:12 +02:00
# Blockquote
pattern = re.compile(r'\\begin{blockquote}(.*?)\\end{blockquote}', re.DOTALL)
string = pattern.sub(replace_blockquote, string)
# Code
2023-03-17 20:06:11 +01:00
string = string.replace('\\begin{code}', '```')
string = string.replace('\\end{code}', '```')
string = re.sub(r"(.)```", r"\1\n```", string)
2023-04-16 23:00:12 +02:00
result = ''
is_code = False
for line in string.split('\n'):
if line.lstrip(' ').startswith('```'):
is_code = not is_code
result += line
2023-05-10 18:41:23 +02:00
if is_code or line.startswith('|'): # Don't add an extra \n for tables or code
result += '\n'
else:
result += '\n\n'
if is_code:
result = result + '```' # Unfinished code block
string = result.strip()
2023-05-10 18:41:23 +02:00
return markdown.markdown(string, extensions=['fenced_code', 'tables'])
2023-03-17 20:06:11 +01:00
def generate_basic_html(string):
string = convert_to_markdown(string)
string = f'<style>{readable_css}</style><div class="container">{string}</div>'
return string
2023-01-15 20:43:31 +01:00
2023-01-07 03:14:08 +01:00
def process_post(post, c):
t = post.split('\n')
number = t[0].split(' ')[1]
if len(t) > 1:
src = '\n'.join(t[1:])
else:
src = ''
src = re.sub('>', '&gt;', src)
src = re.sub('(&gt;&gt;[0-9]*)', '<span class="quote">\\1</span>', src)
src = re.sub('\n', '<br>\n', src)
src = f'<blockquote class="message">{src}\n'
src = f'<span class="name">Anonymous </span> <span class="number">No.{number}</span>\n{src}'
return src
2023-01-11 05:10:11 +01:00
def generate_4chan_html(f):
2023-01-07 03:14:08 +01:00
posts = []
post = ''
c = -2
for line in f.splitlines():
line += "\n"
if line == '-----\n':
continue
elif line.startswith('--- '):
c += 1
if post != '':
src = process_post(post, c)
posts.append(src)
post = line
else:
post += line
if post != '':
src = process_post(post, c)
posts.append(src)
for i in range(len(posts)):
if i == 0:
posts[i] = f'<div class="op">{posts[i]}</div>\n'
else:
posts[i] = f'<div class="reply">{posts[i]}</div>\n'
2023-01-07 03:14:08 +01:00
output = ''
2023-03-15 18:19:28 +01:00
output += f'<style>{_4chan_css}</style><div id="parent"><div id="container">'
2023-01-07 03:14:08 +01:00
for post in posts:
output += post
output += '</div></div>'
2023-01-07 03:14:08 +01:00
output = output.split('\n')
for i in range(len(output)):
2023-01-07 05:20:10 +01:00
output[i] = re.sub(r'^(&gt;(.*?)(<br>|</div>))', r'<span class="greentext">\1</span>', output[i])
output[i] = re.sub(r'^<blockquote class="message">(&gt;(.*?)(<br>|</div>))', r'<blockquote class="message"><span class="greentext">\1</span>', output[i])
2023-01-07 03:14:08 +01:00
output = '\n'.join(output)
return output
2023-04-05 04:03:58 +02:00
def make_thumbnail(image):
image = image.resize((350, round(image.size[1] / image.size[0] * 350)), Image.Resampling.LANCZOS)
2023-04-05 04:03:58 +02:00
if image.size[1] > 470:
image = ImageOps.fit(image, (350, 470), Image.ANTIALIAS)
return image
def get_image_cache(path):
cache_folder = Path("cache")
if not cache_folder.exists():
cache_folder.mkdir()
mtime = os.stat(path).st_mtime
if (path in image_cache and mtime != image_cache[path][0]) or (path not in image_cache):
2023-04-05 04:03:58 +02:00
img = make_thumbnail(Image.open(path))
output_file = Path(f'cache/{path.name}_cache.png')
img.convert('RGB').save(output_file, format='PNG')
image_cache[path] = [mtime, output_file.as_posix()]
return image_cache[path][1]
2023-04-05 16:49:59 +02:00
def generate_instruct_html(history):
output = f'<style>{instruct_css}</style><div class="chat" id="chat">'
for i, _row in enumerate(history[::-1]):
2023-04-05 16:49:59 +02:00
row = [convert_to_markdown(entry) for entry in _row]
output += f"""
<div class="assistant-message">
<div class="text">
<div class="message-body">
{row[1]}
</div>
</div>
</div>
"""
if len(row[0]) == 0: # don't display empty user messages
2023-04-05 16:49:59 +02:00
continue
output += f"""
<div class="user-message">
<div class="text">
<div class="message-body">
{row[0]}
</div>
</div>
</div>
"""
output += "</div>"
return output
def generate_cai_chat_html(history, name1, name2, style, reset_cache=False):
output = f'<style>{chat_styles[style]}</style><div class="chat" id="chat">'
2023-03-24 21:18:27 +01:00
# We use ?name2 and ?time.time() to force the browser to reset caches
img_bot = f'<img src="file/cache/pfp_character.png?{name2}">' if Path("cache/pfp_character.png").exists() else ''
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[::-1]):
2023-03-17 20:06:11 +01:00
row = [convert_to_markdown(entry) for entry in _row]
2023-03-24 21:18:27 +01:00
output += f"""
<div class="message">
<div class="circle-bot">
{img_bot}
</div>
<div class="text">
<div class="username">
{name2}
</div>
<div class="message-body">
{row[1]}
</div>
</div>
</div>
"""
if len(row[0]) == 0: # don't display empty user messages
continue
2023-03-24 21:18:27 +01:00
output += f"""
<div class="message">
<div class="circle-you">
{img_me}
</div>
<div class="text">
<div class="username">
{name1}
</div>
<div class="message-body">
{row[0]}
</div>
</div>
</div>
"""
output += "</div>"
return output
2023-04-05 16:49:59 +02:00
2023-04-16 21:44:50 +02:00
def generate_chat_html(history, name1, name2, reset_cache=False):
output = f'<style>{chat_styles["wpp"]}</style><div class="chat" id="chat">'
2023-04-16 21:44:50 +02:00
for i, _row in enumerate(history[::-1]):
row = [convert_to_markdown(entry) for entry in _row]
output += f"""
<div class="message">
<div class="text-bot">
<div class="message-body">
{row[1]}
</div>
</div>
</div>
"""
if len(row[0]) == 0: # don't display empty user messages
continue
output += f"""
<div class="message">
<div class="text-you">
<div class="message-body">
{row[0]}
</div>
</div>
</div>
"""
output += "</div>"
return output
2023-04-05 16:49:59 +02:00
def chat_html_wrapper(history, name1, name2, mode, style, reset_cache=False):
if mode == 'instruct':
2023-04-05 16:49:59 +02:00
return generate_instruct_html(history)
elif style == 'wpp':
return generate_chat_html(history, name1, name2)
2023-04-05 16:49:59 +02:00
else:
return generate_cai_chat_html(history, name1, name2, style, reset_cache)