mistralrs_core/pipeline/
mod.rs

1mod amoe;
2mod auto;
3pub mod chat_template;
4mod diffusion;
5mod ggml;
6mod gguf;
7mod inputs_processor;
8mod isq;
9pub(crate) mod llg;
10mod loaders;
11mod macros;
12mod normal;
13mod paths;
14mod processing;
15mod response;
16mod sampling;
17mod speculative;
18mod speech;
19mod vision;
20
21pub use super::diffusion_models::DiffusionGenerationParams;
22use crate::amoe::{AnyMoeConfig, AnyMoeExpertType, AnyMoeTrainingInputs, AnyMoeTrainingResult};
23use crate::device_map::DeviceMapper;
24use crate::paged_attention::{CacheConfig, CacheEngine, ModelConfigLike};
25use crate::prefix_cacher::PrefixCacheManagerV2;
26pub use amoe::{AnyMoeLoader, AnyMoePipeline};
27pub use auto::{AutoLoader, AutoLoaderBuilder};
28use chat_template::ChatTemplate;
29pub use diffusion::{DiffusionLoader, DiffusionLoaderBuilder};
30pub use ggml::{GGMLLoader, GGMLLoaderBuilder, GGMLSpecificConfig};
31pub use gguf::{GGUFLoader, GGUFLoaderBuilder, GGUFSpecificConfig};
32use image::DynamicImage;
33pub use inputs_processor::InputProcessorOutput;
34pub(crate) use isq::IsqModelLoader;
35pub use isq::{parse_isq_value, IsqModel, IsqOrganization, UQFF_MULTI_FILE_DELIMITER};
36use llguidance::toktrie::TokEnv;
37pub use loaders::{
38    AdapterKind, AutoDeviceMapParams, AutoNormalLoader, AutoVisionLoader, DeepSeekV2Loader,
39    DeepSeekV3Loader, DeviceMappedModelLoader, DiffusionLoaderType, DiffusionModel,
40    DiffusionModelLoader, FluxLoader, GLM4Loader, Gemma2Loader, Gemma3Loader, Gemma3nLoader,
41    GemmaLoader, Idefics2Loader, Idefics3Loader, LLaVALoader, LLaVANextLoader, LlamaLoader, Loader,
42    LocalModelPaths, MiniCpmOLoader, Mistral3Loader, MistralLoader, MixtralLoader, ModelKind,
43    ModelPaths, NormalLoaderType, NormalLoadingMetadata, NormalModel, NormalModelLoader,
44    Phi2Loader, Phi3Loader, Phi3VLoader, Phi3_5MoELoader, Phi4MMLoader, PrettyName,
45    QuantizationKind, Qwen2Loader, Qwen2VLLoader, Qwen2_5VLLoader, Qwen3Loader, Qwen3MoELoader,
46    SmolLm3Loader, Starcoder2Loader, TokenSource, VLlama4Loader, VLlamaLoader, VisionLoaderType,
47    VisionModel, VisionModelLoader,
48};
49use mistralrs_quant::IsqType;
50pub use normal::{NormalLoader, NormalLoaderBuilder, NormalSpecificConfig};
51pub(crate) use paths::{get_chat_template, get_model_paths, get_xlora_paths};
52pub use paths::{AdapterPaths, LoraAdapterPaths};
53pub(crate) use processing::{
54    apply_chat_template, BasicProcessor, MessagesAction, Processor, ProcessorCreator,
55};
56use rand_isaac::Isaac64Rng;
57pub use speculative::{SpeculativeConfig, SpeculativeLoader, SpeculativePipeline};
58pub use speech::{SpeechLoader, SpeechPipeline};
59use std::any::Any;
60use std::collections::HashMap;
61use std::fmt::Debug;
62use std::sync::Arc;
63use std::time::{Duration, Instant};
64use tokenizers::Tokenizer;
65pub use vision::{VisionLoader, VisionLoaderBuilder, VisionSpecificConfig};
66
67use anyhow::Result;
68use candle_core::{DType, Device, IndexOp, Tensor, Var};
69
70use crate::sequence::Sequence;
71
72pub use self::inputs_processor::{
73    text_models_inputs_processor, InputsProcessor, InputsProcessorType,
74};
75use self::text_models_inputs_processor::PagedAttentionMeta;
76pub use crate::kv_cache::{
77    Cache, CacheManager, EitherCache, KvCache, LayerCaches, NormalCache, NormalCacheType,
78};
79
80#[derive(Clone, PartialEq, Eq)]
81pub enum SupportedModality {
82    Text,
83    Audio,
84    Vision,
85}
86
87impl Debug for SupportedModality {
88    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
89        match self {
90            Self::Text => write!(f, "📝 Text"),
91            Self::Audio => write!(f, "🔊 Audio"),
92            Self::Vision => write!(f, "🖼️ Vision"),
93        }
94    }
95}
96
97#[derive(Debug, Clone)]
98pub struct Modalities {
99    pub input: Vec<SupportedModality>,
100    pub output: Vec<SupportedModality>,
101}
102
103pub struct GeneralMetadata {
104    pub max_seq_len: usize,
105    /// Only None if it doesn't make sense for the model
106    pub llg_factory: Option<Arc<llguidance::ParserFactory>>,
107    pub no_kv_cache: bool,
108    pub no_prefix_cache: bool,
109    pub num_hidden_layers: usize,
110    pub eos_tok: Vec<u32>,
111    pub kind: ModelKind,
112    // TODO: Replace is_xlora queries to check via kind instead:
113    pub is_xlora: bool,
114    pub activation_dtype: DType,
115    pub sliding_window: Option<usize>,
116    // PagedAttention stuff
117    pub cache_config: Option<CacheConfig>,
118    pub cache_engine: Option<CacheEngine>,
119    pub model_metadata: Option<Arc<dyn ModelConfigLike + Send + Sync>>,
120    pub modalities: Modalities,
121}
122
123impl GeneralMetadata {
124    pub fn tok_env(&self) -> Option<TokEnv> {
125        self.llg_factory.as_ref().map(|f| f.tok_env().clone())
126    }
127}
128
129pub enum CacheInstruction {
130    In,
131    Out,
132    /// load_preallocated_cache means to load the preallocated cache, if applicable.
133    Reset {
134        load_preallocated_cache: bool,
135        reset_non_granular: bool,
136    },
137    Nothing,
138}
139
140pub trait PreProcessingMixin: MetadataMixin {
141    fn get_processor(&self) -> Arc<dyn Processor> {
142        Arc::new(BasicProcessor)
143    }
144    /// Only None if it doesnt make sense for the model
145    fn get_chat_template(&self) -> Option<Arc<ChatTemplate>>;
146    fn get_input_processor_config(&self) -> Option<Arc<dyn Any>>;
147}
148
149pub trait IsqPipelineMixin {
150    fn re_isq_model(&mut self, dtype: IsqType) -> Result<()>;
151}
152
153pub trait CacheManagerMixin {
154    /// Clone the cache FROM the sequences' cache TO the model cache. Only called for completion seqs.
155    /// It is not a guarantee that this will be called for each completion step.
156    fn clone_in_cache(&self, seqs: &mut [&mut Sequence]);
157    /// Clone the cache FROM the model cache TO the sequences. Called for prompt and completion seqs.
158    /// It is not a guarantee that this will be called for each step.
159    fn clone_out_cache(&self, seqs: &mut [&mut Sequence]);
160    /// Set the model cache to all None. Only called for prompt seqs.
161    /// It is not a guarantee that this will be called for each prompt step.
162    /// This may also reset the non granular state if applicable.
163    fn set_none_cache(
164        &self,
165        seqs: &mut [&mut Sequence],
166        reset_non_granular: bool,
167        modify_draft_cache: bool,
168        load_preallocated_cache: bool,
169    );
170    fn cache(&self) -> &EitherCache;
171    fn do_preallocated_cache(&self) -> bool {
172        matches!(self.cache(), EitherCache::Normal(_))
173    }
174}
175
176pub trait MetadataMixin {
177    fn device(&self) -> Device;
178    /// Only None if it doesnt make sense for the model
179    fn tokenizer(&self) -> Option<Arc<Tokenizer>>;
180    fn name(&self) -> String;
181    fn reset_non_granular_state(&self);
182    fn get_metadata(&self) -> Arc<GeneralMetadata>;
183    fn device_mapper(&self) -> Option<&dyn DeviceMapper>;
184}
185
186/// Implemented by the base model of an AnyMoe.
187pub trait AnyMoePipelineMixin {
188    /// Get vars for each gating layer
189    fn amoe_layer_vars(&self) -> Vec<Vec<Var>> {
190        unreachable!()
191    }
192    fn amoe_finish_training(&mut self, _gate_model_id: Option<String>) -> candle_core::Result<()> {
193        unreachable!()
194    }
195    fn amoe_base_model_trainable_params(&self) -> usize {
196        unreachable!()
197    }
198    fn amoe_supported(&self) -> bool {
199        false
200    }
201    /// Per-layer cached outputs.
202    fn amoe_take_cached_gating_outputs(&mut self) -> Vec<Tensor> {
203        unreachable!()
204    }
205    /// Inject the MoE layers
206    #[allow(clippy::too_many_arguments)]
207    fn amoe_create_layers(
208        &mut self,
209        _model_ids: Vec<String>,
210        _token: &TokenSource,
211        _revision: Option<String>,
212        _match_regex: &str,
213        _config: AnyMoeConfig,
214        _dtype: DType,
215        _dev: &Device,
216        (_prefix, _mlp): (String, String),
217        _layers: Vec<usize>,
218        _expert_type: AnyMoeExpertType,
219        _silent: bool,
220        _gate_model_id: Option<String>,
221    ) -> candle_core::Result<()> {
222        unreachable!()
223    }
224    /// Pre-train the gating layers
225    #[allow(clippy::too_many_arguments)]
226    fn amoe_pre_train(
227        &self,
228        _inputs: AnyMoeTrainingInputs,
229        (_prefix, _mlp): (String, String),
230        _model_ids: Vec<String>,
231        _token: TokenSource,
232        _revision: Option<String>,
233        _layers: Vec<usize>,
234        _silent: bool,
235    ) -> Result<Option<AnyMoeTrainingResult>, candle_core::Error> {
236        unreachable!()
237    }
238}
239
240/// Category of the model. This can also be used to extract model-category specific tools,
241/// such as the vision model prompt prefixer.
242#[derive(Clone)]
243pub enum ModelCategory {
244    Text,
245    Vision {
246        prefixer: Arc<dyn MultimodalPromptPrefixer>,
247    },
248    Diffusion,
249    Audio,
250    Speech,
251}
252
253impl std::fmt::Debug for ModelCategory {
254    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
255        match self {
256            ModelCategory::Text => write!(f, "ModelCategory::Text"),
257            ModelCategory::Vision { .. } => write!(f, "ModelCategory::Vision {{ prefixer: .. }}"),
258            ModelCategory::Diffusion => write!(f, "ModelCategory::Diffusion"),
259            ModelCategory::Audio => write!(f, "ModelCategory::Audio"),
260            ModelCategory::Speech => write!(f, "ModelCategory::Speech"),
261        }
262    }
263}
264
265impl PartialEq for ModelCategory {
266    fn eq(&self, other: &Self) -> bool {
267        match (self, other) {
268            (Self::Text, Self::Text) => true,
269            (Self::Vision { .. }, Self::Vision { .. }) => true,
270            (Self::Audio, Self::Audio) => true,
271            (Self::Speech, Self::Speech) => true,
272            (Self::Diffusion, Self::Diffusion) => true,
273            (
274                Self::Text | Self::Vision { .. } | Self::Diffusion | Self::Audio | Self::Speech,
275                _,
276            ) => false,
277        }
278    }
279}
280
281/// Prepend a vision tag appropriate for the model to the prompt. Image indexing is assumed that start at 0.
282pub trait MultimodalPromptPrefixer: Send + Sync {
283    /// Prefix for inclusion in messages (may do nothing if the chat template handles it).
284    fn prefix_image(&self, _image_indices: Vec<usize>, prompt: &str) -> String {
285        prompt.to_string()
286    }
287    /// Prefix for inclusion in messages (may do nothing if the chat template handles it).
288    fn prefix_audio(&self, _audio_indexes: Vec<usize>, prompt: &str) -> String {
289        prompt.to_string()
290    }
291}
292
293pub enum CacheBackendMetadata {
294    DefaultInstructions {
295        pre_op: CacheInstruction,
296        post_op: CacheInstruction,
297    },
298    PagedAttention {
299        metadata: PagedAttentionMeta,
300        blocks_to_copy: HashMap<usize, Vec<usize>>,
301    },
302}
303
304#[derive(Clone, Debug)]
305pub enum ForwardInputsResult {
306    RawLogits {
307        logits: Tensor,
308    },
309    CausalGeneration {
310        logits: Tensor,
311    },
312    Image {
313        images: Vec<DynamicImage>,
314    },
315    Speech {
316        pcms: Vec<Arc<Vec<f32>>>,
317        rates: Vec<usize>,
318        channels: Vec<usize>,
319    },
320}
321
322impl ForwardInputsResult {
323    fn index_bs(&self, bs_idx: usize) -> candle_core::Result<Self> {
324        match self {
325            Self::CausalGeneration { logits } => Ok(Self::CausalGeneration {
326                logits: logits.i(bs_idx)?,
327            }),
328            Self::RawLogits { logits } => Ok(Self::RawLogits {
329                logits: logits.i(bs_idx)?,
330            }),
331            Self::Image { images } => Ok(Self::Image {
332                images: vec![images[bs_idx].clone()],
333            }),
334            Self::Speech {
335                pcms,
336                rates,
337                channels,
338            } => Ok(Self::Speech {
339                pcms: vec![pcms[bs_idx].clone()],
340                rates: vec![rates[bs_idx]],
341                channels: vec![channels[bs_idx]],
342            }),
343        }
344    }
345
346    fn to_device(&self, device: &Device) -> candle_core::Result<Self> {
347        match self {
348            Self::CausalGeneration { logits } => Ok(Self::CausalGeneration {
349                logits: logits.to_device(device)?,
350            }),
351            Self::RawLogits { logits } => Ok(Self::RawLogits {
352                logits: logits.to_device(device)?,
353            }),
354            Self::Image { .. } => Ok(self.clone()),
355            Self::Speech { .. } => Ok(self.clone()),
356        }
357    }
358}
359
360#[derive(serde::Serialize, serde::Deserialize)]
361pub(crate) struct FileListCache {
362    files: Vec<String>,
363}
364
365#[async_trait::async_trait]
366pub trait Pipeline:
367    Send
368    + Sync
369    + PreProcessingMixin
370    + IsqPipelineMixin
371    + CacheManagerMixin
372    + MetadataMixin
373    + AnyMoePipelineMixin
374{
375    fn forward_inputs(
376        &mut self,
377        inputs: Box<dyn Any>,
378        return_raw_logits: bool,
379    ) -> Result<ForwardInputsResult, candle_core::Error>;
380
381    /// Returns the total of model execution time.
382    #[allow(clippy::too_many_arguments)]
383    async fn step(
384        &mut self,
385        input_seqs: &mut [&mut Sequence],
386        is_prompt: bool,
387        return_raw_logits: bool,
388        prefix_cacher: &mut PrefixCacheManagerV2,
389        disable_eos_stop: bool,
390        rng: Arc<std::sync::Mutex<Isaac64Rng>>,
391        backend_metadata: CacheBackendMetadata,
392    ) -> Result<Duration, candle_core::Error> {
393        match backend_metadata {
394            CacheBackendMetadata::DefaultInstructions { pre_op, post_op } => {
395                let inputs_iter =
396                    std::iter::once(self.get_processor().inputs_processor().process_inputs(
397                        self.tokenizer(),
398                        input_seqs,
399                        is_prompt,
400                        self.get_metadata().is_xlora,
401                        &self.device(),
402                        self.get_metadata().no_kv_cache,
403                        None,
404                        return_raw_logits,
405                        self.get_input_processor_config(),
406                        None,
407                        self.device_mapper(),
408                    ));
409
410                let mut logits = vec![None; input_seqs.len()];
411                let len_inputs = 1;
412                let mut raw_out_logits = vec![vec![None; len_inputs]; input_seqs.len()];
413
414                let mut exec_duration = Duration::ZERO;
415                for (i, inputs) in inputs_iter.into_iter().enumerate() {
416                    let InputProcessorOutput {
417                        inputs,
418                        seq_indices,
419                    } = inputs.map_err(candle_core::Error::msg)?;
420                    if i == 0 {
421                        match pre_op {
422                            CacheInstruction::In => self.clone_in_cache(input_seqs),
423                            CacheInstruction::Nothing => (),
424                            CacheInstruction::Reset {
425                                load_preallocated_cache,
426                                reset_non_granular,
427                            } => self.set_none_cache(
428                                input_seqs,
429                                reset_non_granular,
430                                false,
431                                load_preallocated_cache,
432                            ),
433                            _ => unreachable!("Unreachable PRE cache op."),
434                        }
435                    }
436
437                    let start = Instant::now();
438                    let raw_logits = self.forward_inputs(inputs, return_raw_logits)?;
439                    let end = Instant::now();
440                    exec_duration += end.duration_since(start);
441
442                    for (logit_idx, seq_idx) in seq_indices.into_iter().enumerate() {
443                        if let ForwardInputsResult::RawLogits { logits } = &raw_logits {
444                            raw_out_logits[seq_idx][i] =
445                                Some(logits.i(logit_idx)?.to_device(&Device::Cpu)?);
446                        } else {
447                            logits[seq_idx] = Some(raw_logits.index_bs(logit_idx)?);
448                        }
449                    }
450                }
451
452                match post_op {
453                    CacheInstruction::Out => self.clone_out_cache(input_seqs),
454                    CacheInstruction::Nothing => (),
455                    CacheInstruction::Reset {
456                        load_preallocated_cache,
457                        reset_non_granular,
458                    } => self.set_none_cache(
459                        input_seqs,
460                        reset_non_granular,
461                        false,
462                        load_preallocated_cache,
463                    ),
464                    _ => unreachable!("Unreachable POST cache op."),
465                }
466
467                if raw_out_logits[0][0].is_some() {
468                    let start = Instant::now();
469                    response::send_raw_responses(
470                        input_seqs,
471                        raw_out_logits
472                            .into_iter()
473                            .map(|raw| raw.into_iter().flatten().collect::<Vec<_>>())
474                            .collect(),
475                    )
476                    .await?;
477                    let end = Instant::now();
478                    exec_duration += end.duration_since(start);
479
480                    return Ok(exec_duration);
481                }
482
483                let start = Instant::now();
484                let logits_on_cpu = logits.len() > 1;
485                let logits = logits
486                    .into_iter()
487                    .map(|l| {
488                        let l = l.expect("Did not get any inputs. This is shocking.");
489                        if logits_on_cpu {
490                            l.to_device(&Device::Cpu)
491                        } else {
492                            Ok(l)
493                        }
494                    })
495                    .collect::<candle_core::Result<Vec<_>>>()?;
496
497                match &logits[0] {
498                    ForwardInputsResult::RawLogits { .. } => unreachable!(),
499                    ForwardInputsResult::CausalGeneration { .. } => {
500                        self.sample_causal_gen(
501                            input_seqs,
502                            logits
503                                .into_iter()
504                                .map(|r| {
505                                    #[allow(irrefutable_let_patterns)]
506                                    let ForwardInputsResult::CausalGeneration { logits } = r
507                                    else {
508                                        unreachable!(
509                                            "All results must have same type, `CausalGeneration`"
510                                        )
511                                    };
512                                    logits
513                                })
514                                .collect::<Vec<_>>(),
515                            prefix_cacher,
516                            disable_eos_stop,
517                            rng,
518                        )
519                        .await?;
520                    }
521                    ForwardInputsResult::Image { .. } => {
522                        response::send_image_responses(
523                            input_seqs,
524                            logits
525                                .into_iter()
526                                .map(|r| {
527                                    #[allow(irrefutable_let_patterns)]
528                                    let ForwardInputsResult::Image { images } = r
529                                    else {
530                                        unreachable!("All results must have same type, `Image`")
531                                    };
532                                    images
533                                        .into_iter()
534                                        .next()
535                                        .expect("Must have at least 1 element.")
536                                })
537                                .collect::<Vec<_>>(),
538                        )
539                        .await?;
540                    }
541                    ForwardInputsResult::Speech { .. } => {
542                        let rates = logits
543                            .iter()
544                            .map(|r| {
545                                #[allow(irrefutable_let_patterns)]
546                                let ForwardInputsResult::Speech { rates, .. } = r
547                                else {
548                                    unreachable!("All results must have same type, `Speech`")
549                                };
550                                assert_eq!(rates.len(), 1, "Each sequence must have 1 PCM output.");
551                                *rates.first().unwrap()
552                            })
553                            .collect::<Vec<_>>();
554                        let channels = logits
555                            .iter()
556                            .map(|r| {
557                                #[allow(irrefutable_let_patterns)]
558                                let ForwardInputsResult::Speech { channels, .. } = r
559                                else {
560                                    unreachable!("All results must have same type, `Speech`")
561                                };
562                                assert_eq!(
563                                    channels.len(),
564                                    1,
565                                    "Each sequence must have 1 PCM output."
566                                );
567                                *channels.first().unwrap()
568                            })
569                            .collect::<Vec<_>>();
570                        let pcms = logits
571                            .into_iter()
572                            .map(|r| {
573                                #[allow(irrefutable_let_patterns)]
574                                let ForwardInputsResult::Speech { pcms, .. } = r
575                                else {
576                                    unreachable!("All results must have same type, `Speech`")
577                                };
578                                assert_eq!(pcms.len(), 1, "Each sequence must have 1 PCM output.");
579                                pcms.into_iter().nth(0).unwrap()
580                            })
581                            .collect::<Vec<_>>();
582                        response::send_speech_responses(input_seqs, &pcms, &rates, &channels)
583                            .await?;
584                    }
585                }
586                let end = Instant::now();
587                exec_duration += end.duration_since(start);
588
589                Ok(exec_duration)
590            }
591            CacheBackendMetadata::PagedAttention {
592                metadata,
593                blocks_to_copy,
594            } => {
595                // Cloning might be bad?
596                self.get_metadata()
597                    .cache_engine
598                    .as_ref()
599                    .expect("PagedAttention must have cache engines.")
600                    .execute_scheduler_ops(&blocks_to_copy)?;
601
602                let inputs_iter =
603                    std::iter::once(self.get_processor().inputs_processor().process_inputs(
604                        self.tokenizer(),
605                        input_seqs,
606                        is_prompt,
607                        self.get_metadata().is_xlora,
608                        &self.device(),
609                        self.get_metadata().no_kv_cache,
610                        None,
611                        return_raw_logits,
612                        self.get_input_processor_config(),
613                        Some(metadata),
614                        self.device_mapper(),
615                    ));
616
617                let mut logits = vec![None; input_seqs.len()];
618                let len_inputs = 1;
619                let mut raw_out_logits = vec![vec![None; len_inputs]; input_seqs.len()];
620
621                let mut exec_duration = Duration::ZERO;
622                for (i, inputs) in inputs_iter.into_iter().enumerate() {
623                    let InputProcessorOutput {
624                        inputs,
625                        seq_indices,
626                    } = inputs.map_err(candle_core::Error::msg)?;
627
628                    let start = Instant::now();
629                    let raw_logits = self.forward_inputs(inputs, return_raw_logits)?;
630                    let end = Instant::now();
631                    exec_duration += end.duration_since(start);
632
633                    for (logit_idx, seq_idx) in seq_indices.into_iter().enumerate() {
634                        if let ForwardInputsResult::RawLogits { logits } = &raw_logits {
635                            raw_out_logits[seq_idx][i] =
636                                Some(logits.i(logit_idx)?.to_device(&Device::Cpu)?);
637                        } else {
638                            logits[seq_idx] = Some(raw_logits.index_bs(logit_idx)?);
639                        }
640                    }
641                }
642
643                if raw_out_logits[0][0].is_some() {
644                    let start = Instant::now();
645                    response::send_raw_responses(
646                        input_seqs,
647                        raw_out_logits
648                            .into_iter()
649                            .map(|raw| raw.into_iter().flatten().collect::<Vec<_>>())
650                            .collect(),
651                    )
652                    .await?;
653                    let end = Instant::now();
654                    exec_duration += end.duration_since(start);
655
656                    return Ok(exec_duration);
657                }
658
659                let start = Instant::now();
660                let logits_on_cpu = logits.len() > 1;
661                let logits = logits
662                    .into_iter()
663                    .map(|l| {
664                        let l = l.expect("Did not get any inputs. This is shocking.");
665                        if logits_on_cpu {
666                            l.to_device(&Device::Cpu)
667                        } else {
668                            Ok(l)
669                        }
670                    })
671                    .collect::<candle_core::Result<Vec<_>>>()?;
672
673                match &logits[0] {
674                    ForwardInputsResult::RawLogits { .. } => unreachable!(),
675                    ForwardInputsResult::CausalGeneration { .. } => {
676                        self.sample_causal_gen(
677                            input_seqs,
678                            logits
679                                .into_iter()
680                                .map(|r| {
681                                    #[allow(irrefutable_let_patterns)]
682                                    let ForwardInputsResult::CausalGeneration { logits } = r
683                                    else {
684                                        unreachable!("All results must have same type")
685                                    };
686                                    logits
687                                })
688                                .collect::<Vec<_>>(),
689                            prefix_cacher,
690                            disable_eos_stop,
691                            rng,
692                        )
693                        .await?;
694                    }
695                    ForwardInputsResult::Image { .. } => {
696                        response::send_image_responses(
697                            input_seqs,
698                            logits
699                                .into_iter()
700                                .map(|r| {
701                                    #[allow(irrefutable_let_patterns)]
702                                    let ForwardInputsResult::Image { images } = r
703                                    else {
704                                        unreachable!("All results must have same type, `Image`")
705                                    };
706                                    images
707                                        .into_iter()
708                                        .next()
709                                        .expect("Must have at least 1 element.")
710                                })
711                                .collect::<Vec<_>>(),
712                        )
713                        .await?;
714                    }
715                    ForwardInputsResult::Speech { .. } => {
716                        let rates = logits
717                            .iter()
718                            .map(|r| {
719                                #[allow(irrefutable_let_patterns)]
720                                let ForwardInputsResult::Speech { rates, .. } = r
721                                else {
722                                    unreachable!("All results must have same type, `Speech`")
723                                };
724                                assert_eq!(rates.len(), 1, "Each sequence must have 1 PCM output.");
725                                *rates.first().unwrap()
726                            })
727                            .collect::<Vec<_>>();
728                        let channels = logits
729                            .iter()
730                            .map(|r| {
731                                #[allow(irrefutable_let_patterns)]
732                                let ForwardInputsResult::Speech { channels, .. } = r
733                                else {
734                                    unreachable!("All results must have same type, `Speech`")
735                                };
736                                assert_eq!(
737                                    channels.len(),
738                                    1,
739                                    "Each sequence must have 1 PCM output."
740                                );
741                                *channels.first().unwrap()
742                            })
743                            .collect::<Vec<_>>();
744                        let pcms = logits
745                            .into_iter()
746                            .map(|r| {
747                                #[allow(irrefutable_let_patterns)]
748                                let ForwardInputsResult::Speech { pcms, .. } = r
749                                else {
750                                    unreachable!("All results must have same type, `Speech`")
751                                };
752                                assert_eq!(pcms.len(), 1, "Each sequence must have 1 PCM output.");
753                                pcms.into_iter().nth(0).unwrap()
754                            })
755                            .collect::<Vec<_>>();
756                        response::send_speech_responses(input_seqs, &pcms, &rates, &channels)
757                            .await?;
758                    }
759                }
760                let end = Instant::now();
761                exec_duration += end.duration_since(start);
762
763                Ok(exec_duration)
764            }
765        }
766    }
767
768    async fn sample_causal_gen(
769        &self,
770        seqs: &mut [&mut Sequence],
771        logits: Vec<Tensor>,
772        prefix_cacher: &mut PrefixCacheManagerV2,
773        disable_eos_stop: bool,
774        rng: Arc<std::sync::Mutex<Isaac64Rng>>,
775    ) -> Result<(), candle_core::Error>;
776
777    fn category(&self) -> ModelCategory;
778}
779
780pub(crate) fn extract_logits(
781    logits: &Tensor,
782    context_lens: Vec<(usize, usize)>,
783) -> candle_core::Result<Tensor> {
784    let mut toks = Vec::new();
785    for (dim, (start, len)) in logits.chunk(logits.dims()[0], 0)?.iter().zip(context_lens) {
786        toks.push(dim.narrow(1, start, len)?);
787    }
788    Tensor::cat(&toks, 0)
789}
790
791#[cfg(test)]
792mod tests {
793    use crate::MessageContent;
794    use either::Either;
795    use indexmap::IndexMap;
796    use serde_json::Value;
797
798    macro_rules! hashmap {
799        (@single $($x:tt)*) => (());
800        (@count $($rest:expr),*) => (<[()]>::len(&[$(hashmap!(@single $rest)),*]));
801
802        ($($key:expr => $value:expr,)+) => { hashmap!($($key => $value),+) };
803        ($($key:expr => $value:expr),*) => {
804            {
805                let _cap = hashmap!(@count $($key),*);
806                let mut _map = ::indexmap::IndexMap::with_capacity(_cap);
807                $(
808                    let _ = _map.insert($key, Value::String($value));
809                )*
810                _map
811            }
812        };
813    }
814
815    #[cfg(test)]
816    #[track_caller]
817    fn test_with_inputs(
818        templates: &[(bool, &str, &str, &str, &str)],
819        expected_outputs: &[&str],
820        inputs: Vec<IndexMap<String, MessageContent>>,
821    ) {
822        use crate::pipeline::chat_template::ChatTemplateValue;
823
824        use super::chat_template::apply_chat_template_to;
825        let mut failed = Vec::new();
826        let n_templates = templates.len();
827        for ((has_system, bos, eos, unk, template), expected) in
828            templates.iter().zip(expected_outputs)
829        {
830            let output = match apply_chat_template_to(
831                if !has_system {
832                    inputs[1..].to_vec()
833                } else {
834                    inputs.clone()
835                },
836                true,
837                None,
838                &ChatTemplateValue(Either::Left(template.to_string())),
839                Some(bos.to_string()),
840                Some(eos.to_string()),
841                Some(unk.to_string()),
842                Vec::new(),
843            ) {
844                Ok(v) => v,
845                Err(e) => {
846                    failed.push(format!("Failed with {e}."));
847                    continue;
848                }
849            };
850            if output != *expected {
851                failed.push(format!(
852                    "Expected: `{}` \n\nGot:      `{}`",
853                    expected.replace('\n', "\\n"),
854                    output.replace('\n', "\\n")
855                ));
856            }
857        }
858        if !failed.is_empty() {
859            for (i, line) in failed.iter().enumerate() {
860                println!("------------ Template {i} ------------");
861                println!("{line}");
862            }
863            println!("------------------------");
864            panic!("{}/{n_templates} chat templates failed.", failed.len());
865        }
866    }
867
868    #[test]
869    /// Generating these cases:
870    /// ```py
871    /// >>> t=transformers.AutoTokenizer.from_pretrained(...)
872    /// # If non-system prompt model
873    /// >>> t.apply_chat_template([{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there"},{"role":"user","content":"Who are you"},{"role":"assistant","content":"   I am an assistant   "},{"role":"user","content":"Another question"}], add_generation_prompt=True, tokenize=False)
874    /// # If system prompt model
875    /// >>> t.apply_chat_template([{"role":"system","content":"You are a helpful assistant"},{"role":"user","content":"Hello"},{"role":"assistant","content":"Hi there"},{"role":"user","content":"Who are you"},{"role":"assistant","content":"   I am an assistant   "},{"role":"user","content":"Another question"}], add_generation_prompt=True, tokenize=False)
876    /// ```
877    fn test_chat_templates() {
878        let templates = [
879            // ChatML: https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B
880            (true, "<s>", "</s>", "<unk>", "{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"),
881            // mistralai/Mistral-7B-Instruct-v0.1
882            (false, "<s>", "</s>", "<unk>", "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token + ' ' }}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}"),
883            // meta-llama/Llama-2-13b-chat-hf
884            (true, "<s>", "</s>", "<unk>", "{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' '  + content.strip() + ' ' + eos_token }}{% endif %}{% endfor %}"),
885            // mistralai/Mixtral-8x7B-Instruct-v0.1
886            (false, "<s>", "</s>", "<unk>", "{{ bos_token }}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if message['role'] == 'user' %}{{ '[INST] ' + message['content'] + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ message['content'] + eos_token}}{% else %}{{ raise_exception('Only user and assistant roles are supported!') }}{% endif %}{% endfor %}"),
887            // google/gemma-7b-it
888            (false, "<bos>", "<eos>", "<unk>", "{{ bos_token }}{% if messages[0]['role'] == 'system' %}{{ raise_exception('System role not supported') }}{% endif %}{% for message in messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if (message['role'] == 'assistant') %}{% set role = 'model' %}{% else %}{% set role = message['role'] %}{% endif %}{{ '<start_of_turn>' + role + '\n' + message['content'] | trim + '<end_of_turn>\n' }}{% endfor %}{% if add_generation_prompt %}{{'<start_of_turn>model\n'}}{% endif %}"),
889            // HuggingFaceM4/idefics2-8b-chatty
890            (true, "<s>", "</s>", "<unk>", "{% for message in messages %}{{message['role'].capitalize()}}{% if message['content'][0]['type'] == 'image' %}{{':'}}{% else %}{{': '}}{% endif %}{% for line in message['content'] %}{% if line['type'] == 'text' %}{{line['text']}}{% elif line['type'] == 'image' %}{{ '<image>' }}{% endif %}{% endfor %}<end_of_utterance>\n{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"),
891        ];
892        let expected_outputs = [
893            // ChatML: https://huggingface.co/teknium/OpenHermes-2.5-Mistral-7B
894            "<|im_start|>system\nYou are a helpful assistant<|im_end|>\n<|im_start|>user\nHello<|im_end|>\n<|im_start|>assistant\nHi there<|im_end|>\n<|im_start|>user\nWho are you<|im_end|>\n<|im_start|>assistant\n   I am an assistant   <|im_end|>\n<|im_start|>user\nAnother question<|im_end|>\n<|im_start|>assistant\n",
895            // mistralai/Mistral-7B-Instruct-v0.1
896            "<s>[INST] Hello [/INST]Hi there</s> [INST] Who are you [/INST]   I am an assistant   </s> [INST] Another question [/INST]",
897            // meta-llama/Llama-2-13b-chat-hf
898            "<s>[INST] <<SYS>>\nYou are a helpful assistant\n<</SYS>>\n\nHello [/INST] Hi there </s><s>[INST] Who are you [/INST] I am an assistant </s><s>[INST] Another question [/INST]",
899            // mistralai/Mixtral-8x7B-Instruct-v0.1
900            "<s>[INST] Hello [/INST]Hi there</s>[INST] Who are you [/INST]   I am an assistant   </s>[INST] Another question [/INST]",
901            // google/gemma-7b-it
902            "<bos><start_of_turn>user\nHello<end_of_turn>\n<start_of_turn>model\nHi there<end_of_turn>\n<start_of_turn>user\nWho are you<end_of_turn>\n<start_of_turn>model\nI am an assistant<end_of_turn>\n<start_of_turn>user\nAnother question<end_of_turn>\n<start_of_turn>model\n",
903        ];
904        let messages = [
905            ["system", "You are a helpful assistant"],
906            ["user", "Hello"],
907            ["assistant", "Hi there"],
908            ["user", "Who are you"],
909            ["assistant", "   I am an assistant   "],
910            ["user", "Another question"],
911        ];
912        let mut inputs = Vec::new();
913        for [role, content] in messages {
914            let mut message: IndexMap<String, Either<String, Vec<IndexMap<String, Value>>>> =
915                IndexMap::new();
916            message.insert("role".to_string(), Either::Left(role.to_string()));
917            message.insert("content".to_string(), Either::Left(content.to_string()));
918            inputs.push(message);
919        }
920        test_with_inputs(&templates, &expected_outputs, inputs);
921    }
922
923    #[test]
924    /// Generating these cases:
925    /// ```py
926    /// >>> processor=transformers.AutoProcessor.from_pretrained(...)
927    /// >>> processor.apply_chat_template([
928    ///         {"role":"system","content":[{"type":"text", "text": "You are a helpful assistant"}]},
929    ///         {"role":"user","content":[{"type":"image"}, {"type":"text", "text": "Hello, please describe the above."}]},
930    ///         {"role":"assistant","content":[{"type":"text", "text": "Hi there"}]},
931    ///         {"role":"user","content":[{"type":"text", "text": "Who are you"}]},
932    ///         {"role":"assistant","content":[{"type":"text", "text": "   I am an assistant   "}]},
933    ///         {"role":"user","content":[{"type":"text", "text": "Another question"}]}
934    ///     ], add_generation_prompt=True, tokenize=False)
935    /// ```
936    fn test_image_chat_templates() {
937        let templates = [
938            // HuggingFaceM4/idefics2-8b-chatty
939            (true, "<s>", "</s>", "<unk>", "{% for message in messages %}{{message['role'].capitalize()}}{% if message['content'][0]['type'] == 'image' %}{{':'}}{% else %}{{': '}}{% endif %}{% for line in message['content'] %}{% if line['type'] == 'text' %}{{line['text']}}{% elif line['type'] == 'image' %}{{ '<image>' }}{% endif %}{% endfor %}<end_of_utterance>\n{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"),
940        ];
941        let expected_outputs = [
942            // HuggingFaceM4/idefics2-8b-chatty
943            "System: You are a helpful assistant<end_of_utterance>\nUser:<image>Hello, please describe the above.<end_of_utterance>\nAssistant: Hi there<end_of_utterance>\nUser:<image>This is me, who are you<end_of_utterance>\nAssistant:    I am an assistant   <end_of_utterance>\nUser:<image>Another question, what is this?<end_of_utterance>\nAssistant:",
944        ];
945
946        let mut inputs = Vec::new();
947
948        let mut message: IndexMap<String, Either<String, Vec<IndexMap<String, Value>>>> =
949            IndexMap::new();
950        message.insert("role".to_string(), Either::Left("system".to_string()));
951        message.insert(
952            "content".to_string(),
953            Either::Right(vec![hashmap! {
954                "type".to_string() => "text".to_string(),
955                "text".to_string() => "You are a helpful assistant".to_string()
956            }]),
957        );
958        inputs.push(message);
959
960        let mut message: IndexMap<String, Either<String, Vec<IndexMap<String, Value>>>> =
961            IndexMap::new();
962        message.insert("role".to_string(), Either::Left("user".to_string()));
963        message.insert(
964            "content".to_string(),
965            Either::Right(vec![
966                hashmap! {
967                    "type".to_string() => "image".to_string()
968                },
969                hashmap! {
970                    "type".to_string() => "text".to_string(),
971                    "text".to_string() => "Hello, please describe the above.".to_string()
972                },
973            ]),
974        );
975        inputs.push(message);
976
977        let mut message: IndexMap<String, Either<String, Vec<IndexMap<String, Value>>>> =
978            IndexMap::new();
979        message.insert("role".to_string(), Either::Left("assistant".to_string()));
980        message.insert(
981            "content".to_string(),
982            Either::Right(vec![hashmap! {
983                "type".to_string() => "text".to_string(),
984                "text".to_string() => "Hi there".to_string()
985            }]),
986        );
987        inputs.push(message);
988
989        let mut message: IndexMap<String, Either<String, Vec<IndexMap<String, Value>>>> =
990            IndexMap::new();
991        message.insert("role".to_string(), Either::Left("user".to_string()));
992        message.insert(
993            "content".to_string(),
994            Either::Right(vec![
995                hashmap! {
996                    "type".to_string() => "image".to_string()
997                },
998                hashmap! {
999                    "type".to_string() => "text".to_string(),
1000                    "text".to_string() => "This is me, who are you".to_string()
1001                },
1002            ]),
1003        );
1004        inputs.push(message);
1005
1006        let mut message: IndexMap<String, Either<String, Vec<IndexMap<String, Value>>>> =
1007            IndexMap::new();
1008        message.insert("role".to_string(), Either::Left("assistant".to_string()));
1009        message.insert(
1010            "content".to_string(),
1011            Either::Right(vec![hashmap! {
1012                "type".to_string() => "text".to_string(),
1013                "text".to_string() => "   I am an assistant   ".to_string()
1014            }]),
1015        );
1016        inputs.push(message);
1017
1018        let mut message: IndexMap<String, Either<String, Vec<IndexMap<String, Value>>>> =
1019            IndexMap::new();
1020        message.insert("role".to_string(), Either::Left("user".to_string()));
1021        message.insert(
1022            "content".to_string(),
1023            Either::Right(vec![
1024                hashmap! {
1025                    "type".to_string() => "image".to_string()
1026                },
1027                hashmap! {
1028                    "type".to_string() => "text".to_string(),
1029                    "text".to_string() => "Another question, what is this?".to_string()
1030                },
1031            ]),
1032        );
1033        inputs.push(message);
1034
1035        test_with_inputs(&templates, &expected_outputs, inputs);
1036    }
1037}
mistralrs_core/pipeline/mod.rs

mistralrs_core/pipeline/
mod.rs
