[display_colored_text_demo.py] text coloring, useful when working on RAG

display_colored_text_demo.py

"""Demo for displaying colored text in Jupyter Notebook."""

import matplotlib
import matplotlib.pyplot as plt
from IPython.display import HTML, display


def importance_to_color(importance):
    cmap = plt.get_cmap("Greys")  # YlOrRd, Greys
    rgba = cmap(importance)
    return matplotlib.colors.rgb2hex(rgba)


def display_colored_text(items, sepatate_lines=False):
    if sepatate_lines:
        colored_text = [
            f'<p style="color: {importance_to_color(item["importance"])}">{item["text"]}</p>'
            for item in items
        ]
    else:
        colored_text = [
            f'<span style="color: {importance_to_color(item["importance"])}">{item["text"]}</span>'
            for item in items
        ]
    display(HTML(" ".join(colored_text)))


items = [
    {"text": "Hello, world!", "importance": 0.5},
    {"text": "How are you?", "importance": 0.7},
    {"text": "Goodbye", "importance": 0.15},
    {"text": "How it is going?", "importance": 0.99},
    {"text": "Goodbye", "importance": 0.2},
    {"text": "How are you?", "importance": 0.7},
]

#  display the colored text
display_colored_text(items)
# items