penzai.toolshed.token_visualization.show_token_scores(tokens: pz.nx.NamedArrayBase | jax.Array | np.ndarray, scores: pz.nx.NamedArrayBase | jax.Array | np.ndarray, vocab: SentencePieceProcessorLike, axis_name_order: Sequence[str] | None = None, vmax: float | None = None)[source]#

Renders an array of token scores, but with the scored tokens as words.

  • tokens – An array of token IDs.

  • scores – An array of token scores of the same shape as tokens.

  • vocab – Vocabulary, usually a SentencePiece tokenizer.

  • axis_name_order – Optional ordering of the axis names. The last axis name will be laid out horizontally. Inferred if not given.

  • vmax – Value at which to truncate the colormap.


A visualization of the token array.