Captum: Small extension for visualize_text function

[heytitle]

Hi,

As a user, it would be great if visualize_text can produce legend for me. Although it might be obvious that green is positive and red is for negative attributions, I still think it might be useful to show it explicitly via legend in some context or situation.

Here is my prototype for the functionality purposed.

What do you think?

[bilalsal]

Hi @heytitle, that is a great suggestion, thank you!

We agree on offering the option to show such legend.
If your prototype uses matplotlib, we would much appreciate if you create a PR, that we can review to offer this functionality in the next release.

Thanks!
Bilal

[heytitle]

Because visualize_text is based on HTML DOMs, my prototype is also based on that.

This is the code I added into visualize_text:

  if legend:
      dom.append('<div style="border-top: 1px solid; margin-top: 5px; padding-top: 5px; display: inline-block">')
      dom.append('<b>Legend: </b>')
      for value, label in zip([-1, 0, 1], ["Negative", "Neutral", "Positive Attribution"]):
        dom.append(
            (
              f'<span style="display: inline-block; width: 10px; height: 10px; border: 1px solid; background-color: {visualization._get_color(value)}"></span>'
              f' {label}'
              f'  '
            )
        )
      dom.append('</div>')

[bilalsal]

Thank you very much @heytitle! You are right, it is implemented in HTML.
We will include your contribution in the next release.

[heytitle]

hi everyone,

Thanks for the feedback! I'll create a PR soon.