Fix treescope html saving

inseq-team · Aug 15, 2024 · a33fae2 · a33fae2
1 parent 9de007e
commit a33fae2
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 7 deletions.
diff --git a/inseq/commands/attribute_context/attribute_context_viz_helpers.py b/inseq/commands/attribute_context/attribute_context_viz_helpers.py
@@ -326,6 +326,7 @@ def visualize_attribute_context_treescope(
                         rounding=10,
                         colormap=cmap_cci,
                         strip_chars=replace_chars,
+                        show_empty_tokens=False,
                     )
                 )
                 cci_parts.append(rp.text("\n\n"))
@@ -341,6 +342,7 @@ def visualize_attribute_context_treescope(
                         rounding=10,
                         colormap=cmap_cci,
                         strip_chars=replace_chars,
+                        show_empty_tokens=False,
                     )
                 )
                 cci_parts.append(rp.text("\n\n"))

diff --git a/inseq/data/viz.py b/inseq/data/viz.py
@@ -234,7 +234,7 @@ def show_granular_attributions(
     if display:
         ts.show(fig)
     if return_html:
-        return ts.render_to_html(fig)
+        return ts.lowering.render_to_html_as_root(fg.treescope_part_from_display_object(fig))
 
 
 def show_token_attributions(
@@ -372,7 +372,7 @@ def show_token_attributions(
     if display:
         ts.show(fig)
     if return_html:
-        return ts.render_to_html(fig)
+        return ts.lowering.render_to_html_as_root(fg.treescope_part_from_display_object(fig))
 
 
 def get_attribution_colors(
@@ -621,12 +621,12 @@ def get_single_token_heatmap_treescope(
             parts.append(rp.text(repl))
     if (show_empty_tokens and token != "") or curr_tok != "":
         show_token = token if show_empty_tokens and curr_tok == "" else curr_tok
-        highlighted_text = fg.treescope_part_from_display_object(
-            fg.text_on_color(show_token, value=round(score, rounding), vmin=min_val, vmax=max_val, colormap=colormap)
-        )
-        parts[idx_highlight] = highlighted_text
     else:
-        parts.pop(idx_highlight)
+        show_token = " "
+    highlighted_text = fg.treescope_part_from_display_object(
+        fg.text_on_color(show_token, value=round(score, rounding), vmin=min_val, vmax=max_val, colormap=colormap)
+    )
+    parts[idx_highlight] = highlighted_text
     if return_highlighted_idx:
         return parts, idx_highlight, show_token
     return parts