rten_generate/

generator.rs

//! Tools to run the generation loop for an auto-regressive model.

use std::error::Error;
use std::fmt;
use std::ops::Range;

use rten::{Dimension, NodeId, RunOptions, Value, ValueOrView, ValueView};
use rten_tensor::prelude::*;
use rten_tensor::{NdTensor, Tensor};

#[cfg(feature = "text-decoder")]
use rten_text::{Tokenizer, TokenizerError};

use crate::filter::LogitsFilter;
use crate::logits::Logits;
use crate::metrics::Metrics;
use crate::model::Model;
use crate::sampler::{ArgMax, Sampler};

#[cfg(feature = "text-decoder")]
use crate::text_decoder::TextDecoder;

/// Integer type used to represent token IDs.
pub type TokenId = u32;

/// Errors that occur when creating or running a [`Generator`].
#[derive(Debug)]
pub enum GeneratorError {
    /// An expected model input was not found.
    InputNotFound(String),

    /// An expected model output was not found.
    OutputNotFound(String),

    /// An input or output did not have the expected shape.
    ShapeMismatch(String),

    /// An error occurred while generating the next token.
    GenerateError(Box<dyn Error>),

    /// An error occurred while decoding tokens.
    #[cfg(feature = "text-decoder")]
    DecodeError(TokenizerError),
}

impl fmt::Display for GeneratorError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match self {
            GeneratorError::InputNotFound(name) => write!(f, "model input not found: {}", name),
            GeneratorError::OutputNotFound(name) => write!(f, "model output not found: {}", name),
            GeneratorError::ShapeMismatch(err) => write!(f, "shape mismatch: {}", err),
            GeneratorError::GenerateError(err) => write!(f, "generation error: {}", err),
            #[cfg(feature = "text-decoder")]
            GeneratorError::DecodeError(err) => write!(f, "decode error: {}", err),
        }
    }
}

impl Error for GeneratorError {}

/// Wraps an error with associated context for debugging.
#[derive(Debug)]
struct ErrorContext {
    error: Box<dyn Error>,
    context: String,
}

impl Error for ErrorContext {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        Some(self.error.as_ref())
    }
}

impl std::fmt::Display for ErrorContext {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}: {}", self.context, self.error)
    }
}

enum KvCacheData {
    /// Key-value cache with shape `[batch, seq_len, channels]`.
    ///
    /// In this configuration the channels for all heads are combined into the
    /// last dimension.
    BatchSeqChans(NdTensor<f32, 3>),
    /// Key-value cache with shape `[batch, heads, seq_len, channels]`.
    BatchHeadSeqChans(NdTensor<f32, 4>),
}

impl KvCacheData {
    /// Allocate a KV cache buffer with the given batch size, number of heads
    /// and embed size.
    ///
    /// The buffer initially has capacity to be extended to a sequence length
    /// of `seq_len_capacity`.
    fn with_capacity(
        batch_size: usize,
        n_heads: Option<usize>,
        size: usize,
        seq_len_capacity: usize,
    ) -> KvCacheData {
        if let Some(n_heads) = n_heads {
            KvCacheData::BatchHeadSeqChans(NdTensor::with_capacity(
                [batch_size, n_heads, seq_len_capacity, size],
                2, /* seq dim */
            ))
        } else {
            KvCacheData::BatchSeqChans(NdTensor::with_capacity(
                [batch_size, seq_len_capacity, size],
                1, /* seq dim */
            ))
        }
    }

    /// Return the current sequence length of the cache.
    fn sequence_len(&self) -> usize {
        match self {
            KvCacheData::BatchSeqChans(data) => data.size(1),
            KvCacheData::BatchHeadSeqChans(data) => data.size(2),
        }
    }

    /// Return true if the KV cache has capacity for a given sequence length.
    fn has_capacity(&self, sequence_len: usize) -> bool {
        match self {
            KvCacheData::BatchSeqChans(data) => {
                data.has_capacity(1 /* seq dim */, sequence_len)
            }
            KvCacheData::BatchHeadSeqChans(data) => {
                data.has_capacity(2 /* seq dim */, sequence_len)
            }
        }
    }

    /// Clone this cache into a new buffer with space to store sequences of
    /// a given size.
    fn clone_with_capacity(&self, max_sequence_len: usize) -> KvCacheData {
        let max_sequence_len = max_sequence_len.max(self.sequence_len());
        match self {
            KvCacheData::BatchSeqChans(data) => {
                let [batch, _seq, chans] = data.shape();
                let mut new_data =
                    NdTensor::with_capacity([batch, max_sequence_len, chans], 1 /* seq dim */);
                new_data.append(1, data).expect("should have capacity");
                KvCacheData::BatchSeqChans(new_data)
            }
            KvCacheData::BatchHeadSeqChans(data) => {
                let [batch, n_heads, _seq, chans] = data.shape();
                let mut new_data = NdTensor::with_capacity(
                    [batch, n_heads, max_sequence_len, chans],
                    2, /* seq dim */
                );
                new_data.append(2, data).expect("should have capacity");
                KvCacheData::BatchHeadSeqChans(new_data)
            }
        }
    }
}

/// Key-value cache for a single layer of a transformer model.
struct KvCache {
    /// Input ID for this cache entry.
    input_id: NodeId,

    /// Output ID for this cache entry.
    output_id: NodeId,

    /// The cached keys and values. This is set to `None` during inference, as
    /// the model temporarily takes ownership of it.
    cache: Option<KvCacheData>,
}

impl KvCache {
    fn size(&self) -> Option<usize> {
        self.cache.as_ref().map(|c| c.sequence_len())
    }
}

/// Specifies a pattern for the name of a key-value cache input or output.
///
/// These inputs are expected to have the form `{prefix}{layer_number}{suffix}`,
/// with one input and output per layer for the key cache and the value cache.
pub struct KVCachePattern<'a> {
    pub prefix: &'a str,
    pub suffix: &'a str,
}

impl<'a> From<(&'a str, &'a str)> for KVCachePattern<'a> {
    /// Construct a [`KVCachePattern`] from a `(prefix, suffix)` tuple.
    fn from(value: (&'a str, &'a str)) -> Self {
        let (prefix, suffix) = value;
        KVCachePattern { prefix, suffix }
    }
}

/// Specifies a pair of patterns for corresponding input and output key-value
/// cache entries.
pub struct KVCachePair<'a> {
    /// The pattern for the model input name.
    pub input: KVCachePattern<'a>,

    /// The pattern for the model output name.
    pub output: KVCachePattern<'a>,

    /// Specifies whether this cache is used for a cross-attention ("encoder")
    /// KV cache.
    ///
    /// Encoder KV-cache entries are computed only on the first run of the
    /// model and reused in subsequent runs.
    pub encoder: bool,
}

/// Specifies the names of model inputs and outputs.
///
/// The [`Default`] impl for this struct returns an instance whose names
/// follow the configuration of Hugging Face's Optimum tool.
///
/// Any inputs that are not present in the model are ignored.
pub struct ModelInputsConfig<'a> {
    /// Model input that contains the token IDs of the prompt and output
    /// generated so far.
    pub input_ids: &'a str,

    /// Model output that contains logits.
    pub logits: &'a str,

    /// Model input that contains an attention mask.
    pub attention_mask: &'a str,

    /// Model input that contains KV cache positions for each position.
    ///
    /// This input does not have a batch dimension.
    pub cache_position: &'a str,

    /// Patterns for inputs and outputs used for key-value caches.
    pub kv_caches: Vec<KVCachePair<'a>>,

    /// Model input that contains position IDs for each position.
    ///
    /// This input has a batch dimension.
    pub position_ids: &'a str,

    /// Boolean input that is set to false on the first run and true on
    /// subsequent runs.
    pub use_cache_flag: &'a str,
}

/// Contains essential configuration needed for a `Generator` to execute a
/// model, such as the roles of different inputs and outputs.
#[derive(Default)]
pub struct GeneratorConfig<'a> {
    /// Specifies names and roles of model inputs and outputs.
    pub model_inputs: ModelInputsConfig<'a>,

    /// Reserve capacity in the KV cache for a given number of tokens.
    ///
    /// By default the KV cache starts off with a small initial capacity and is
    /// re-allocated if needed after a token is generated, with the capacity
    /// being doubled each time to amortize overhead.
    ///
    /// If the maximum number of tokens that will be generated is known in
    /// advance, this can be used to reduce the overhead of re-allocating the
    /// KV cache during inference.
    pub kv_cache_capacity: Option<usize>,
}

impl Default for ModelInputsConfig<'_> {
    /// Return default model input names.
    ///
    /// These are based on [Hugging Face's
    /// Optimum](https://huggingface.co/docs/optimum/en/index) model exporter.
    fn default() -> Self {
        ModelInputsConfig {
            input_ids: "input_ids",
            logits: "logits",
            attention_mask: "attention_mask",
            cache_position: "cache_position",
            position_ids: "position_ids",
            use_cache_flag: "use_cache_branch",

            // Patterns are matched in order, so patterns with longer prefixes/
            // suffixes are listed first to ensure we match them.
            kv_caches: [
                // "Merged" decoders exported by Optimum for encoder-decoder
                // models. These have KV caches for both the self-attention and
                // cross-attention modules.
                KVCachePair {
                    input: ("past_key_values.", ".decoder.key").into(),
                    output: ("present.", ".decoder.key").into(),
                    encoder: false,
                },
                KVCachePair {
                    input: ("past_key_values.", ".decoder.value").into(),
                    output: ("present.", ".decoder.value").into(),
                    encoder: false,
                },
                KVCachePair {
                    input: ("past_key_values.", ".encoder.key").into(),
                    output: ("present.", ".encoder.key").into(),
                    encoder: true,
                },
                KVCachePair {
                    input: ("past_key_values.", ".encoder.value").into(),
                    output: ("present.", ".encoder.value").into(),
                    encoder: true,
                },
                // Decoder-only models exported by Optimum.
                KVCachePair {
                    input: ("past_key_values.", ".key").into(),
                    output: ("present.", ".key").into(),
                    encoder: false,
                },
                KVCachePair {
                    input: ("past_key_values.", ".value").into(),
                    output: ("present.", ".value").into(),
                    encoder: false,
                },
            ]
            .into(),
        }
    }
}

/// Generates a token ID sequence using a transformer decoder model.
///
/// This is an iterator that runs the model on each call to [`Iterator::next`]
/// and yields a result containing the next token ID or an error.
///
/// The token ID sequence can be converted to text using the
/// [`decode`](GeneratorUtils::decode) method of the [`GeneratorUtils`] trait.
///
/// The `GeneratorUtils` trait also provides useful wrappers for the output,
/// such as stopping generation when an end-of-text token is reached. You can
/// also use all of the standard iterator adapters. For example
/// `generator.take(30)` will return an iterator that stops generation after 30
/// tokens have been produced.
///
/// ## Processing pipeline
///
/// Each call to [`next`](Iterator::next) performs the following steps:
///
/// 1. Compute model inputs. This includes the token IDs (prompt on first run,
///    most recently sampled token after that), constant inputs and
///    position-varying inputs (eg. attention mask).
/// 2. Run the model
/// 3. Apply filters to model outputs ("logits")
/// 4. Sample a token from the logits
/// 5. Save the sampled token and KV-caches from the model for the next
///    generation step.
/// 6. Return the sampled token as the iterator output
///
/// ## Logit filters
///
/// The raw model outputs can be modified before sampling by configuring a
/// [`LogitsFilter`] using [`with_logits_filter`](Generator::with_logits_filter).
///
/// A chain of filters can be created using [`Chain`](crate::filter::Chain).
/// When setting up the chain, a best practice is to put the most selective
/// filters first (ie. those which eliminate the most tokens).
///
///
/// ```no_run
/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
/// use rten::Model;
/// use rten_generate::Generator;
/// use rten_generate::filter::Chain;
/// use rten_generate::sampler::Multinomial;
///
/// let model = Model::load_file("model.onnx")?;
/// let mut generator = Generator::from_model(&model)?
///   .with_logits_filter(
///     Chain::new()
///       .top_k(30) // Remove all tokens except the 30 with the highest score
///       .temperature(0.7) // Scale scores using `score / 0.7`.
///       .top_p(0.9) // Take the top N tokens whose cumulative probability exceeds 0.9.
///   )
///   // Select a random token from the filtered set according to the probability
///   // of each.
///   .with_sampler(Multinomial::new());
///
/// # Ok(()) }
/// ```
///
/// ## Sampling
///
/// The token ID is sampled from the outputs of the model (the "logits") using
/// a [`Sampler`]. By default this is [`ArgMax`] which simply chooses
/// the token with the highest probability. The sampler can be configured using
/// [`with_sampler`](Self::with_sampler).
///
/// ## Key-value caches and generation performance
///
/// To enable efficient decoding, the model should have inputs and outputs for
/// the [key-value
/// cache](https://peterchng.com/blog/2024/06/11/what-is-the-transformer-kv-cache/).
/// The generator will work with models that do not have cache inputs, but
/// decoding of long output sequences will be much slower.
pub struct Generator<'a> {
    model: &'a dyn Model,

    run_options: Option<RunOptions>,

    /// Additional constant model inputs (eg. encoder outputs) passed to the
    /// model at each step.
    constant_inputs: Vec<(NodeId, ValueOrView<'a>)>,

    /// Additional model inputs computed using constant propagation. This
    /// effectively caches parts of the graph that don't change in each
    /// generation step. This is `None` if the cache is out of date.
    constant_prop_inputs: Option<Vec<(NodeId, Value)>>,

    /// Additional varying model inputs computed and passed to the model at
    /// each step. The functions receive `(batch_size, sequence_positions)` as inputs.
    #[allow(clippy::type_complexity)]
    varying_inputs: Vec<(NodeId, &'a dyn Fn(usize, Range<usize>) -> ValueOrView<'a>)>,

    /// Input token IDs for the next run of the model.
    input_ids: Vec<TokenId>,

    /// Position ID associated with the first token in `input_ids`.
    input_offset: usize,

    /// Input node IDs
    input_ids_input: NodeId,

    /// Output node IDs
    logits_output: NodeId,

    /// Filter used to modify logits before sampling.
    logits_filter: Option<Box<dyn LogitsFilter + 'a>>,

    /// Sampler used to get the next token ID from the output logits.
    sampler: Box<dyn Sampler + 'a>,

    /// Previously sampled tokens. These are retained for conditional filtering
    /// and sampling.
    prev_tokens: Vec<u32>,

    /// Self-attention key-value cache. This is extended on each iteration.
    kv_cache: Vec<KvCache>,

    /// Cross-attention key-value cache.
    ///
    /// This is used by encoder-decoder models. The cross-attention values
    /// are computed on the first run and reused in subsequent runs.
    encoder_kv_cache: Vec<KvCache>,
}

impl<'a> Generator<'a> {
    /// Create a generator that iteratively produces tokens using a model.
    ///
    /// This function assumes default names for model inputs and outputs
    /// based on the conventions of Hugging Face's Optimum exporter. These
    /// can be customized using [`from_model_config`](Self::from_model_config).
    ///
    /// The model must have the required inputs:
    ///
    ///  - `input_ids` - (batch, sequence) tensor of token IDs
    ///
    /// The model may have the optional inputs:
    ///
    ///  - `attention_mask` - (batch, sequence) tensor of booleans
    ///  - `cache_position` - (sequence) tensor of KV-cache positions. Usually the same as `position_ids`
    ///  - `position_ids` - (batch, sequence) tensor of position indices
    ///  - `past_key_values.N.key` - (batch, head, past_seq_len, size) key vector cache
    ///    where `N` is the layer index
    ///  - `past_key_values.N.value` - (batch, head, past_key_values, size) value vector cache,
    ///    where `N` is the layer index
    ///
    /// **Warning:** Generation of long sequences will be much slower in models without
    /// key-value caches.
    ///
    /// The model must have the outputs:
    ///
    ///  - `logits` - output (batch, sequence, vocab) tensor of next token probabilities
    ///
    /// The model may have the optional outputs:
    ///
    ///  - `present.N.key` - (batch, head, past_seq_len + 1, size) updated key vector cache
    ///  - `present.N.value` - (batch, head, past_seq_len + 1, size) updated value vector cache
    pub fn from_model(model: &'a dyn Model) -> Result<Generator<'a>, GeneratorError> {
        let config = GeneratorConfig {
            model_inputs: ModelInputsConfig::default(),
            kv_cache_capacity: None,
        };
        Self::from_model_config(model, config)
    }

    /// Create a generator that iteratively produces tokens using a model.
    ///
    /// This is a variant of [`from_model`](Self::from_model) that allows
    /// specifying custom names for model inputs.
    pub fn from_model_config(
        model: &'a dyn Model,
        config: GeneratorConfig,
    ) -> Result<Generator<'a>, GeneratorError> {
        let model_inputs = &config.model_inputs;

        let input_ids_input =
            model
                .find_node(model_inputs.input_ids)
                .ok_or(GeneratorError::InputNotFound(
                    model_inputs.input_ids.to_string(),
                ))?;

        let logits_output =
            model
                .find_node(model_inputs.logits)
                .ok_or(GeneratorError::OutputNotFound(
                    model_inputs.logits.to_string(),
                ))?;

        // Find inputs and corresponding outputs for key-value cache.
        let batch_size = 1;
        let mut kv_cache = Vec::new();
        let mut encoder_kv_cache = Vec::new();
        for &input_id in model.input_ids() {
            let input_info = model
                .node_info(input_id)
                .ok_or(GeneratorError::InputNotFound(format!(
                    "input ID {}",
                    input_id
                )))?;

            let name = input_info.name();

            let Some(kv_pattern) = model_inputs
                .kv_caches
                .iter()
                .find(|pat| name.starts_with(pat.input.prefix) && name.ends_with(pat.input.suffix))
            else {
                continue;
            };

            let (n_heads, size) = match *input_info.shape() {
                [_, Dimension::Fixed(n_heads), _, Dimension::Fixed(size)] => (Some(n_heads), size),
                [_, _, Dimension::Fixed(size)] => (None, size),
                _ => {
                    return Err(GeneratorError::ShapeMismatch(format!(
                        "input \"{}\" has unexpected shape. expected (batch, past_seq_len, chans) or (batch, heads, past_seq_len, chans) where `heads` and `size` are fixed",
                        name
                    )));
                }
            };

            let prefix = kv_pattern.input.prefix;

            let layer_index_start = prefix.len();
            let layer_index_end = name.len() - kv_pattern.input.suffix.len();
            let layer_index_str = &name[layer_index_start..layer_index_end];
            let Ok(layer_index) = layer_index_str.parse::<u32>() else {
                continue;
            };

            let output_prefix = kv_pattern.output.prefix;
            let output_suffix = kv_pattern.output.suffix;

            let output_name = format!("{}{}{}", output_prefix, layer_index, output_suffix);
            let output_id = model
                .find_node(&output_name)
                .ok_or(GeneratorError::OutputNotFound(output_name))?;

            // Initial sequence length capacity for KV cache buffer.
            //
            // For models that execute different operations on the first vs
            // subsequent iterations (eg. Hugging Face "merged" models with
            // past and no-past branches) the input buffer may not be used in
            // the first iteration. Instead we need to reserve capacity once
            // the model returns the initial KV cache.
            //
            // For other simpler models the input KV cache buffer is used for
            // all iterations, in which case we would ideally reserve capacity
            // up-front based on the max expected sequence length.
            let max_seq_len = config.kv_cache_capacity.unwrap_or(1);

            let kv_cache_entry = KvCache {
                input_id,
                output_id,
                cache: Some(KvCacheData::with_capacity(
                    batch_size,
                    n_heads,
                    size,
                    max_seq_len,
                )),
            };

            if kv_pattern.encoder {
                encoder_kv_cache.push(kv_cache_entry);
            } else {
                kv_cache.push(kv_cache_entry);
            }
        }

        let mut generator = Generator {
            model,
            run_options: None,

            constant_inputs: Vec::new(),
            varying_inputs: Vec::new(),

            // Constant propagation is performed as a graph optimization when
            // the model is loaded, so we only need to re-do it if additional
            // constant inputs are added.
            constant_prop_inputs: Some(Vec::new()),

            logits_filter: None,
            input_ids: vec![],
            input_ids_input,
            input_offset: 0,
            logits_output,
            kv_cache,
            encoder_kv_cache,
            prev_tokens: Vec::new(),
            sampler: Box::new(ArgMax::new()),
        };

        let attention_mask_input = model.find_node(model_inputs.attention_mask);
        if let Some(attention_mask_input) = attention_mask_input {
            generator = generator
                .with_varying_input(attention_mask_input, &|batch_size, positions| {
                    NdTensor::full([batch_size, positions.end], 1i32).into()
                });
        }

        let position_ids_input = model.find_node(model_inputs.position_ids);
        if let Some(position_ids_input) = position_ids_input {
            generator =
                generator.with_varying_input(position_ids_input, &|batch_size, positions| {
                    NdTensor::from_fn([batch_size, positions.len()], |[_batch, pos]| {
                        (positions.start + pos) as i32
                    })
                    .into()
                });
        }

        let cache_position_input = model.find_node(model_inputs.cache_position);
        if let Some(cache_position_input) = cache_position_input {
            generator =
                generator.with_varying_input(cache_position_input, &|_batch_size, positions| {
                    NdTensor::from_fn([positions.len()], |[pos]| (positions.start + pos) as i32)
                        .into()
                });
        }

        let use_cache_input = model.find_node(model_inputs.use_cache_flag);
        if let Some(use_cache_input) = use_cache_input {
            generator = generator.with_varying_input(use_cache_input, &|_batch_size, positions| {
                Tensor::from(if positions.start == 0 { 0i32 } else { 1 }).into()
            });
        }

        Ok(generator)
    }

    /// Set the initial sequence of tokens (aka. the prompt) passed to the model
    /// when it is first run.
    ///
    /// To add new inputs after the initial generation, use
    /// [`append_prompt`](Self::append_prompt) instead.
    pub fn with_prompt(mut self, prompt: &[TokenId]) -> Self {
        self.input_ids = prompt.to_vec();
        self
    }

    /// Add input tokens to be included in the next iteration of the model.
    ///
    /// This is useful in applications such as chat where the model's input
    /// alternates between encoded user input and model-generated output.
    pub fn append_prompt(&mut self, prompt: &[TokenId]) {
        self.input_ids.extend(prompt);
    }

    /// Clear the pending prompt for the next generation.
    ///
    /// This includes tokens added using [`with_prompt`](Self::with_prompt) and
    /// [`append_prompt`](Self::append_prompt) as well as the single token
    /// sampled by the most recent call to [`Iterator::next`](Self::next).
    ///
    /// This does not affect the state resulting from tokens that have already
    /// been generated. In other words, it does not "rewind" the conversation.
    pub fn clear_prompt(&mut self) {
        self.input_ids.clear();
    }

    /// Return the prompt that will be used for the next generation.
    pub fn prompt(&self) -> &[TokenId] {
        &self.input_ids
    }

    /// Return the tokens that have been generated so far, including the prompt.
    pub fn prev_tokens(&self) -> &[TokenId] {
        &self.prev_tokens
    }

    /// Return the current decoder KV-cache length.
    ///
    /// Returns `None` if the model does not use a KV cache.
    pub fn kv_cache_len(&self) -> Option<usize> {
        self.kv_cache.first()?.size()
    }

    /// Add a constant input which is provided to the model at each iteration.
    ///
    /// A common use case is to pass the outputs of an encoder model to
    /// an auto-regressive decoder.
    pub fn with_constant_input(mut self, input_id: NodeId, value: ValueView<'a>) -> Self {
        self.constant_prop_inputs = None;
        self.constant_inputs.push((input_id, value.into()));
        self
    }

    /// Add an input which varies with the sequence position.
    ///
    /// `value_fn` receives `(batch_size, sequence_positions)` as input and
    /// computes the value for the input at the given positions.
    ///
    /// A common use case is to pass position embeddings, if they are not
    /// computed internally by the model.
    pub fn with_varying_input<F: Fn(usize, Range<usize>) -> ValueOrView<'a>>(
        mut self,
        input_id: NodeId,
        value_fn: &'a F,
    ) -> Self {
        self.varying_inputs.push((input_id, value_fn));
        self
    }

    /// Set the filter used to process model output logits before passing them
    /// to the sampler to select a token ID.
    ///
    /// To combine multiple filters, use a [`Chain`](crate::filter::Chain).
    pub fn with_logits_filter<F: LogitsFilter + 'a>(mut self, filter: F) -> Self {
        self.logits_filter = Some(Box::new(filter));
        self
    }

    /// Set the sampler used to sample the next token ID from the output logits.
    ///
    /// The default sampler picks the token with the highest probability
    /// (aka. greedy sampling). The most common alternative is
    /// [`Multinomial`](crate::sampler::Multinomial) which samples tokens
    /// according to their respective probabilities.
    pub fn with_sampler<S: Sampler + 'a>(mut self, sampler: S) -> Self {
        self.sampler = Box::new(sampler);
        self
    }

    /// Set execution options used when running model inference.
    pub fn with_run_options(mut self, opts: Option<RunOptions>) -> Self {
        self.run_options = opts;
        self
    }

    /// Feed the current prompt into the model and update the KV cache.
    ///
    /// If `generate_logits` is true, the model's logits output is computed and
    /// returned as a `(batch, sequence, vocab)` tensor.
    fn generate_impl(
        &mut self,
        generate_logits: bool,
    ) -> Result<Option<NdTensor<f32, 3>>, GeneratorError> {
        let batch_size = 1;
        let input_ids: NdTensor<i32, 2> = self
            .input_ids
            .iter()
            .map(|id| *id as i32)
            .collect::<Tensor<_>>()
            .into_shape([batch_size, self.input_ids.len()]);

        let input_positions = self.input_offset..self.input_offset + self.input_ids.len();

        let mut model_inputs: Vec<(NodeId, ValueOrView)> =
            vec![(self.input_ids_input, input_ids.view().into())];

        // Propagate constants on the first run.
        if self.constant_prop_inputs.is_none() {
            let inputs = match self.model.partial_run(
                self.constant_inputs.clone(),
                &[self.logits_output],
                self.run_options.clone(),
            ) {
                Ok(inputs) => inputs,
                Err(err) => {
                    return Err(wrap_error(
                        err,
                        "failed to partially evaluate model with constant inputs",
                    ));
                }
            };
            self.constant_prop_inputs = Some(inputs);
        }

        if let Some(constants) = self.constant_prop_inputs.as_ref() {
            model_inputs.extend(
                constants
                    .iter()
                    .map(|(node_id, output)| (*node_id, output.as_view().into())),
            );
        }

        if !self.varying_inputs.is_empty() {
            model_inputs.extend(self.varying_inputs.iter().map(|(node_id, value_fn)| {
                (*node_id, value_fn(batch_size, input_positions.clone()))
            }));
        }

        // Add key-value cache from previous run. The model takes ownership
        // of the KV-cache tensor during the run so it can efficiently append
        // the entry for the current step, without copying the existing buffer.
        for entry in self.kv_cache.iter_mut() {
            let cache = entry.cache.take();
            match cache {
                Some(KvCacheData::BatchSeqChans(cache)) => {
                    model_inputs.push((entry.input_id, cache.into()));
                }
                Some(KvCacheData::BatchHeadSeqChans(cache)) => {
                    model_inputs.push((entry.input_id, cache.into()));
                }
                None => {}
            }
        }

        // Add cross-attention key-value cache.
        for entry in self.encoder_kv_cache.iter() {
            match &entry.cache {
                Some(KvCacheData::BatchSeqChans(cache)) => {
                    model_inputs.push((entry.input_id, cache.into()));
                }
                Some(KvCacheData::BatchHeadSeqChans(cache)) => {
                    model_inputs.push((entry.input_id, cache.into()));
                }
                None => {}
            }
        }

        // Run the model and collect updated KV cache and logits.
        let mut model_outputs: Vec<NodeId> = self
            .kv_cache
            .iter()
            .map(|entry| entry.output_id)
            .chain(self.encoder_kv_cache.iter().map(|entry| entry.output_id))
            .collect();

        if generate_logits {
            model_outputs.push(self.logits_output);
        }

        let mut outputs = self
            .model
            .run(model_inputs, &model_outputs, self.run_options.clone())
            .map_err(|e| wrap_error(e, "failed to run model"))?;

        // Update the self-attention key-value cache.
        //
        // The KV cache tensors returned from the model should be the same as
        // the passed in tensors, but extended by one element along the sequence
        // axis.
        for cache_entry in self.kv_cache.iter_mut() {
            let output = outputs.remove(0);

            let err_context = "failed to save self-attention KV-cache";
            let mut kv_cache = match output.ndim() {
                3 => KvCacheData::BatchSeqChans(
                    output.try_into().map_err(|e| wrap_error(e, err_context))?,
                ),
                4 => KvCacheData::BatchHeadSeqChans(
                    output.try_into().map_err(|e| wrap_error(e, err_context))?,
                ),
                ndim => {
                    return Err(wrap_error(
                        format!("KV cache has {} dims, expected 3 or 4", ndim),
                        err_context,
                    ));
                }
            };

            // Grow the KV cache buffer if it has reached the limit of its
            // pre-allocated sequence length.
            //
            // Double the capacity each time to amortize the costs of copying
            // the previous buffer.
            if !kv_cache.has_capacity(kv_cache.sequence_len() + 1) {
                kv_cache = kv_cache.clone_with_capacity(kv_cache.sequence_len() * 2);
            }

            cache_entry.cache = Some(kv_cache);
        }

        // Update the cross-attention key-value cache.
        for cache_entry in self.encoder_kv_cache.iter_mut() {
            let output = outputs.remove(0);
            if output.is_empty() {
                // Optimum-exported models only return encoder KV-cache tensors
                // on the first run and dummy empty tensors on subsequent runs.
                // Ignore these and continue to use the value from the first run.
                continue;
            }

            let err_context = "failed to save cross-attention KV-cache";
            let kv_cache = match output.ndim() {
                3 => KvCacheData::BatchSeqChans(
                    output.try_into().map_err(|e| wrap_error(e, err_context))?,
                ),
                4 => KvCacheData::BatchHeadSeqChans(
                    output.try_into().map_err(|e| wrap_error(e, err_context))?,
                ),
                ndim => {
                    return Err(wrap_error(
                        format!("KV cache has {} dims, expected 3 or 4", ndim),
                        err_context,
                    ));
                }
            };
            cache_entry.cache = Some(kv_cache);
        }

        // Save prompt for use in logit filters.
        if self.prev_tokens.is_empty() {
            self.prev_tokens.extend(self.input_ids.iter());
        }

        // Clear the prompt for the next generation.
        if !self.kv_cache.is_empty() {
            self.input_offset += self.input_ids.len();
            self.input_ids.clear();
        }

        if generate_logits {
            // Apply filtering to model outputs.
            let logits: NdTensor<f32, 3> = outputs
                .remove(0)
                .try_into()
                .map_err(|e| wrap_error(e, "failed to extract logits from model outputs"))?;
            Ok(Some(logits))
        } else {
            Ok(None)
        }
    }

    /// Run the model and update the KV cache.
    ///
    /// Unlike calling [`next`](Self::next) this does not generate the logits
    /// output and sample a token.
    pub fn process_prompt(&mut self) -> Result<(), GeneratorError> {
        self.generate_impl(false).map(|_| ())
    }

    /// Run the model and generate the next token.
    ///
    /// The generated token is automatically added to the prompt for the next
    /// generation.
    fn generate_next_token(&mut self) -> Result<TokenId, GeneratorError> {
        let logits = self.generate_impl(true)?.expect("should have logits");
        let last_logits = Logits::dense(logits.slice((0, -1)).to_contiguous().to_vec());
        let filtered_logits = if let Some(filter) = self.logits_filter.as_ref() {
            filter.filter(last_logits, &self.prev_tokens)
        } else {
            last_logits
        };

        // If filtering removed all the tokens, we have nothing to sample from.
        if filtered_logits.is_empty() {
            return Err(GeneratorError::GenerateError(
                "filtered logits are empty".into(),
            ));
        }

        // Sample output token.
        let next_id = self.sampler.sample(&filtered_logits);

        // Append token to prompt for next generation.
        self.prev_tokens.push(next_id);
        self.input_ids.push(next_id);

        Ok(next_id)
    }
}

fn wrap_error<E>(error: E, context: &str) -> GeneratorError
where
    E: Into<Box<dyn Error>>,
{
    let error_ctx = ErrorContext {
        error: error.into(),
        context: context.to_string(),
    };
    GeneratorError::GenerateError(error_ctx.into())
}

/// Output items from a [`Generator`].
pub type GeneratorItem = Result<TokenId, GeneratorError>;

impl Iterator for Generator<'_> {
    type Item = Result<TokenId, GeneratorError>;

    /// Run the model and generate the next output token.
    ///
    /// The generated token is added to the prompt. This enables calling `next`
    /// repeatedly to generate a sequence of tokens. The prompt can be extended
    /// or cleared using [`append_prompt`](Self::append_prompt) or
    /// [`clear_prompt`](Self::clear_prompt) respectively.
    fn next(&mut self) -> Option<Self::Item> {
        Some(self.generate_next_token())
    }
}

/// Iterator utilities that wrap a [`Generator`] to perform common tasks such
/// as stopping generation when an end-of-text token is encountered.
pub trait GeneratorUtils: Iterator<Item = GeneratorItem> + Sized {
    /// Stop the generator when any token in `eos_tokens` is encountered.
    fn stop_on_tokens<A: AsRef<[u32]>>(self, eos_tokens: A) -> impl Iterator<Item = GeneratorItem> {
        self.take_while(move |tok| match tok {
            Ok(tok_id) => !eos_tokens.as_ref().contains(tok_id),
            _ => true,
        })
    }

    /// Decode the tokens to text using a tokenizer.
    ///
    /// To get both the decoded text and token IDs, call
    /// [`with_ids`](TextDecoder::with_ids) on the result.
    #[cfg(feature = "text-decoder")]
    fn decode(self, tokenizer: &Tokenizer) -> TextDecoder<'_, Self> {
        TextDecoder::wrap(self, tokenizer)
    }

    /// Record timing metrics.
    ///
    /// Metrics such as the number of tokens generated per second will be
    /// available from `metrics` after generation has finished.
    fn profile(self, metrics: &mut Metrics) -> impl Iterator<Item = Self::Item> {
        Profiler::wrap(self, metrics)
    }
}

impl<I: Iterator<Item = GeneratorItem>> GeneratorUtils for I {}

/// Wraps a [`Generator`] to record timing metrics into a [`Metrics`] struct.
struct Profiler<'a, G: Iterator> {
    generator: G,
    metrics: &'a mut Metrics,
}

impl<'a, G: Iterator> Profiler<'a, G> {
    fn wrap(generator: G, metrics: &'a mut Metrics) -> Profiler<'a, G> {
        Profiler { generator, metrics }
    }
}

impl<G: Iterator> Iterator for Profiler<'_, G> {
    type Item = G::Item;

    fn next(&mut self) -> Option<Self::Item> {
        let start = std::time::Instant::now();
        let item = self.generator.next()?;
        self.metrics.add_step_duration(start.elapsed());
        Some(item)
    }
}

#[cfg(test)]
mod tests {
    use std::cell::{Cell, RefCell};
    use std::collections::HashMap;
    use std::error::Error;
    use std::rc::Rc;

    use rten::{Dimension, NodeId, RunOptions, Value, ValueOrView};
    use rten_tensor::NdTensor;
    use rten_tensor::prelude::*;

    use super::{Generator, GeneratorUtils, Logits};
    use crate::filter::LogitsFilter;
    use crate::metrics::Metrics;
    use crate::model::{Model, NodeInfo};

    struct FakeModel {
        nodes: Vec<NodeInfo>,
        input_ids: Vec<NodeId>,
        output_ids: Vec<NodeId>,

        // Next inference step
        step: Cell<usize>,

        // Inference outputs for each step
        outputs: Vec<HashMap<NodeId, Value>>,

        // Inference inputs for each step
        inputs: RefCell<Vec<HashMap<NodeId, Value>>>,

        // Run options for most recent inference
        run_opts: Cell<Option<RunOptions>>,
    }

    impl FakeModel {
        /// Return a model with a given set of inputs and outputs.
        fn with_inputs_and_outputs(inputs: &[NodeInfo], outputs: &[NodeInfo]) -> FakeModel {
            let node_infos = [inputs, outputs].concat();
            let input_ids = (0..inputs.len())
                .map(|id| NodeId::from_u32(id as u32))
                .collect();
            let output_ids = (inputs.len()..(inputs.len() + outputs.len()))
                .map(|id| NodeId::from_u32(id as u32))
                .collect();

            FakeModel {
                input_ids,
                output_ids,
                nodes: node_infos,
                step: Cell::new(0),
                inputs: RefCell::new(vec![]),
                outputs: vec![],
                run_opts: Cell::new(None),
            }
        }

        /// Add inference outputs for one run of the model.
        fn add_outputs(&mut self, outputs: HashMap<NodeId, Value>) {
            self.outputs.push(outputs)
        }

        /// Get an input for the `step`th run of the model.
        fn get_inputs(&self, step: usize, node_id: NodeId) -> Option<Value> {
            self.inputs
                .borrow()
                .get(step)
                .map(|step_inputs| step_inputs.get(&node_id))
                .flatten()
                .cloned()
        }
    }

    impl Model for FakeModel {
        fn find_node(&self, name: &str) -> Option<NodeId> {
            self.nodes
                .iter()
                .position(|info| info.name() == name)
                .map(|pos| NodeId::from_u32(pos as u32))
        }

        fn node_info(&self, id: NodeId) -> Option<NodeInfo> {
            self.nodes.get(id.as_usize()).cloned()
        }

        fn input_ids(&self) -> &[NodeId] {
            &self.input_ids
        }

        fn run(
            &self,
            inputs: Vec<(NodeId, ValueOrView)>,
            outputs: &[NodeId],
            opts: Option<RunOptions>,
        ) -> Result<Vec<Value>, Box<dyn Error>> {
            if let Some((input_id, _)) = inputs.iter().find(|(id, _)| !self.input_ids.contains(id))
            {
                return Err(format!("invalid input ID {}", input_id).into());
            }
            for &expected_input in self.input_ids.iter() {
                if !inputs.iter().any(|&(id, _)| id == expected_input) {
                    return Err(format!("missing input ID {}", expected_input).into());
                }
            }

            if let Some(output_id) = outputs.iter().find(|id| !self.output_ids.contains(id)) {
                return Err(format!("invalid output ID {}", output_id).into());
            }

            self.inputs.borrow_mut().push(
                inputs
                    .into_iter()
                    .map(|(id, input_or_output)| (id, input_or_output.to_owned()))
                    .collect(),
            );

            let result = outputs
                .iter()
                .map(|id| {
                    let step_outputs = self
                        .outputs
                        .get(self.step.get())
                        .expect("outputs not specified for step");

                    step_outputs
                        .get(id)
                        .cloned()
                        .expect("invalid output node ID")
                })
                .collect();

            self.step.set(self.step.get() + 1);
            self.run_opts.set(opts);

            Ok(result)
        }

        fn partial_run(
            &self,
            _inputs: Vec<(NodeId, ValueOrView)>,
            _outputs: &[NodeId],
            _opts: Option<RunOptions>,
        ) -> Result<Vec<(NodeId, Value)>, Box<dyn Error>> {
            Ok(Vec::new())
        }
    }

    /// Generate `[batch, sequence, n_vocab]` tensor for `logits` output.
    fn generate_logits(n_vocab: usize, token_ids: &[u32]) -> NdTensor<f32, 3> {
        let mut logits = NdTensor::zeros([1, token_ids.len(), n_vocab]);
        for (idx, id) in token_ids.iter().copied().enumerate() {
            logits[[0, idx, id as usize]] = 1.0;
        }
        logits
    }

    #[derive(Copy, Clone, PartialEq)]
    struct TransformerParams {
        /// Number of layers. This determines the number of KV-cache inputs
        /// and outputs.
        n_layers: usize,
        n_heads: usize,
        n_embed: usize,

        /// Vocabulary size. This is the size of the last dimension of the
        /// logits output.
        n_vocab: usize,
    }

    impl Default for TransformerParams {
        fn default() -> Self {
            Self {
                n_layers: 5,
                n_heads: 3,
                n_vocab: 5,
                n_embed: 8,
            }
        }
    }

    #[derive(Copy, Clone, PartialEq)]
    enum KvCacheType {
        /// Add KV-cache inputs and outputs for self-attention.
        Decoder,
        /// Add KV-cache inputs and outputs for self-attention and cross-
        /// attention.
        EncoderDecoder,
    }

    /// Create a fake transformer model using the default names for inputs and
    /// outputs.
    fn fake_transformer_model(
        params: TransformerParams,
        kv_cache: Option<KvCacheType>,
        prompt_len: usize,
        output_token_ids: &[u32],
    ) -> FakeModel {
        let TransformerParams {
            n_layers,
            n_heads,
            n_vocab,
            n_embed,
        } = params;

        // Add inputs and outputs using the standard names.
        let mut inputs = vec![
            NodeInfo::from_name_shape("input_ids", &[]),
            NodeInfo::from_name_shape("cache_position", &[]),
            NodeInfo::from_name_shape("position_ids", &[]),
            NodeInfo::from_name_shape("attention_mask", &[]),
        ];
        let mut outputs = vec![NodeInfo::from_name_shape("logits", &[])];

        // Add KV-cache inputs and outputs.
        let mut kv_cache_output_names = Vec::new();
        if let Some(kv_cache_type) = kv_cache {
            let dims = [
                Dimension::Symbolic("batch".to_string()),
                Dimension::Fixed(n_heads as usize),
                Dimension::Symbolic("seq".to_string()),
                Dimension::Fixed(n_embed),
            ];
            let make_name_info = |name: &str| NodeInfo::from_name_shape(name, &dims);

            for layer in 0..n_layers {
                let past_names: Vec<String>;
                let present_names: Vec<String>;

                match kv_cache_type {
                    KvCacheType::Decoder => {
                        past_names = [
                            format!("past_key_values.{}.key", layer),
                            format!("past_key_values.{}.value", layer),
                        ]
                        .into();
                        present_names = [
                            format!("present.{}.key", layer),
                            format!("present.{}.value", layer),
                        ]
                        .into();
                    }
                    KvCacheType::EncoderDecoder => {
                        past_names = [
                            format!("past_key_values.{}.decoder.key", layer),
                            format!("past_key_values.{}.decoder.value", layer),
                            format!("past_key_values.{}.encoder.key", layer),
                            format!("past_key_values.{}.encoder.value", layer),
                        ]
                        .into();

                        present_names = [
                            format!("present.{}.decoder.key", layer),
                            format!("present.{}.decoder.value", layer),
                            format!("present.{}.encoder.key", layer),
                            format!("present.{}.encoder.value", layer),
                        ]
                        .into();
                    }
                }

                inputs.extend(past_names.iter().map(|name| make_name_info(&name)));
                outputs.extend(present_names.iter().map(|name| make_name_info(&name)));
                kv_cache_output_names.extend(present_names);
            }

            if kv_cache_type == KvCacheType::EncoderDecoder {
                inputs.push(NodeInfo::from_name_shape("use_cache_branch", &[]));
            }
        }

        let mut model = FakeModel::with_inputs_and_outputs(&inputs, &outputs);
        let logits_id = model.find_node("logits").unwrap();

        for (step, output_token_id) in output_token_ids.iter().copied().enumerate() {
            assert!(
                output_token_id < n_vocab as u32,
                "token ID is invalid for vocab size"
            );

            let logits = if kv_cache.is_some() {
                generate_logits(n_vocab, &[output_token_id])
            } else {
                generate_logits(n_vocab, &output_token_ids[..=step])
            };

            let mut outputs = HashMap::new();
            outputs.insert(logits_id, Value::FloatTensor(logits.into()));

            // Add KV cache outputs
            for kv_output in kv_cache_output_names.iter() {
                let kv_output_id = model.find_node(&kv_output).unwrap();
                let context_len = if step == 0 {
                    prompt_len
                } else {
                    prompt_len + step - 1
                };

                let is_encoder = model
                    .node_info(kv_output_id)
                    .as_ref()
                    .map(|ni| ni.name())
                    .unwrap_or("")
                    .contains("encoder");

                let output_n_embed = if is_encoder && step > 0 {
                    // Encoder KV cache outputs are only used on the first run.
                    // On subsequent runs return a dummy output, which should
                    // be ignored.
                    0
                } else {
                    n_embed
                };

                outputs.insert(
                    kv_output_id,
                    Value::FloatTensor(
                        NdTensor::zeros([1, n_heads, context_len, output_n_embed]).into(),
                    ),
                );
            }

            model.add_outputs(outputs);
        }

        model
    }

    fn test_generator_impl(kv_cache_type: Option<KvCacheType>) -> Result<(), Box<dyn Error>> {
        let params = TransformerParams::default();
        let expected_token_ids = [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 0, 0];
        let prompt = [1, 2, 3, 1, 2, 3];
        let model =
            fake_transformer_model(params, kv_cache_type, prompt.len(), &expected_token_ids);

        let generator = Generator::from_model(&model)?;
        let generation_len = 10;

        let output_token_ids: Vec<_> = generator
            .with_prompt(&prompt)
            .take(generation_len)
            .map(|id| id.expect("generation failed"))
            .collect();

        // Check generator outputs
        assert_eq!(output_token_ids.len(), generation_len);
        assert_eq!(output_token_ids, &expected_token_ids[..generation_len]);

        // Check model inputs
        let input_id = model.find_node("input_ids").unwrap();
        let position_ids = model.find_node("position_ids").unwrap();
        let attention_mask = model.find_node("attention_mask").unwrap();
        let cache_branch = model.find_node("use_cache_branch");
        let cache_position = model.find_node("cache_position").unwrap();

        for step in 0..generation_len {
            let step_inputs = model.get_inputs(step, input_id).unwrap();
            let step_inputs: NdTensor<i32, 2> = step_inputs.try_into().unwrap();

            let step_pos_ids = model.get_inputs(step, position_ids).unwrap();
            let step_pos_ids: NdTensor<i32, 2> = step_pos_ids.try_into().unwrap();

            let step_cache_pos = model.get_inputs(step, cache_position).unwrap();
            let step_cache_pos: NdTensor<i32, 1> = step_cache_pos.try_into().unwrap();

            let step_attn_mask = model.get_inputs(step, attention_mask).unwrap();
            let step_attn_mask: NdTensor<i32, 2> = step_attn_mask.try_into().unwrap();

            let cache_branch = cache_branch.map(|cb_id| {
                let cb = model.get_inputs(step, cb_id).unwrap();
                let cb: NdTensor<i32, 0> = cb.try_into().unwrap();
                cb
            });

            if step == 0 {
                assert_eq!(step_inputs.size(1), prompt.len());
                assert!(
                    step_inputs
                        .iter()
                        .map(|x| *x as u32)
                        .eq(prompt.iter().copied())
                );

                assert_eq!(step_attn_mask.size(1), prompt.len());
                assert!(step_attn_mask.iter().all(|x| *x == 1));

                assert_eq!(step_pos_ids.size(1), prompt.len());
                assert!(step_pos_ids.iter().map(|x| *x as usize).eq(0..prompt.len()));

                assert_eq!(step_cache_pos.size(0), prompt.len());
                assert!(
                    step_cache_pos
                        .iter()
                        .map(|x| *x as usize)
                        .eq(0..prompt.len())
                );

                if let Some(cache_branch) = cache_branch {
                    assert_eq!(cache_branch.item(), Some(&0));
                }
            } else if kv_cache_type.is_some() {
                assert_eq!(step_inputs.size(1), 1);
                assert_eq!(step_inputs[[0, 0]] as u32, expected_token_ids[step - 1]);

                assert_eq!(step_attn_mask.size(1), prompt.len() + step);
                assert_eq!(step_attn_mask[[0, 0]], 1);

                assert_eq!(step_pos_ids.size(1), 1);
                assert_eq!(step_pos_ids[[0, 0]], (prompt.len() + step - 1) as i32);

                assert_eq!(step_cache_pos.size(0), 1);
                assert_eq!(step_cache_pos[[0]], (prompt.len() + step - 1) as i32);

                if let Some(cache_branch) = cache_branch {
                    assert_eq!(cache_branch.item(), Some(&1));
                }
            } else {
                let expected_inputs: Vec<i32> = prompt
                    .iter()
                    .copied()
                    .chain(expected_token_ids)
                    .take(prompt.len() + step)
                    .map(|x| x as i32)
                    .collect();
                assert_eq!(
                    step_inputs,
                    NdTensor::from_data([1, expected_inputs.len()], expected_inputs)
                );

                let expected_attn_mask = vec![1i32; prompt.len() + step];
                assert_eq!(
                    step_attn_mask,
                    NdTensor::from_data([1, expected_attn_mask.len()], expected_attn_mask)
                );

                let expected_pos_ids: Vec<i32> =
                    (0..prompt.len() + step).map(|x| x as i32).collect();
                assert_eq!(
                    step_pos_ids,
                    NdTensor::from_data([1, expected_pos_ids.len()], expected_pos_ids.clone())
                );
                assert_eq!(
                    step_cache_pos,
                    NdTensor::from_data([expected_pos_ids.len()], expected_pos_ids)
                );
            }
        }

        Ok(())
    }

    #[test]
    fn test_generator_with_decoder_kv_cache() -> Result<(), Box<dyn Error>> {
        test_generator_impl(Some(KvCacheType::Decoder))
    }

    #[test]
    fn test_generator_with_encoder_decoder_kv_cache() -> Result<(), Box<dyn Error>> {
        test_generator_impl(Some(KvCacheType::EncoderDecoder))
    }

    #[test]
    fn test_generator_without_kv_cache() -> Result<(), Box<dyn Error>> {
        test_generator_impl(None)
    }

    #[test]
    fn test_generator_append_prompt() -> Result<(), Box<dyn Error>> {
        let mut params = TransformerParams::default();
        params.n_vocab = 110;
        let output_token_ids = [0, 1, 2, 3, 4, 5, 6, 7, 8];
        let prompt = [99];
        let model = fake_transformer_model(
            params,
            Some(KvCacheType::Decoder),
            prompt.len(),
            &output_token_ids,
        );

        let mut generator = Generator::from_model(&model)?.with_prompt(&prompt);

        generator.next();
        generator.append_prompt(&[100]);
        generator.next();
        generator.append_prompt(&[101, 102]);
        generator.next();

        let input_id = model.find_node("input_ids").unwrap();

        // The input to the first step is just the prompt.
        let inputs = model.get_inputs(0, input_id).unwrap();
        let inputs: NdTensor<i32, 2> = inputs.try_into().unwrap();
        assert_eq!(inputs, NdTensor::from([[99]]));

        // The inputs for the next steps are the output followed by the inputs
        // added with `append_prompt`.
        let inputs = model.get_inputs(1, input_id).unwrap();
        let inputs: NdTensor<i32, 2> = inputs.try_into().unwrap();
        assert_eq!(inputs, NdTensor::from([[0, 100]]));

        let inputs = model.get_inputs(2, input_id).unwrap();
        let inputs: NdTensor<i32, 2> = inputs.try_into().unwrap();
        assert_eq!(inputs, NdTensor::from([[1, 101, 102]]));

        Ok(())
    }

    #[test]
    fn test_stop_on_tokens() -> Result<(), Box<dyn Error>> {
        let params = TransformerParams::default();
        let expected_token_ids = [0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 0, 0, 0];
        let prompt = [1, 2, 3, 1, 2, 3];
        let model = fake_transformer_model(
            params,
            Some(KvCacheType::Decoder),
            prompt.len(),
            &expected_token_ids,
        );

        let generator = Generator::from_model(&model)?;

        let output_token_ids: Vec<_> = generator
            .with_prompt(&prompt)
            .stop_on_tokens([4])
            .map(|id| id.expect("generation failed"))
            .collect();

        assert_eq!(output_token_ids, &[0, 1, 2, 3]);

        Ok(())
    }

    #[test]
    fn test_process_prompt() -> Result<(), Box<dyn Error>> {
        let params = TransformerParams::default();
        let expected_token_ids = [0];
        let prompt = [1, 2, 3, 1, 2, 3];
        let model = fake_transformer_model(
            params,
            Some(KvCacheType::Decoder),
            prompt.len(),
            &expected_token_ids,
        );

        let mut generator = Generator::from_model(&model)?.with_prompt(&prompt);
        assert_eq!(generator.prompt(), prompt);
        assert!(generator.prev_tokens().is_empty());
        assert_eq!(generator.kv_cache_len(), Some(0));

        generator.process_prompt().unwrap();

        assert!(generator.prompt().is_empty());
        assert_eq!(generator.prev_tokens(), prompt);
        assert_eq!(generator.kv_cache_len(), Some(prompt.len()));

        Ok(())
    }

    #[test]
    fn test_profile() -> Result<(), Box<dyn Error>> {
        let params = TransformerParams::default();
        let expected_token_ids = [0, 1, 2, 3, 4];
        let prompt = [1, 2, 3, 1, 2, 3];
        let model = fake_transformer_model(
            params,
            Some(KvCacheType::Decoder),
            prompt.len(),
            &expected_token_ids,
        );

        let generator = Generator::from_model(&model)?;
        let mut metrics = Metrics::new();

        let output_token_ids: Vec<_> = generator
            .with_prompt(&prompt)
            .profile(&mut metrics)
            .take(expected_token_ids.len())
            .map(|id| id.expect("generation failed"))
            .collect();

        assert_eq!(output_token_ids, expected_token_ids);
        assert!(metrics.warmup_duration().is_some());
        assert_eq!(metrics.step_durations().len(), output_token_ids.len() - 1);

        Ok(())
    }

    #[test]
    fn test_filter() -> Result<(), Box<dyn Error>> {
        let mut params = TransformerParams::default();
        params.n_vocab = 8; // Must be >2x the max token ID in `expected_token_ids`.

        let expected_token_ids = [0, 1, 2, 3];
        let prompt = [5, 6, 7];
        let model = fake_transformer_model(
            params,
            Some(KvCacheType::Decoder),
            prompt.len(),
            &expected_token_ids,
        );

        let generator = Generator::from_model(&model)?;

        // Filter that modifies logits to double the selected token ID.
        struct DoubleIndexFilter {
            prev_tokens: Rc<RefCell<Vec<u32>>>,
        }
        impl LogitsFilter for DoubleIndexFilter {
            fn filter(&self, logits: Logits, prev_tokens: &[u32]) -> Logits {
                self.prev_tokens.replace(prev_tokens.to_vec());

                let max_idx = logits
                    .logits()
                    .iter()
                    .zip(logits.indices())
                    .max_by(|(x, _i), (y, _j)| x.total_cmp(y))
                    .map(|(_x, i)| i)
                    .unwrap();

                Logits::sparse(vec![1.0], vec![max_idx * 2])
            }
        }

        let prev_tokens = Rc::new(RefCell::new(Vec::new()));
        let output_token_ids: Vec<_> = generator
            .with_prompt(&prompt)
            .with_logits_filter(DoubleIndexFilter {
                prev_tokens: prev_tokens.clone(),
            })
            .take(expected_token_ids.len())
            .map(|id| id.expect("generation failed"))
            .collect();

        assert_eq!(output_token_ids, [0, 2, 4, 6]);
        assert_eq!(prev_tokens.borrow().as_slice(), [5, 6, 7, 0, 2, 4]);

        Ok(())
    }

    #[test]
    fn test_empty_filter_output() {
        let params = TransformerParams::default();
        let prompt = [1];
        let model = fake_transformer_model(
            params,
            Some(KvCacheType::Decoder),
            prompt.len(),
            &[0, 1, 2, 3],
        );

        struct RemoveAllFilter;
        impl LogitsFilter for RemoveAllFilter {
            fn filter(&self, _logits: Logits, _prev_tokens: &[u32]) -> Logits {
                Logits::dense(vec![])
            }
        }

        let mut generator = Generator::from_model(&model)
            .unwrap()
            .with_logits_filter(RemoveAllFilter);
        let err = generator.next().unwrap().err().unwrap();
        assert!(err.to_string().contains("filtered logits are empty"));
    }

    #[test]
    fn test_run_options() -> Result<(), Box<dyn Error>> {
        let params = TransformerParams::default();
        let expected_token_ids = [0, 1, 2, 3, 4];
        let prompt = [1, 2, 3, 1, 2, 3];
        let model = fake_transformer_model(
            params,
            Some(KvCacheType::Decoder),
            prompt.len(),
            &expected_token_ids,
        );

        let generator = Generator::from_model(&model)?;

        let run_opts = RunOptions::default().with_verbose(true);
        let output_token_ids: Vec<_> = generator
            .with_prompt(&prompt)
            .with_run_options(Some(run_opts.clone()))
            .take(expected_token_ids.len())
            .map(|id| id.expect("generation failed"))
            .collect();

        assert_eq!(output_token_ids, expected_token_ids);
        assert_eq!(model.run_opts.take(), Some(run_opts));

        Ok(())
    }
}