diff --git a/inseq/data/viz.py b/inseq/data/viz.py index 30078f0..99a13fe 100644 --- a/inseq/data/viz.py +++ b/inseq/data/viz.py @@ -298,7 +298,6 @@ def show_token_attributions( title = "Generated text:\n\n" if step_score_highlight is not None: step_scores = attr.step_scores[step_score_highlight] - scores_vmin = step_scores.min().item() scores_vmax = step_scores.max().item() title = f"Generated text with {step_score_highlight} highlights:\n\n" generated_token_parts.append(rp.custom_style(rp.text(title), css_style="font-weight: bold;")) @@ -333,7 +332,7 @@ def show_token_attributions( if step_scores is not None: gen_tok_label = fg.treescope_part_from_display_object( fg.text_on_color( - curr_gen_tok, value=round(step_scores[gen_idx].item(), 4), vmin=scores_vmin, vmax=scores_vmax + curr_gen_tok, value=round(step_scores[gen_idx].item(), 4), vmax=scores_vmax ) ) else: