diff --git a/llama-cpp-bindings-tests/tests/multimodal_vision.rs b/llama-cpp-bindings-tests/tests/multimodal_vision.rs index 874868a4..8a330da9 100644 --- a/llama-cpp-bindings-tests/tests/multimodal_vision.rs +++ b/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 { diff --git a/llama-cpp-bindings/src/error/eval_multimodal_chunks_error.rs b/llama-cpp-bindings/src/error/eval_multimodal_chunks_error.rs index d5d485f7..7874616f 100644 --- a/llama-cpp-bindings/src/error/eval_multimodal_chunks_error.rs +++ b/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), } diff --git a/llama-cpp-bindings/src/ingest_prompt_chunk.rs b/llama-cpp-bindings/src/ingest_prompt_chunk.rs deleted file mode 100644 index 3f569306..00000000 --- a/llama-cpp-bindings/src/ingest_prompt_chunk.rs +++ /dev/null @@ -1,26 +0,0 @@ -use crate::mtmd::MtmdInputChunk; -use crate::mtmd::MtmdInputChunkType; -use crate::mtmd::MtmdInputChunkTypeError; -use crate::sampled_token_classifier::SampledTokenClassifier; - -/// # Errors -/// Returns [`MtmdInputChunkTypeError`] when the chunk reports a type unknown -/// to this binding. Counters are not updated on error. -pub fn ingest_prompt_chunk( - classifier: &mut SampledTokenClassifier<'_>, - chunk: &MtmdInputChunk, -) -> Result<(), MtmdInputChunkTypeError> { - let n_tokens = chunk.n_tokens() as u64; - match chunk.chunk_type()? { - MtmdInputChunkType::Text => { - classifier.record_prompt_tokens(n_tokens); - if let Some(tokens) = chunk.text_tokens()? { - classifier.ingest_prompt_tokens(tokens); - } - } - MtmdInputChunkType::Image => classifier.record_input_image_tokens(n_tokens), - MtmdInputChunkType::Audio => classifier.record_input_audio_tokens(n_tokens), - } - - Ok(()) -} diff --git a/llama-cpp-bindings/src/lib.rs b/llama-cpp-bindings/src/lib.rs index 387af605..b9899f37 100644 --- a/llama-cpp-bindings/src/lib.rs +++ b/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; diff --git a/llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunk.rs b/llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunk.rs index 2023fbb1..f5186da7 100644 --- a/llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunk.rs +++ b/llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunk.rs @@ -3,6 +3,7 @@ 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; @@ -10,6 +11,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::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, { + 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, diff --git a/llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunks.rs b/llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunks.rs index fbc6ad5e..7ae238ad 100644 --- a/llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunks.rs +++ b/llama-cpp-bindings/src/mtmd/fitted_mtmd_input_chunks.rs @@ -3,6 +3,7 @@ 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)] @@ -10,6 +11,7 @@ pub struct FittedMtmdInputChunks where TChunks: Borrow, { + chunk_types: Vec, chunks: TChunks, micro_batch_tokens: u32, } @@ -19,7 +21,7 @@ where TChunks: Borrow, { /// 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 { - 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 { + 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] diff --git a/llama-cpp-bindings/src/mtmd/mtmd_input_chunk.rs b/llama-cpp-bindings/src/mtmd/mtmd_input_chunk.rs index 4011d418..dce0ac17 100644 --- a/llama-cpp-bindings/src/mtmd/mtmd_input_chunk.rs +++ b/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, 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 { + 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) } } diff --git a/llama-cpp-bindings/src/mtmd/mtmd_input_chunks.rs b/llama-cpp-bindings/src/mtmd/mtmd_input_chunks.rs index 22b43101..113ada7f 100644 --- a/llama-cpp-bindings/src/mtmd/mtmd_input_chunks.rs +++ b/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, 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] diff --git a/llama-cpp-bindings/src/sampled_token_classifier.rs b/llama-cpp-bindings/src/sampled_token_classifier.rs index e95ce8fc..8c9c8c0c 100644 --- a/llama-cpp-bindings/src/sampled_token_classifier.rs +++ b/llama-cpp-bindings/src/sampled_token_classifier.rs @@ -21,6 +21,7 @@ use crate::model::LlamaModel; use crate::mtmd::FittedMtmdInputChunk; use crate::mtmd::FittedMtmdInputChunks; use crate::mtmd::MtmdContext; +use crate::mtmd::MtmdInputChunkType; use crate::mtmd::MtmdInputChunks; use crate::mtmd::micro_batch_tokens; use crate::mtmd::positive_batch_tokens::positive_batch_tokens; @@ -455,11 +456,9 @@ impl<'model> SampledTokenClassifier<'model> { /// them, so a chunk that cannot be decoded moves neither the KV cache nor the usage counters. /// /// # Errors - /// Returns [`EvalMultimodalChunksError::EvalFailed`] when `params.n_batch` is not positive or - /// a chunk cannot be decoded (both before any chunk is evaluated), or when evaluating a - /// chunk fails, - /// [`EvalMultimodalChunksError::UnknownChunkType`] when a chunk reports a - /// type unknown to this binding, or + /// Returns [`EvalMultimodalChunksError::EvalFailed`] when `params.n_batch` is not positive, a + /// chunk cannot be decoded or reports a type unknown to this binding (all before any chunk is + /// evaluated), or when evaluating a chunk fails, or /// [`EvalMultimodalChunksError::ChunkOutOfBounds`] when a valid index returns /// `None` from `chunks.get`. pub fn eval_multimodal_chunks( @@ -490,12 +489,25 @@ impl<'model> SampledTokenClassifier<'model> { params.n_batch, logits_for_this_chunk, )?; - crate::ingest_prompt_chunk::ingest_prompt_chunk(self, fitted_chunk.chunk())?; + self.ingest_fitted_chunk(&fitted_chunk); } Ok(next_position) } + pub fn ingest_fitted_chunk(&mut self, fitted_chunk: &FittedMtmdInputChunk<'_>) { + let n_tokens = fitted_chunk.chunk().n_tokens() as u64; + + match fitted_chunk.chunk_type() { + MtmdInputChunkType::Text => { + self.record_prompt_tokens(n_tokens); + self.ingest_prompt_tokens(fitted_chunk.text_tokens()); + } + MtmdInputChunkType::Image => self.record_input_image_tokens(n_tokens), + MtmdInputChunkType::Audio => self.record_input_audio_tokens(n_tokens), + } + } + pub const fn record_prompt_tokens(&mut self, count: u64) { self.usage.record_prompt_tokens(count); }