visualization.py

from typing import Iterable

import numpy as np
import plotly.express as px  # type: ignore


def highlight(word: str, attn: float) -> str:
    color = "#%02X%02X%02X" % (255, int(255 * (1 - attn)), int(255 * (1 - attn)))
    return f'<span style="background-color: {color}">{word}</span>'


def html_hext(words_attn: Iterable[tuple[str, float]]) -> str:
    return " ".join(highlight(word, attn) for word, attn in words_attn)


def heatmap(data: np.ndarray):
    y_labels = [
        "嘲笑や特性を理解されない",
        "特性や能力への攻撃",
        "学校や職場で受け入れられない",
        "特性をおかしいとみなされる",
        "障害への差別や苦悩をなかったことにされる",
        "うまくコミュニケーションがとれない",
        "障害について理解されない",
        "侮蔑される，認められない",
        "周囲の理解不足",
        "障害をなかったことにされる，責められる",
    ]
    fig = px.imshow(data, labels=dict(x="判定", y="名称"), y=y_labels)
    fig.update_layout(coloraxis_colorbar=dict(title="得点"))
    return fig