aha/models/common/
mod.rs

1use anyhow::Result;
2use candle_core::Tensor;
3pub mod embedding;
4pub mod generate;
5pub mod gguf;
6pub mod model_mapping;
7pub mod modules;
8pub mod reranker;
9pub mod sample;
10
11/// 多模态模型的特征数据
12/// 每个模型数据不一样
13/// 需按顺序存放与取用
14#[derive(Clone, Debug)]
15pub struct MultiModalData {
16    pub data_vec: Vec<Option<Tensor>>,
17}
18impl MultiModalData {
19    pub fn new(data_vec: Vec<Option<Tensor>>) -> Self {
20        Self { data_vec }
21    }
22}
23
24#[allow(unused)]
25pub trait InferenceModel {
26    /// 初始前向传播（考虑多模态输入）
27    /// 默认实现无特殊数据
28    fn forward_initial(
29        &mut self,
30        input_ids: &Tensor,
31        seqlen_offset: usize,
32        data: MultiModalData,
33    ) -> Result<Tensor> {
34        Self::forward_step(self, input_ids, seqlen_offset)
35    }
36
37    /// 后续前向传播（自回归步骤）
38    fn forward_step(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result<Tensor>;
39
40    /// 清理 KV cache
41    fn clear_cache(&mut self);
42
43    /// 获取结束 token IDs
44    fn stop_token_ids(&self) -> Vec<u32>;
45}
aha/models/common/mod.rs

aha/models/common/
mod.rs
