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

llama_cpp_bindings/mtmd/
mtmd_input_chunk.rs
