From a33fae2e2aae49ddbc45d912944ca8f35b347381 Mon Sep 17 00:00:00 2001
From: Gabriele Sarti <gabriele.sarti996@gmail.com>
Date: Thu, 15 Aug 2024 10:31:27 +0200
Subject: [PATCH] Fix treescope html saving

---
 .../attribute_context_viz_helpers.py               |  2 ++
 inseq/data/viz.py                                  | 14 +++++++-------
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/inseq/commands/attribute_context/attribute_context_viz_helpers.py b/inseq/commands/attribute_context/attribute_context_viz_helpers.py
index bba42ac..babed31 100644
--- 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
index 86018cd..362ac98 100644
--- 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