Keep fitted chunk types and ingest fitted chunks without classifying them again

llama-cpp-bindings-tests/tests/multimodal_vision.rs

@@ -5,7 +5,6 @@ use llama_cpp_bindings::EvalMultimodalChunksParams;
 use llama_cpp_bindings::TokenUsage;
 use llama_cpp_bindings::context::LlamaContext;
 use llama_cpp_bindings::error::EvalMultimodalChunksError;
-use llama_cpp_bindings::ingest_prompt_chunk::ingest_prompt_chunk;
 use llama_cpp_bindings::llama_batch::LlamaBatch;
 use llama_cpp_bindings::mtmd::FittedMtmdInputChunk;
 use llama_cpp_bindings::mtmd::FittedMtmdInputChunks;
@@ -1407,20 +1406,29 @@ fn text_chunk_records_prompt_tokens(fixture: &LlamaFixture<'_>) -> Result<()> {
         add_special: false,
         parse_special: false,
     };
-    let chunks = mtmd_ctx.tokenize(input_text, &[])?;
+    let llama_ctx = LlamaContext::from_model(
+        model,
+        fixture.backend,
+        (*fixture.context_params).into_llama_context_params(),
+    )?;
+    let fitted_chunks = FittedMtmdInputChunks::fit(
+        mtmd_ctx.tokenize(input_text, &[])?,
+        mtmd_ctx,
+        llama_ctx.n_ubatch(),
+    )?;
 
-    let text_chunk = (0..chunks.len())
-        .filter_map(|index| chunks.get(index))
-        .find(|chunk| chunk.chunk_type() == Ok(MtmdInputChunkType::Text))
+    let text_chunk = (0..fitted_chunks.len())
+        .filter_map(|index| FittedMtmdInputChunk::at(&fitted_chunks, index))
+        .find(|chunk| chunk.chunk_type() == MtmdInputChunkType::Text)
         .ok_or_else(|| {
             anyhow::anyhow!("text-only tokenization should produce at least one text chunk")
         })?;
 
-    let n_tokens = u64::try_from(text_chunk.n_tokens())?;
+    let n_tokens = u64::try_from(text_chunk.chunk().n_tokens())?;
 
     let mut classifier = model.sampled_token_classifier(BareJsonToolCalls::Detect)?;
 
-    ingest_prompt_chunk(&mut classifier, &text_chunk)?;
+    classifier.ingest_fitted_chunk(&text_chunk);
 
     let usage = classifier.usage();
     if usage.prompt_tokens != n_tokens {
@@ -1472,21 +1480,30 @@ fn image_chunk_records_input_image_tokens_only(fixture: &LlamaFixture<'_>) -> Re
         add_special: false,
         parse_special: true,
     };
-    let chunks = mtmd_ctx.tokenize(input_text, &[&bitmap])?;
+    let llama_ctx = LlamaContext::from_model(
+        model,
+        fixture.backend,
+        (*fixture.context_params).into_llama_context_params(),
+    )?;
+    let fitted_chunks = FittedMtmdInputChunks::fit(
+        mtmd_ctx.tokenize(input_text, &[&bitmap])?,
+        mtmd_ctx,
+        llama_ctx.n_ubatch(),
+    )?;
 
-    let image_chunk = (0..chunks.len())
-        .filter_map(|index| chunks.get(index))
-        .find(|chunk| chunk.chunk_type() == Ok(MtmdInputChunkType::Image))
+    let image_chunk = (0..fitted_chunks.len())
+        .filter_map(|index| FittedMtmdInputChunk::at(&fitted_chunks, index))
+        .find(|chunk| chunk.chunk_type() == MtmdInputChunkType::Image)
         .ok_or_else(|| anyhow::anyhow!("multimodal tokenization should produce an image chunk"))?;
 
-    let n_tokens = u64::try_from(image_chunk.n_tokens())?;
+    let n_tokens = u64::try_from(image_chunk.chunk().n_tokens())?;
     if n_tokens == 0 {
         anyhow::bail!("image chunk should report at least one token");
     }
 
     let mut classifier = model.sampled_token_classifier(BareJsonToolCalls::Detect)?;
 
-    ingest_prompt_chunk(&mut classifier, &image_chunk)?;
+    classifier.ingest_fitted_chunk(&image_chunk);
 
     let usage = classifier.usage();
     if usage.input_image_tokens != n_tokens {
@@ -1531,15 +1548,23 @@ fn text_chunk_drives_marker_state_machine_to_reasoning(fixture: &LlamaFixture<'_
         add_special: false,
         parse_special: true,
     };
-    let chunks = mtmd_ctx.tokenize(input_text, &[])?;
+    let llama_ctx = LlamaContext::from_model(
+        model,
+        fixture.backend,
+        (*fixture.context_params).into_llama_context_params(),
+    )?;
+    let fitted_chunks = FittedMtmdInputChunks::fit(
+        mtmd_ctx.tokenize(input_text, &[])?,
+        mtmd_ctx,
+        llama_ctx.n_ubatch(),
+    )?;
 
     let mut classifier = model.sampled_token_classifier(BareJsonToolCalls::Detect)?;
 
-    for index in 0..chunks.len() {
-        let chunk = chunks
-            .get(index)
+    for index in 0..fitted_chunks.len() {
+        let chunk = FittedMtmdInputChunk::at(&fitted_chunks, index)
             .ok_or_else(|| anyhow::anyhow!("chunk index {index} must exist"))?;
-        ingest_prompt_chunk(&mut classifier, &chunk)?;
+        classifier.ingest_fitted_chunk(&chunk);
     }
 
     if classifier.current_section() != llama_cpp_bindings::SampledTokenSection::Reasoning {

llama-cpp-bindings/src/error/eval_multimodal_chunks_error.rs

@@ -1,12 +1,9 @@
 use crate::mtmd::MtmdEvalError;
-use crate::mtmd::mtmd_input_chunk_type_error::MtmdInputChunkTypeError;
 
 #[derive(Debug, thiserror::Error)]
 pub enum EvalMultimodalChunksError {
     #[error("{0}")]
     EvalFailed(#[from] MtmdEvalError),
-    #[error("{0}")]
-    UnknownChunkType(#[from] MtmdInputChunkTypeError),
     #[error("chunk index {0} out of bounds during post-eval walk")]
     ChunkOutOfBounds(usize),
 }

llama-cpp-bindings/src/lib.rs

@@ -21,7 +21,6 @@ pub mod gguf_tensor_f32;
 pub mod gguf_type;
 pub mod grammar_matcher;
 pub mod ingest_outcome;
-pub mod ingest_prompt_chunk;
 pub mod invalid_numa_strategy;
 pub mod json_probe_outcome;
 pub mod json_schema_to_grammar;
@@ -105,7 +104,6 @@ pub use synthetic_tool_call_renders::SyntheticToolCallRenders;
 pub use token_piece::TokenPiece;
 
 pub use ggml_time_us::ggml_time_us;
-pub use ingest_prompt_chunk::ingest_prompt_chunk;
 pub use json_schema_to_grammar::json_schema_to_grammar;
 pub use llama_time_us::llama_time_us;
 pub use max_devices::max_devices;

llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunk.rs

@@ -3,13 +3,15 @@ use std::ffi::c_char;
 use std::marker::PhantomData;
 
 use crate::context::LlamaContext;
+use crate::token::LlamaToken;
 use llama_cpp_ffi_status::read_and_free_cpp_string;
 
 use super::fitted_mtmd_input_chunks::FittedMtmdInputChunks;
 use super::micro_batch_tokens::micro_batch_tokens;
 use super::mtmd_context::MtmdContext;
 use super::mtmd_eval_error::MtmdEvalError;
 use super::mtmd_input_chunk::MtmdInputChunk;
+use super::mtmd_input_chunk_type::MtmdInputChunkType;
 use super::mtmd_input_chunks::MtmdInputChunks;
 use super::positive_batch_tokens::positive_batch_tokens;
 
@@ -81,6 +83,7 @@ fn eval_chunk_single_status_to_result(
 #[derive(Debug)]
 pub struct FittedMtmdInputChunk<'fitted> {
     chunk: MtmdInputChunk,
+    chunk_type: MtmdInputChunkType,
     fitted_chunks: PhantomData<&'fitted MtmdInputChunks>,
     micro_batch_tokens: u32,
 }
@@ -94,8 +97,11 @@ impl<'fitted> FittedMtmdInputChunk<'fitted> {
     where
         TChunks: Borrow<MtmdInputChunks>,
     {
+        let chunk_type = fitted_chunks.chunk_type(chunk_index)?;
+
         fitted_chunks.get(chunk_index).map(|chunk| Self {
             chunk,
+            chunk_type,
             fitted_chunks: PhantomData,
             micro_batch_tokens: fitted_chunks.micro_batch_tokens(),
         })
@@ -106,6 +112,18 @@ impl<'fitted> FittedMtmdInputChunk<'fitted> {
         &self.chunk
     }
 
+    #[must_use]
+    pub const fn chunk_type(&self) -> MtmdInputChunkType {
+        self.chunk_type
+    }
+
+    #[must_use]
+    pub fn text_tokens(&self) -> &[LlamaToken] {
+        self.chunk
+            .text_tokens_of(self.chunk_type)
+            .unwrap_or_default()
+    }
+
     /// # Errors
     ///
     /// Returns [`MtmdEvalError::NonPositiveBatchSize`] when `n_batch` is not positive,

llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunks.rs

@@ -3,13 +3,15 @@ use std::borrow::Borrow;
 use super::mtmd_context::MtmdContext;
 use super::mtmd_eval_error::MtmdEvalError;
 use super::mtmd_input_chunk::MtmdInputChunk;
+use super::mtmd_input_chunk_type::MtmdInputChunkType;
 use super::mtmd_input_chunks::MtmdInputChunks;
 
 #[derive(Debug)]
 pub struct FittedMtmdInputChunks<TChunks>
 where
     TChunks: Borrow<MtmdInputChunks>,
 {
+    chunk_types: Vec<MtmdInputChunkType>,
     chunks: TChunks,
     micro_batch_tokens: u32,
 }
@@ -19,7 +21,7 @@ where
     TChunks: Borrow<MtmdInputChunks>,
 {
     /// Checks every chunk once with [`MtmdInputChunks::fit_to_micro_batch`], so the chunks can
-    /// be evaluated without checking them again.
+    /// be evaluated and ingested without checking or classifying them again.
     ///
     /// # Errors
     ///
@@ -29,11 +31,12 @@ where
         mtmd_ctx: &MtmdContext,
         micro_batch_tokens: u32,
     ) -> Result<Self, MtmdEvalError> {
-        chunks
+        let chunk_types = chunks
             .borrow()
             .fit_to_micro_batch(mtmd_ctx, micro_batch_tokens)?;
 
         Ok(Self {
+            chunk_types,
             chunks,
             micro_batch_tokens,
         })
@@ -45,13 +48,18 @@ where
     }
 
     #[must_use]
-    pub fn len(&self) -> usize {
-        self.chunks.borrow().len()
+    pub fn chunk_type(&self, chunk_index: usize) -> Option<MtmdInputChunkType> {
+        self.chunk_types.get(chunk_index).copied()
     }
 
     #[must_use]
-    pub fn is_empty(&self) -> bool {
-        self.chunks.borrow().is_empty()
+    pub const fn len(&self) -> usize {
+        self.chunk_types.len()
+    }
+
+    #[must_use]
+    pub const fn is_empty(&self) -> bool {
+        self.chunk_types.is_empty()
     }
 
     #[must_use]

llama-cpp-bindings/src/mtmd/mtmd_input_chunk.rs

@@ -71,8 +71,13 @@ impl MtmdInputChunk {
     /// binding does not know, so an unclassifiable chunk is never mistaken for a
     /// non-text chunk.
     pub fn text_tokens(&self) -> Result<Option<&[LlamaToken]>, MtmdInputChunkTypeError> {
-        if self.chunk_type()? != MtmdInputChunkType::Text {
-            return Ok(None);
+        Ok(self.text_tokens_of(self.chunk_type()?))
+    }
+
+    #[must_use]
+    pub fn text_tokens_of(&self, chunk_type: MtmdInputChunkType) -> Option<&[LlamaToken]> {
+        if chunk_type != MtmdInputChunkType::Text {
+            return None;
         }
 
         let mut n_tokens = 0usize;
@@ -83,7 +88,7 @@ impl MtmdInputChunk {
             )
         };
 
-        Ok(unsafe { tokens_from_raw_ptr(tokens_ptr, n_tokens) })
+        unsafe { tokens_from_raw_ptr(tokens_ptr, n_tokens) }
     }
 
     #[must_use]
@@ -119,9 +124,9 @@ impl MtmdInputChunk {
         Ok(Self { chunk, owned: true })
     }
 
-    /// Checks that this chunk can be evaluated in decodes of `micro_batch_tokens`. llama.cpp
-    /// splits a causal chunk across decodes, while a media chunk it decodes non-causally has to
-    /// fit a single decode.
+    /// Checks that this chunk can be evaluated in decodes of `micro_batch_tokens` and returns the
+    /// chunk type it classified. llama.cpp splits a causal chunk across decodes, while a media
+    /// chunk it decodes non-causally has to fit a single decode.
     ///
     /// # Errors
     ///
@@ -132,16 +137,17 @@ impl MtmdInputChunk {
         &self,
         mtmd_ctx: &MtmdContext,
         micro_batch_tokens: u32,
-    ) -> Result<(), MtmdEvalError> {
+    ) -> Result<MtmdInputChunkType, MtmdEvalError> {
+        let chunk_type = self.chunk_type()?;
         let decodes_non_causally =
-            self.chunk_type()? != MtmdInputChunkType::Text && mtmd_ctx.decode_use_non_causal(self);
+            chunk_type != MtmdInputChunkType::Text && mtmd_ctx.decode_use_non_causal(self);
 
         non_causal_chunk_micro_batch_error(
             decodes_non_causally,
             self.n_tokens(),
             micro_batch_tokens,
         )
-        .map_or(Ok(()), Err)
+        .map_or(Ok(chunk_type), Err)
     }
 }
 

llama-cpp-bindings/src/mtmd/mtmd_input_chunks.rs

@@ -6,6 +6,7 @@ use super::micro_batch_tokens::micro_batch_tokens;
 use super::mtmd_context::MtmdContext;
 use super::mtmd_eval_error::MtmdEvalError;
 use super::mtmd_input_chunk::MtmdInputChunk;
+use super::mtmd_input_chunk_type::MtmdInputChunkType;
 use super::mtmd_input_chunks_error::MtmdInputChunksError;
 use super::positive_batch_tokens::positive_batch_tokens;
 
@@ -63,7 +64,7 @@ impl MtmdInputChunks {
     }
 
     /// Checks every chunk with [`MtmdInputChunk::fit_to_micro_batch`] before any of them is
-    /// evaluated.
+    /// evaluated, and returns the chunk types it classified, in chunk order.
     ///
     /// # Errors
     ///
@@ -75,17 +76,17 @@ impl MtmdInputChunks {
         &self,
         mtmd_ctx: &MtmdContext,
         micro_batch_tokens: u32,
-    ) -> Result<(), MtmdEvalError> {
-        for index in 0..self.len() {
-            self.get(index)
-                .ok_or(crate::FfiContractError {
-                    operation: "mtmd_input_chunks_get",
-                    detail: "returned a null chunk within the chunk count",
-                })?
-                .fit_to_micro_batch(mtmd_ctx, micro_batch_tokens)?;
-        }
-
-        Ok(())
+    ) -> Result<Vec<MtmdInputChunkType>, MtmdEvalError> {
+        (0..self.len())
+            .map(|index| {
+                self.get(index)
+                    .ok_or(crate::FfiContractError {
+                        operation: "mtmd_input_chunks_get",
+                        detail: "returned a null chunk within the chunk count",
+                    })?
+                    .fit_to_micro_batch(mtmd_ctx, micro_batch_tokens)
+            })
+            .collect()
     }
 
     #[must_use]