llama_cpp_bindings/mtmd/
mtmd_input_chunk.rs

1use std::ffi::CStr;
2use std::ffi::c_char;
3use std::ptr::NonNull;
4use std::slice;
5
6use crate::context::LlamaContext;
7use crate::ffi_error_reader::read_and_free_cpp_error;
8use crate::token::LlamaToken;
9
10use super::image_chunk_batch_size_mismatch::ImageChunkBatchSizeMismatch;
11use super::mtmd_context::MtmdContext;
12use super::mtmd_eval_error::MtmdEvalError;
13use super::mtmd_input_chunk_error::MtmdInputChunkError;
14use super::mtmd_input_chunk_type::MtmdInputChunkType;
15use super::mtmd_input_chunk_type_error::MtmdInputChunkTypeError;
16
17/// # Safety
18///
19/// `tokens_ptr` must point to at least `n_tokens` valid `llama_token` values
20/// that remain valid for the lifetime `'chunk`.
21const unsafe fn tokens_from_raw_ptr<'chunk>(
22    tokens_ptr: *const llama_cpp_bindings_sys::llama_token,
23    n_tokens: usize,
24) -> Option<&'chunk [LlamaToken]> {
25    if tokens_ptr.is_null() || n_tokens == 0 {
26        None
27    } else {
28        unsafe {
29            Some(slice::from_raw_parts(
30                tokens_ptr.cast::<LlamaToken>(),
31                n_tokens,
32            ))
33        }
34    }
35}
36
37fn eval_chunk_single_status_to_result(
38    status: llama_cpp_bindings_sys::llama_rs_mtmd_eval_chunk_single_status,
39    final_position: llama_cpp_bindings_sys::llama_pos,
40    out_vendored_return_code: i32,
41    out_error: *mut c_char,
42) -> Result<llama_cpp_bindings_sys::llama_pos, MtmdEvalError> {
43    match status {
44        llama_cpp_bindings_sys::LLAMA_RS_MTMD_EVAL_CHUNK_SINGLE_OK => Ok(final_position),
45        llama_cpp_bindings_sys::LLAMA_RS_MTMD_EVAL_CHUNK_SINGLE_VENDORED_RETURNED_NONZERO_CODE => {
46            Err(MtmdEvalError::EvalFailed {
47                code: out_vendored_return_code,
48            })
49        }
50        llama_cpp_bindings_sys::LLAMA_RS_MTMD_EVAL_CHUNK_SINGLE_ERROR_STRING_ALLOCATION_FAILED => {
51            Err(MtmdEvalError::NotEnoughMemory)
52        }
53        llama_cpp_bindings_sys::LLAMA_RS_MTMD_EVAL_CHUNK_SINGLE_VENDORED_THREW_CXX_EXCEPTION => {
54            let message = unsafe { read_and_free_cpp_error(out_error) };
55            Err(MtmdEvalError::Reported { message })
56        }
57        other => {
58            unreachable!("llama_rs_mtmd_eval_chunk_single returned unrecognized status: {other}")
59        }
60    }
61}
62
63fn image_chunk_batch_size_error(
64    is_image_chunk: bool,
65    chunk_token_count: usize,
66    n_batch: i32,
67) -> Option<MtmdEvalError> {
68    if is_image_chunk
69        && i64::try_from(chunk_token_count).is_ok_and(|tokens| tokens > i64::from(n_batch))
70    {
71        #[expect(
72            clippy::cast_possible_truncation,
73            clippy::cast_sign_loss,
74            reason = "image token counts and n_batch are model-bounded and fit in u32"
75        )]
76        return Some(MtmdEvalError::ImageChunkExceedsBatchSize(
77            ImageChunkBatchSizeMismatch {
78                image_tokens: chunk_token_count as u32,
79                n_batch: n_batch as u32,
80            },
81        ));
82    }
83
84    None
85}
86
87#[derive(Debug)]
88pub struct MtmdInputChunk {
89    pub chunk: NonNull<llama_cpp_bindings_sys::mtmd_input_chunk>,
90    pub owned: bool,
91}
92
93impl MtmdInputChunk {
94    /// # Errors
95    /// Returns an error if the chunk type is unknown.
96    pub fn chunk_type(&self) -> Result<MtmdInputChunkType, MtmdInputChunkTypeError> {
97        let chunk_type =
98            unsafe { llama_cpp_bindings_sys::mtmd_input_chunk_get_type(self.chunk.as_ptr()) };
99        MtmdInputChunkType::try_from(chunk_type)
100    }
101
102    #[must_use]
103    pub fn text_tokens(&self) -> Option<&[LlamaToken]> {
104        if self.chunk_type() != Ok(MtmdInputChunkType::Text) {
105            return None;
106        }
107
108        let mut n_tokens = 0usize;
109        let tokens_ptr = unsafe {
110            llama_cpp_bindings_sys::mtmd_input_chunk_get_tokens_text(
111                self.chunk.as_ptr(),
112                &raw mut n_tokens,
113            )
114        };
115
116        unsafe { tokens_from_raw_ptr(tokens_ptr, n_tokens) }
117    }
118
119    #[must_use]
120    pub fn n_tokens(&self) -> usize {
121        unsafe { llama_cpp_bindings_sys::mtmd_input_chunk_get_n_tokens(self.chunk.as_ptr()) }
122    }
123
124    #[must_use]
125    pub fn n_positions(&self) -> i32 {
126        unsafe { llama_cpp_bindings_sys::mtmd_input_chunk_get_n_pos(self.chunk.as_ptr()) }
127    }
128
129    #[must_use]
130    pub fn id(&self) -> Option<String> {
131        let ptr = unsafe { llama_cpp_bindings_sys::mtmd_input_chunk_get_id(self.chunk.as_ptr()) };
132        if ptr.is_null() {
133            None
134        } else {
135            unsafe { CStr::from_ptr(ptr) }
136                .to_string_lossy()
137                .into_owned()
138                .into()
139        }
140    }
141
142    /// # Errors
143    ///
144    /// Returns `MtmdInputChunkError::ChunkOperationFailed` if copying fails.
145    pub fn copy(&self) -> Result<Self, MtmdInputChunkError> {
146        let chunk = unsafe { llama_cpp_bindings_sys::mtmd_input_chunk_copy(self.chunk.as_ptr()) };
147        let chunk = NonNull::new(chunk).ok_or(MtmdInputChunkError::ChunkOperationFailed)?;
148
149        Ok(Self { chunk, owned: true })
150    }
151
152    /// # Errors
153    ///
154    /// Returns [`MtmdEvalError::ImageChunkExceedsBatchSize`] when this is an
155    /// image chunk whose token count exceeds `n_batch`. Returns
156    /// [`MtmdEvalError::EvalFailure`] if the underlying encode or decode step
157    /// fails.
158    pub fn eval_single(
159        &self,
160        mtmd_ctx: &MtmdContext,
161        llama_ctx: &LlamaContext,
162        start_position: llama_cpp_bindings_sys::llama_pos,
163        seq_id: llama_cpp_bindings_sys::llama_seq_id,
164        n_batch: i32,
165        logits_last: bool,
166    ) -> Result<llama_cpp_bindings_sys::llama_pos, MtmdEvalError> {
167        let chunk_token_count = self.n_tokens();
168
169        if let Some(error) = image_chunk_batch_size_error(
170            matches!(self.chunk_type(), Ok(MtmdInputChunkType::Image)),
171            chunk_token_count,
172            n_batch,
173        ) {
174            return Err(error);
175        }
176
177        let mut final_position: llama_cpp_bindings_sys::llama_pos = start_position;
178        let mut out_vendored_return_code: i32 = 0;
179        let mut out_error: *mut c_char = std::ptr::null_mut();
180
181        let status = unsafe {
182            llama_cpp_bindings_sys::llama_rs_mtmd_eval_chunk_single(
183                mtmd_ctx.context.as_ptr(),
184                llama_ctx.context.as_ptr(),
185                self.chunk.as_ptr(),
186                start_position,
187                seq_id,
188                n_batch,
189                logits_last,
190                &raw mut final_position,
191                &raw mut out_vendored_return_code,
192                &raw mut out_error,
193            )
194        };
195
196        eval_chunk_single_status_to_result(
197            status,
198            final_position,
199            out_vendored_return_code,
200            out_error,
201        )
202    }
203}
204
205impl Drop for MtmdInputChunk {
206    fn drop(&mut self) {
207        if self.owned {
208            unsafe { llama_cpp_bindings_sys::mtmd_input_chunk_free(self.chunk.as_ptr()) }
209        }
210    }
211}
212
213#[cfg(test)]
214mod unit_tests {
215    use super::eval_chunk_single_status_to_result;
216    use super::image_chunk_batch_size_error;
217    use super::tokens_from_raw_ptr;
218    use crate::mtmd::image_chunk_batch_size_mismatch::ImageChunkBatchSizeMismatch;
219    use crate::mtmd::mtmd_eval_error::MtmdEvalError;
220
221    #[test]
222    fn tokens_from_raw_ptr_returns_none_for_null() {
223        assert!(unsafe { tokens_from_raw_ptr(std::ptr::null(), 5) }.is_none());
224    }
225
226    #[test]
227    fn tokens_from_raw_ptr_returns_none_for_zero_count() {
228        let token: llama_cpp_bindings_sys::llama_token = 42;
229        assert!(unsafe { tokens_from_raw_ptr(&raw const token, 0) }.is_none());
230    }
231
232    #[test]
233    fn tokens_from_raw_ptr_returns_some_for_valid() {
234        let tokens: [llama_cpp_bindings_sys::llama_token; 2] = [1, 2];
235        let result = unsafe { tokens_from_raw_ptr(tokens.as_ptr(), 2) };
236
237        assert!(result.is_some());
238        assert_eq!(result.unwrap().len(), 2);
239    }
240
241    #[test]
242    fn eval_chunk_single_status_ok_returns_final_position() {
243        let result = eval_chunk_single_status_to_result(
244            llama_cpp_bindings_sys::LLAMA_RS_MTMD_EVAL_CHUNK_SINGLE_OK,
245            7,
246            0,
247            std::ptr::null_mut(),
248        );
249
250        assert_eq!(result, Ok(7));
251    }
252
253    #[test]
254    fn eval_chunk_single_status_nonzero_code_maps_to_eval_failed() {
255        let result = eval_chunk_single_status_to_result(
256            llama_cpp_bindings_sys::LLAMA_RS_MTMD_EVAL_CHUNK_SINGLE_VENDORED_RETURNED_NONZERO_CODE,
257            0,
258            -3,
259            std::ptr::null_mut(),
260        );
261
262        assert_eq!(result, Err(MtmdEvalError::EvalFailed { code: -3 }));
263    }
264
265    #[test]
266    fn eval_chunk_single_status_allocation_failed_maps_to_not_enough_memory() {
267        let result = eval_chunk_single_status_to_result(
268            llama_cpp_bindings_sys::LLAMA_RS_MTMD_EVAL_CHUNK_SINGLE_ERROR_STRING_ALLOCATION_FAILED,
269            0,
270            0,
271            std::ptr::null_mut(),
272        );
273
274        assert_eq!(result, Err(MtmdEvalError::NotEnoughMemory));
275    }
276
277    #[test]
278    fn eval_chunk_single_status_cxx_exception_reports_unknown_error_for_null() {
279        let result = eval_chunk_single_status_to_result(
280            llama_cpp_bindings_sys::LLAMA_RS_MTMD_EVAL_CHUNK_SINGLE_VENDORED_THREW_CXX_EXCEPTION,
281            0,
282            0,
283            std::ptr::null_mut(),
284        );
285
286        assert_eq!(
287            result,
288            Err(MtmdEvalError::Reported {
289                message: "unknown error".to_string()
290            })
291        );
292    }
293
294    #[test]
295    #[should_panic(expected = "llama_rs_mtmd_eval_chunk_single returned unrecognized status")]
296    fn eval_chunk_single_status_unrecognized_panics() {
297        let _ = eval_chunk_single_status_to_result(
298            llama_cpp_bindings_sys::llama_rs_mtmd_eval_chunk_single_status::MAX,
299            0,
300            0,
301            std::ptr::null_mut(),
302        );
303    }
304
305    #[test]
306    fn image_chunk_over_batch_size_reports_mismatch() {
307        let error = image_chunk_batch_size_error(true, 9, 4);
308
309        assert_eq!(
310            error,
311            Some(MtmdEvalError::ImageChunkExceedsBatchSize(
312                ImageChunkBatchSizeMismatch {
313                    image_tokens: 9,
314                    n_batch: 4,
315                }
316            ))
317        );
318    }
319
320    #[test]
321    fn non_image_chunk_never_reports_mismatch() {
322        assert!(image_chunk_batch_size_error(false, 9, 4).is_none());
323    }
324
325    #[test]
326    fn image_chunk_within_batch_size_reports_no_mismatch() {
327        assert!(image_chunk_batch_size_error(true, 4, 4).is_none());
328    }
329}
llama_cpp_bindings/mtmd/mtmd_input_chunk.rs

llama_cpp_bindings/mtmd/
mtmd_input_chunk.rs
