Method next_token

cake-core/src/models/common/text_model.rs:397–495

Generate the next token. Assumes `prepare_prompt()` has been called for the first token.

(&mut self, index: usize)

    /// Generate the next token. Assumes `prepare_prompt()` has been called for the first token.
    pub async fn next_token(&mut self, index: usize) -> Result<Token> {
        log::trace!("model.next_token({index})");

        let num_tokens = self.tokens.len();
        let (context_size, context_index) = if self
            .ctx
            .cache
            .as_ref()
            .expect("No cache specified")
            .with_kv_cache()
            && index > 0
        {
            (1, self.index_pos)
        } else {
            (num_tokens, 0)
        };

        let context_offset = num_tokens.saturating_sub(context_size);
        let context_tokens = &self.tokens[context_offset..];
        let num_context_tokens = context_tokens.len();

        let input = Tensor::new(context_tokens, &self.ctx.device)?
            .unsqueeze(0)
            .map_err(|e| anyhow!("error squeezing context tokens: {e}"))?;

        let logits = self
            .forward(&input, context_index)
            .await
            .map_err(|e| anyhow!("error in model.forward: {e}"))?;

        let post_start = std::time::Instant::now();

        let logits = logits
            .squeeze(0)
            .map_err(|e| anyhow!("error squeezing logits: {e}"))?;

        // Apply repeat penalty only to generated tokens (not prompt tokens)
        let penalty_start = std::time::Instant::now();
        let logits = if self.ctx.args.repeat_penalty == 1. {
            logits
        } else {
            let generated_start = self.prompt_len;
            let penalty_tokens = &self.tokens[generated_start..];
            if penalty_tokens.is_empty() {
                logits
            } else {
                let start_at = penalty_tokens
                    .len()
                    .saturating_sub(self.ctx.args.repeat_last_n);
                apply_repeat_penalty_gpu(
                    &logits,
                    self.ctx.args.repeat_penalty,
                    &penalty_tokens[start_at..],
                )?
            }
        };
        let penalty_elapsed = penalty_start.elapsed();
        self.index_pos += num_context_tokens;