Skip to main content

llama_cpp_bindings/
ingest_prompt_chunk.rs

1use crate::mtmd::MtmdInputChunk;
2use crate::mtmd::MtmdInputChunkType;
3use crate::mtmd::MtmdInputChunkTypeError;
4use crate::sampled_token_classifier::SampledTokenClassifier;
5
6/// # Errors
7/// Returns [`MtmdInputChunkTypeError`] when the chunk reports a type unknown
8/// to this binding. Counters are not updated on error.
9pub fn ingest_prompt_chunk(
10    classifier: &mut SampledTokenClassifier<'_>,
11    chunk: &MtmdInputChunk,
12) -> Result<(), MtmdInputChunkTypeError> {
13    let n_tokens = chunk.n_tokens() as u64;
14    match chunk.chunk_type()? {
15        MtmdInputChunkType::Text => {
16            classifier.record_prompt_tokens(n_tokens);
17            if let Some(tokens) = chunk.text_tokens() {
18                classifier.ingest_prompt_tokens(tokens);
19            }
20        }
21        MtmdInputChunkType::Image => classifier.record_input_image_tokens(n_tokens),
22        MtmdInputChunkType::Audio => classifier.record_input_audio_tokens(n_tokens),
23    }
24
25    Ok(())
26}