diff --git a/extensions/perplexity_colors/script.py b/extensions/perplexity_colors/script.py index 2a986ac4..c5fa724a 100644 --- a/extensions/perplexity_colors/script.py +++ b/extensions/perplexity_colors/script.py @@ -1,11 +1,10 @@ -import time - import gradio import numpy as np import torch from transformers import LogitsProcessor +import colorsys -from modules import html_generator, shared +from modules import shared params = { 'active': True, @@ -81,6 +80,7 @@ class PerplexityLogits(LogitsProcessor): # Stores the perplexity and top probabilities +# global ppl_logits_processor ppl_logits_processor = None @@ -109,64 +109,74 @@ def output_modifier(text): gen_tokens = [shared.tokenizer.decode(token_id) for token_id in gen_token_ids] sel_probs = ppl_logits_processor.selected_probs[1:] - end_part = '' if params['probability_dropdown'] else '' # Helps with finding the index after replacing part of the text. + formatted_text = "" - i = 0 for token, prob, ppl, top_tokens, top_probs in zip(gen_tokens, sel_probs, perplexities, top_tokens_list, top_probs_list): color = 'ffffff' - if params['color_by_probability'] and params['color_by_perplexity']: - color = probability_perplexity_color_scale(prob, ppl) - elif params['color_by_perplexity']: + # if params['color_by_probability'] and params['color_by_perplexity']: + # color = probability_perplexity_color_scale(prob, ppl) + if params['color_by_perplexity']: color = perplexity_color_scale(ppl) elif params['color_by_probability']: color = probability_color_scale(prob) - if token in text[i:]: - if params['probability_dropdown']: - text = text[:i] + text[i:].replace(token, add_dropdown_html(token, color, top_tokens, top_probs[0], ppl), 1) - else: - text = text[:i] + text[i:].replace(token, add_color_html(token, color), 1) - i += text[i:].find(end_part) + len(end_part) + + if params['probability_dropdown']: + formatted_token = add_dropdown_html(token, color, top_tokens, top_probs[0], ppl) + else: + formatted_token = add_color_html(token, color) + + formatted_text += formatted_token + + # wonky fix to prevent the later html processing from making the first token into a separate paragraph and disabling its dropdown. + # this adds a dummy entry which gets made into an empty

paragraph instead of the first token + if params['probability_dropdown']: + formatted_text = '' + formatted_text # Use full perplexity list for calculating the average here. print('Average perplexity:', round(np.mean(ppl_logits_processor.perplexities_list[:-1]), 4)) # t1 = time.time() # print(f"Modifier: {(t1-t0):.3f} s") # About 50 ms - return text + return formatted_text def probability_color_scale(prob): ''' Green-yellow-red color scale ''' + # hue (0.0 = red, 0.33 = green) + # saturation (0.0 = gray / white, 1.0 = normal, just leave at 1.0) + # brightness (0.0 = black, 1.0 = brightest, use 0.5 instead of 1.0 for better readability) - rv = 0 - gv = 0 - if prob <= 0.5: - rv = 'ff' - gv = hex(int(255 * prob * 2))[2:] - if len(gv) < 2: - gv = '0' * (2 - len(gv)) + gv - else: - rv = hex(int(255 - 255 * (prob - 0.5) * 2))[2:] - gv = 'ff' - if len(rv) < 2: - rv = '0' * (2 - len(rv)) + rv + hue = prob * 0.33 + rv, gv, bv = colorsys.hsv_to_rgb(hue, 1.0, 0.7) - return rv + gv + '00' + # to hex + hex_col = f"{int(rv*255):02x}{int(gv*255):02x}{int(bv*255):02x}" + + return hex_col def perplexity_color_scale(ppl): ''' Red component only, white for 0 perplexity (sorry if you're not in dark mode) ''' - value = hex(max(int(255.0 - params['ppl_scale'] * (float(ppl) - 1.0)), 0))[2:] - if len(value) < 2: - value = '0' * (2 - len(value)) + value + # hue (0.0 = red) + # saturation (1.0 = red) + # brightness (0.0 = black, 1.0 = red) + # scale saturation from white to red the higher the perplexity - return 'ff' + value + value + ppl = min(ppl, params['ppl_scale']) # clip ppl to 0-params['ppl_scale'] for color scaling. 15 should be fine for clipping and scaling + sat = ppl / params['ppl_scale'] + rv, gv, bv = colorsys.hsv_to_rgb(0.0, sat, 1.0) + + # to hex + hex_col = f"{int(rv*255):02x}{int(gv*255):02x}{int(bv*255):02x}" + + return hex_col +# are you sure this is needed? Since perplexity is proportional to probability, it would make more sense to just pick one of the two def probability_perplexity_color_scale(prob, ppl): ''' Green-yellow-red for probability and blue component for perplexity @@ -193,7 +203,7 @@ def probability_perplexity_color_scale(prob, ppl): def add_color_html(token, color): - return f'{token}' + return f'{token}' # TODO: Major issue: Applying this to too many tokens will cause a permanent slowdown in generation speed until the messages are removed from the history. @@ -201,7 +211,15 @@ def add_color_html(token, color): # Potential solution is maybe to modify the main generation code to send just the internal text and not the visible history, to avoid moving too much around. # I wonder if we can also avoid using deepcopy here. def add_dropdown_html(token, color, top_tokens, top_probs, perplexity=0): - html = f'
{token}
' + html += '' return html # About 750 characters per token... @@ -263,13 +281,6 @@ def custom_css(): pre { white-space: pre-wrap; } - - # TODO: This makes the hover menus extend outside the bounds of the chat area, which is good. - # However, it also makes the scrollbar disappear, which is bad. - # The scroll bar needs to still be present. So for now, we can't see dropdowns that extend past the edge of the chat area. - #.chat { - # overflow-y: auto; - #} """ @@ -282,7 +293,7 @@ def convert_to_markdown(string): return '
' + string + '
' -html_generator.convert_to_markdown = convert_to_markdown +# html_generator.convert_to_markdown = convert_to_markdown def ui():