Skip to content

tinker_cookbook.utils.format_colorized

tinker_cookbook.utils.format_colorized(tokens, weights, tokenizer, draw_newline_arrow)

Colour-code text according to per-token weights.

Tokens are coloured green (weight > 0), yellow (weight == 0), or red (weight < 0). The function minimises ANSI escape sequences by wrapping runs of like-coloured tokens, and decodes each run in a single call so that multi-byte characters (e.g. CJK) render correctly.

Parameters:

  • tokens (list[int]) – Token IDs to decode and display.
  • weights (list[float]) – Per-token weight values (same length as tokens).
  • tokenizer (Tokenizer) – Tokenizer used to decode token IDs to text.
  • draw_newline_arrow (bool) – If True, render newlines as a visible arrow (↵) followed by an actual newline.

Returns: str – ANSI-coloured string suitable for terminal display.

Raises:

  • ValueError: If tokens and weights have different lengths.
text = format_colorized(token_ids, loss_weights, tokenizer)
print(text)