show_token_scores#
- 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.
- Parameters:
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.
- Returns:
A visualization of the token array.