jhqxxx
738 lines
29 KiB
Rust
738 lines
29 KiB
Rust
use anyhow::{Result, anyhow};
|
|
use candle_core::{D, IndexOp, Shape, Tensor};
|
|
use candle_nn::{
|
|
Conv2d, Embedding, LayerNorm, Linear, Module, RmsNorm, VarBuilder, embedding, linear,
|
|
linear_no_bias, rms_norm,
|
|
};
|
|
use num::integer::Roots;
|
|
|
|
use crate::{
|
|
models::{
|
|
common::{
|
|
NaiveAttnGateUpDownMLPBlock, NaiveAttnTwoLinearMLPBlock, get_conv2d, get_layer_norm,
|
|
},
|
|
paddleocr_vl::config::{
|
|
PaddleOCRVLConfig, PaddleOCRVLRopeScalingConfig, PaddleOCRVLVisionConfig,
|
|
},
|
|
},
|
|
position_embed::rope::{Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding},
|
|
utils::tensor_utils::{
|
|
get_vision_next_indices, interpolate_bilinear, masked_scatter_dim0, nonzero_index,
|
|
prepare_causal_attention_mask, zero_index,
|
|
},
|
|
};
|
|
|
|
pub struct Projector {
|
|
merge_size: usize,
|
|
pre_norm: LayerNorm,
|
|
linear_1: Linear,
|
|
linear_2: Linear,
|
|
}
|
|
|
|
impl Projector {
|
|
pub fn new(vb: VarBuilder, config: &PaddleOCRVLConfig) -> Result<Self> {
|
|
let merge_size = config.vision_config.spatial_merge_size;
|
|
let hidden_size = config.vision_config.hidden_size * merge_size * merge_size;
|
|
let pre_norm = get_layer_norm(
|
|
vb.pp("pre_norm"),
|
|
config.rms_norm_eps,
|
|
config.vision_config.hidden_size,
|
|
)?;
|
|
let linear_1 = linear(hidden_size, hidden_size, vb.pp("linear_1"))?;
|
|
let linear_2 = linear(hidden_size, config.hidden_size, vb.pp("linear_2"))?;
|
|
|
|
Ok(Self {
|
|
merge_size,
|
|
pre_norm,
|
|
linear_1,
|
|
linear_2,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor, image_grid_thw: &Tensor) -> Result<Tensor> {
|
|
let img_num = image_grid_thw.dim(0)?;
|
|
let mut processed_features = vec![];
|
|
let start = 0usize;
|
|
for i in 0..img_num {
|
|
let [t, h, w] = image_grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
|
|
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
|
|
};
|
|
let end = start + (t * h * w) as usize;
|
|
let xs_i = xs.i((start..end, ..))?;
|
|
let xs_i = self.pre_norm.forward(&xs_i)?;
|
|
let dim = xs_i.dim(1)?;
|
|
let shape = Shape::from(vec![
|
|
t as usize,
|
|
h as usize / self.merge_size,
|
|
self.merge_size,
|
|
w as usize / self.merge_size,
|
|
self.merge_size,
|
|
dim,
|
|
]);
|
|
let xs_i = xs_i
|
|
.reshape((t as usize, h as usize, w as usize, dim))?
|
|
.reshape(shape)?
|
|
.permute((0, 1, 3, 2, 4, 5))?
|
|
.reshape((
|
|
(t * h * w) as usize / self.merge_size / self.merge_size,
|
|
self.merge_size * self.merge_size * dim,
|
|
))?;
|
|
let xs_i = self.linear_1.forward(&xs_i)?.gelu()?;
|
|
let xs_i = self.linear_2.forward(&xs_i)?;
|
|
processed_features.push(xs_i);
|
|
}
|
|
let xs = Tensor::cat(&processed_features, 0)?;
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
pub struct SiglipVisionEmbeddings {
|
|
embed_dim: usize,
|
|
patch_size: usize,
|
|
patch_embedding: Conv2d,
|
|
num_positions: usize,
|
|
position_embedding: Embedding,
|
|
packing_position_embedding: Embedding,
|
|
}
|
|
|
|
impl SiglipVisionEmbeddings {
|
|
pub fn new(vb: VarBuilder, config: &PaddleOCRVLVisionConfig) -> Result<Self> {
|
|
let embed_dim = config.hidden_size;
|
|
let image_size = config.image_size;
|
|
let patch_size = config.patch_size;
|
|
let patch_embedding = get_conv2d(
|
|
vb.pp("patch_embedding"),
|
|
config.num_channels,
|
|
embed_dim,
|
|
patch_size,
|
|
0,
|
|
patch_size,
|
|
1,
|
|
1,
|
|
true,
|
|
)?;
|
|
let num_positions = (image_size / patch_size).pow(2);
|
|
let position_embedding = embedding(num_positions, embed_dim, vb.pp("position_embedding"))?;
|
|
let packing_position_embedding =
|
|
embedding(32768, embed_dim, vb.pp("packing_position_embedding"))?;
|
|
Ok(Self {
|
|
embed_dim,
|
|
patch_size,
|
|
patch_embedding,
|
|
num_positions,
|
|
position_embedding,
|
|
packing_position_embedding,
|
|
})
|
|
}
|
|
fn interpolate_pos_encoding(
|
|
&self,
|
|
h: usize,
|
|
w: usize,
|
|
is_after_patchify: bool,
|
|
) -> Result<Tensor> {
|
|
let (new_height, new_width) = if is_after_patchify {
|
|
(h, w)
|
|
} else {
|
|
(h / self.patch_size, w / self.patch_size)
|
|
};
|
|
let sqrt_num_positions = self.num_positions.sqrt();
|
|
let patch_pos_embed = self
|
|
.position_embedding
|
|
.embeddings()
|
|
.reshape((1, sqrt_num_positions, sqrt_num_positions, self.embed_dim))?
|
|
.permute((0, 3, 1, 2))?;
|
|
let patch_pos_embed =
|
|
interpolate_bilinear(&patch_pos_embed, (new_height, new_width), Some(false))?;
|
|
let patch_pos_embed =
|
|
patch_pos_embed
|
|
.permute((0, 2, 3, 1))?
|
|
.reshape((1, (), self.embed_dim))?;
|
|
Ok(patch_pos_embed)
|
|
}
|
|
|
|
pub fn forward(
|
|
&self,
|
|
pixel_values: &Tensor,
|
|
position_ids: &Tensor,
|
|
image_grid_thw: &Tensor,
|
|
interpolate_pos_encoding: bool,
|
|
) -> Result<Tensor> {
|
|
let (bs, seq_len, c, h, w) = pixel_values.dims5()?;
|
|
let pixel_values = pixel_values.reshape((bs * seq_len, c, h, w))?;
|
|
let patch_embeds = self.patch_embedding.forward(&pixel_values)?;
|
|
// (bs*seq_len, c)
|
|
let mut embeddings = patch_embeds.squeeze(D::Minus1)?.squeeze(D::Minus1)?;
|
|
if interpolate_pos_encoding {
|
|
let mut tmp_embeddings = vec![];
|
|
let img_num = image_grid_thw.dim(0)?;
|
|
let mut start = 0usize;
|
|
for i in 0..img_num {
|
|
let [t, h, w] = image_grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
|
|
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
|
|
};
|
|
let end = start + (t * h * w) as usize;
|
|
let image_embeddings = embeddings.i(start..end)?;
|
|
let position_embedding = self
|
|
.interpolate_pos_encoding(h as usize, w as usize, true)?
|
|
.squeeze(0)?
|
|
.repeat((t as usize, 1usize))?;
|
|
let image_embeddings = image_embeddings.add(&position_embedding)?;
|
|
tmp_embeddings.push(image_embeddings);
|
|
start = end;
|
|
}
|
|
embeddings = Tensor::cat(&tmp_embeddings, 0)?.unsqueeze(0)?; // add bs dim
|
|
} else {
|
|
let packing_pos_embed = self.packing_position_embedding.forward(position_ids)?;
|
|
embeddings = embeddings.add(&packing_pos_embed)?.unsqueeze(0)?;
|
|
}
|
|
Ok(embeddings)
|
|
}
|
|
}
|
|
|
|
pub struct SiglipEncoder {
|
|
layers: Vec<NaiveAttnTwoLinearMLPBlock>,
|
|
rotary_pos_emb: Qwen2_5VisionRotaryEmbedding,
|
|
}
|
|
|
|
impl SiglipEncoder {
|
|
pub fn new(vb: VarBuilder, config: &PaddleOCRVLVisionConfig) -> Result<Self> {
|
|
let vb_layers = vb.pp("layers");
|
|
let mut layers = vec![];
|
|
for i in 0..config.num_hidden_layers {
|
|
let layer_i = NaiveAttnTwoLinearMLPBlock::new(
|
|
vb_layers.pp(i),
|
|
config.hidden_size,
|
|
config.num_attention_heads,
|
|
None,
|
|
None,
|
|
true,
|
|
"self_attn",
|
|
Some("out_proj"),
|
|
config.intermediate_size,
|
|
config.hidden_act,
|
|
true,
|
|
"mlp",
|
|
"fc1",
|
|
"fc2",
|
|
config.layer_norm_eps,
|
|
"layer_norm1",
|
|
"layer_norm2",
|
|
)?;
|
|
layers.push(layer_i);
|
|
}
|
|
let head_dim = config.hidden_size / config.num_attention_heads;
|
|
let rotary_pos_emb = Qwen2_5VisionRotaryEmbedding::new(head_dim / 2, Some(10000.0));
|
|
Ok(Self {
|
|
layers,
|
|
rotary_pos_emb,
|
|
})
|
|
}
|
|
|
|
pub fn forward(&self, xs: &Tensor, image_grid_thw: &Tensor) -> Result<Tensor> {
|
|
let mut split_hids = vec![];
|
|
let mut split_wids = vec![];
|
|
for i in 0..image_grid_thw.dim(0)? {
|
|
let [t, h, w] = image_grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
|
|
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
|
|
};
|
|
let pos_w: Vec<u32> = (0..h).flat_map(|_| 0u32..w).collect();
|
|
let pos_w = pos_w.repeat(t as usize);
|
|
let pos_w = Tensor::new(pos_w, xs.device())?;
|
|
let pos_h: Vec<u32> = (0..h).flat_map(|h| vec![h; w as usize]).collect();
|
|
let pos_h = pos_h.repeat(t as usize);
|
|
let pos_h = Tensor::new(pos_h, xs.device())?;
|
|
split_hids.push(pos_h);
|
|
split_wids.push(pos_w);
|
|
}
|
|
let width_position_ids = Tensor::cat(&split_wids, 0)?;
|
|
let height_position_ids = Tensor::cat(&split_hids, 0)?;
|
|
let max_grid_size = image_grid_thw.i((.., 1..))?.max_all()?.to_scalar::<u32>()?;
|
|
let rope_emb_max_grid = self
|
|
.rotary_pos_emb
|
|
.forward(max_grid_size as usize, xs.device())?;
|
|
let rotary_pos_emb_h = rope_emb_max_grid.index_select(&height_position_ids, 0)?;
|
|
let rotary_pos_emb_w = rope_emb_max_grid.index_select(&width_position_ids, 0)?;
|
|
let rope_emb = Tensor::cat(&[rotary_pos_emb_h, rotary_pos_emb_w], 1)?.contiguous()?;
|
|
let rope_emb = rope_emb.repeat((1, 2))?;
|
|
let cos = rope_emb.cos()?;
|
|
let sin = rope_emb.sin()?;
|
|
let mut xs = xs.clone();
|
|
for layer in &self.layers {
|
|
xs = layer.forward(&xs, Some(&cos), Some(&sin), None, false)?;
|
|
}
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
pub struct SiglipVisionModel {
|
|
embeddings: SiglipVisionEmbeddings,
|
|
encoder: SiglipEncoder,
|
|
post_layernorm: LayerNorm,
|
|
}
|
|
impl SiglipVisionModel {
|
|
pub fn new(vb: VarBuilder, config: &PaddleOCRVLVisionConfig) -> Result<Self> {
|
|
let vb = vb.pp("vision_model");
|
|
let embeddings = SiglipVisionEmbeddings::new(vb.pp("embeddings"), config)?;
|
|
let encoder = SiglipEncoder::new(vb.pp("encoder"), config)?;
|
|
let post_layernorm = get_layer_norm(
|
|
vb.pp("post_layernorm"),
|
|
config.layer_norm_eps,
|
|
config.hidden_size,
|
|
)?;
|
|
Ok(Self {
|
|
embeddings,
|
|
encoder,
|
|
post_layernorm,
|
|
})
|
|
}
|
|
|
|
pub fn forward(
|
|
&self,
|
|
pixel_values: &Tensor,
|
|
image_grid_thw: &Tensor,
|
|
position_ids: &Tensor,
|
|
interpolate_pos_encoding: bool,
|
|
) -> Result<Tensor> {
|
|
let xs = self.embeddings.forward(
|
|
pixel_values,
|
|
position_ids,
|
|
image_grid_thw,
|
|
interpolate_pos_encoding,
|
|
)?;
|
|
let xs = self.encoder.forward(&xs, image_grid_thw)?;
|
|
let xs = self.post_layernorm.forward(&xs)?;
|
|
Ok(xs)
|
|
}
|
|
}
|
|
|
|
pub struct Ernie4_5Model {
|
|
embed_tokens: Embedding,
|
|
layers: Vec<NaiveAttnGateUpDownMLPBlock>,
|
|
norm: RmsNorm,
|
|
rotary_emb: Qwen2_5VLTextRotaryEmbedding,
|
|
rope_scaling: PaddleOCRVLRopeScalingConfig,
|
|
}
|
|
|
|
impl Ernie4_5Model {
|
|
pub fn new(vb: VarBuilder, config: &PaddleOCRVLConfig) -> Result<Self> {
|
|
let embed_tokens = embedding(config.vocab_size, config.hidden_size, vb.pp("embed_tokens"))?;
|
|
let vb_layers = vb.pp("layers");
|
|
let mut layers = vec![];
|
|
for i in 0..config.num_hidden_layers {
|
|
let layer_i = NaiveAttnGateUpDownMLPBlock::new(
|
|
vb_layers.pp(i),
|
|
config.hidden_size,
|
|
config.num_attention_heads,
|
|
Some(config.num_key_value_heads),
|
|
Some(config.head_dim),
|
|
config.use_bias,
|
|
"self_attn",
|
|
None,
|
|
config.intermediate_size,
|
|
config.hidden_act,
|
|
config.use_bias,
|
|
"mlp",
|
|
config.rms_norm_eps,
|
|
"input_layernorm",
|
|
"post_attention_layernorm",
|
|
)?;
|
|
layers.push(layer_i);
|
|
}
|
|
let norm = rms_norm(config.hidden_size, config.rms_norm_eps, vb.pp("norm"))?;
|
|
let rotary_emb =
|
|
Qwen2_5VLTextRotaryEmbedding::new(config.head_dim, config.rope_theta as f32);
|
|
|
|
Ok(Self {
|
|
embed_tokens,
|
|
layers,
|
|
norm,
|
|
rotary_emb,
|
|
rope_scaling: config.rope_scaling.clone(),
|
|
})
|
|
}
|
|
|
|
pub fn forward(
|
|
&mut self,
|
|
inputs_embeds: &Tensor,
|
|
seqlen_offset: usize,
|
|
position_ids: Option<&Tensor>,
|
|
) -> Result<Tensor> {
|
|
let (b_size, seq_len, _) = inputs_embeds.dims3()?;
|
|
|
|
let position_ids = match position_ids {
|
|
Some(ids) => ids.clone(),
|
|
None => Tensor::arange(
|
|
seqlen_offset as u32,
|
|
(seq_len + seqlen_offset) as u32,
|
|
inputs_embeds.device(),
|
|
)?
|
|
.unsqueeze(0)?
|
|
.unsqueeze(0)?
|
|
.broadcast_as((3, b_size, seq_len))?,
|
|
};
|
|
let (cos, sin) = self.rotary_emb.forward(
|
|
&position_ids,
|
|
inputs_embeds.dtype(),
|
|
self.rope_scaling.mrope_section.clone(),
|
|
)?;
|
|
let mut xs = inputs_embeds.clone();
|
|
let attention_mask: Option<&Tensor> = {
|
|
if seq_len <= 1 {
|
|
None
|
|
} else {
|
|
Some(&prepare_causal_attention_mask(
|
|
b_size,
|
|
seq_len,
|
|
0,
|
|
xs.device(),
|
|
)?)
|
|
}
|
|
};
|
|
for layer in self.layers.iter_mut() {
|
|
xs = layer.forward(&xs, &cos, &sin, attention_mask)?;
|
|
}
|
|
let xs = xs.apply(&self.norm)?;
|
|
Ok(xs)
|
|
}
|
|
|
|
pub fn clear_kv_cache(&mut self) {
|
|
for layer in self.layers.iter_mut() {
|
|
layer.clear_kv_cache()
|
|
}
|
|
}
|
|
}
|
|
|
|
pub struct PaddleOCRVLModel {
|
|
mlp_ar: Projector,
|
|
visual: SiglipVisionModel,
|
|
model: Ernie4_5Model,
|
|
pub cfg: PaddleOCRVLConfig,
|
|
lm_head: Linear,
|
|
rope_deltas: Option<Tensor>,
|
|
}
|
|
|
|
impl PaddleOCRVLModel {
|
|
pub fn new(cfg: PaddleOCRVLConfig, vb: VarBuilder) -> Result<Self> {
|
|
let mlp_ar = Projector::new(vb.pp("mlp_AR"), &cfg)?;
|
|
let visual = SiglipVisionModel::new(vb.pp("visual"), &cfg.vision_config)?;
|
|
let model = Ernie4_5Model::new(vb.pp("model"), &cfg)?;
|
|
let vocab_size = cfg.vocab_size;
|
|
let lm_head = if cfg.tie_word_embeddings {
|
|
Linear::new(model.embed_tokens.embeddings().clone(), None)
|
|
} else {
|
|
linear_no_bias(cfg.hidden_size, vocab_size, vb.pp("lm_head"))?
|
|
};
|
|
|
|
Ok(Self {
|
|
mlp_ar,
|
|
visual,
|
|
model,
|
|
cfg,
|
|
lm_head,
|
|
rope_deltas: None,
|
|
})
|
|
}
|
|
|
|
pub fn get_rope_index(
|
|
&self,
|
|
input_ids: &Tensor,
|
|
image_grid_thw: Option<&Tensor>,
|
|
video_grid_thw: Option<&Tensor>,
|
|
mask: Option<&Tensor>,
|
|
second_per_grid_ts: Option<Vec<f32>>,
|
|
) -> Result<(Tensor, Tensor)> {
|
|
let spatial_merge_size = self.cfg.vision_config.spatial_merge_size;
|
|
let mut mrope_position_deltas: Vec<i64> = Vec::new();
|
|
if image_grid_thw.is_some() || video_grid_thw.is_some() {
|
|
let total_input_ids = input_ids.clone();
|
|
let mask_ = mask
|
|
.cloned()
|
|
.unwrap_or(Tensor::ones_like(&total_input_ids)?);
|
|
let mut position_ids = Tensor::ones(
|
|
(3, input_ids.dim(0)?, input_ids.dim(1)?),
|
|
input_ids.dtype(),
|
|
input_ids.device(),
|
|
)?;
|
|
let mut image_index = 0;
|
|
let mut video_index = 0;
|
|
for i in 0..total_input_ids.dim(0)? {
|
|
let mut input_ids_i = total_input_ids.i(i)?;
|
|
let mask_i = mask_.i(i)?;
|
|
// 推理时, attention_mask如果是全1向量,取非0索引的操作没必要
|
|
if mask_i.sum_all()?.to_scalar::<u32>()? != mask_i.dim(0)? as u32 {
|
|
let nonzero_idx = nonzero_index(&mask_i)?;
|
|
input_ids_i = input_ids_i.gather(&nonzero_idx, 0)?;
|
|
}
|
|
let mut text_start = 0;
|
|
let mut text_end = 0;
|
|
let mut thw = vec![];
|
|
let mut second_per_grid_t = 0_f32;
|
|
let mut llm_pos_ids_list: Vec<Tensor> = Vec::new();
|
|
// vision start的下一个索引
|
|
let vision_indices =
|
|
get_vision_next_indices(&input_ids_i, self.cfg.vision_start_token_id);
|
|
match vision_indices {
|
|
Ok(indeices) => {
|
|
let vision_tokens = input_ids_i.gather(&indeices, 0)?.to_vec1::<u32>()?;
|
|
let vision_indices_vec = indeices.to_vec1::<u32>()?;
|
|
for (j, &token) in vision_tokens.iter().enumerate() {
|
|
if token == self.cfg.image_token_id {
|
|
thw = image_grid_thw.unwrap().i(image_index)?.to_vec1::<u32>()?;
|
|
image_index += 1;
|
|
text_end = vision_indices_vec[j];
|
|
second_per_grid_t = 0.0;
|
|
}
|
|
if token == self.cfg.video_token_id {
|
|
thw = video_grid_thw.unwrap().i(video_index)?.to_vec1::<u32>()?;
|
|
text_end = vision_indices_vec[j];
|
|
second_per_grid_t = match second_per_grid_ts {
|
|
None => 1.0,
|
|
Some(ref vec) => vec[video_index],
|
|
};
|
|
video_index += 1;
|
|
}
|
|
let llm_grid_t = thw[0];
|
|
let llm_grid_h = thw[1] / spatial_merge_size as u32;
|
|
let llm_grid_w = thw[2] / spatial_merge_size as u32;
|
|
let text_len = text_end - text_start;
|
|
let start_idx = if !llm_pos_ids_list.is_empty() {
|
|
llm_pos_ids_list[llm_pos_ids_list.len() - 1]
|
|
.max_all()?
|
|
.to_scalar::<u32>()?
|
|
+ 1
|
|
} else {
|
|
0
|
|
};
|
|
let pos_ids = Tensor::arange(
|
|
start_idx,
|
|
start_idx + text_len,
|
|
input_ids_i.device(),
|
|
)?
|
|
.unsqueeze(0)?
|
|
.broadcast_as((3usize, text_len as usize))?;
|
|
llm_pos_ids_list.push(pos_ids);
|
|
let range_tensor = Tensor::arange(0, llm_grid_t, input_ids_i.device())?
|
|
.unsqueeze(D::Minus1)?;
|
|
let expanded_range = range_tensor.broadcast_as((
|
|
llm_grid_t as usize,
|
|
(llm_grid_h * llm_grid_w) as usize,
|
|
))?;
|
|
let time_tensor = expanded_range
|
|
.broadcast_mul(&Tensor::new(
|
|
vec![
|
|
(second_per_grid_t
|
|
* self.cfg.vision_config.tokens_per_second as f32)
|
|
as u32,
|
|
],
|
|
input_ids_i.device(),
|
|
)?)?
|
|
.broadcast_add(&Tensor::new(
|
|
vec![start_idx + text_len],
|
|
input_ids_i.device(),
|
|
)?)?;
|
|
let t_index = time_tensor.flatten_all()?;
|
|
let h_index = Tensor::arange(
|
|
start_idx + text_len,
|
|
start_idx + text_len + llm_grid_h,
|
|
input_ids_i.device(),
|
|
)?
|
|
.unsqueeze(0)?
|
|
.unsqueeze(D::Minus1)?
|
|
.broadcast_as((
|
|
llm_grid_t as usize,
|
|
llm_grid_h as usize,
|
|
llm_grid_w as usize,
|
|
))?
|
|
.flatten_all()?;
|
|
let w_index = Tensor::arange(
|
|
start_idx + text_len,
|
|
start_idx + text_len + llm_grid_w,
|
|
input_ids_i.device(),
|
|
)?
|
|
.unsqueeze(0)?
|
|
.unsqueeze(0)?
|
|
.broadcast_as((
|
|
llm_grid_t as usize,
|
|
llm_grid_h as usize,
|
|
llm_grid_w as usize,
|
|
))?
|
|
.flatten_all()?;
|
|
|
|
let thw_index = Tensor::stack(&[t_index, h_index, w_index], 0)?;
|
|
llm_pos_ids_list.push(thw_index);
|
|
text_start = text_end + llm_grid_t * llm_grid_h * llm_grid_w;
|
|
}
|
|
}
|
|
Err(e) => {
|
|
println!("get vision_indices err: {}", e);
|
|
}
|
|
};
|
|
|
|
if text_start < input_ids_i.dim(0)? as u32 {
|
|
let start_idx = if !llm_pos_ids_list.is_empty() {
|
|
llm_pos_ids_list[llm_pos_ids_list.len() - 1]
|
|
.max_all()?
|
|
.to_scalar::<u32>()?
|
|
+ 1
|
|
} else {
|
|
0
|
|
};
|
|
let text_len = input_ids_i.dim(0)? as u32 - text_start;
|
|
let pos_ids =
|
|
Tensor::arange(start_idx, start_idx + text_len, input_ids_i.device())?
|
|
.unsqueeze(0)?
|
|
.broadcast_as((3usize, text_len as usize))?;
|
|
llm_pos_ids_list.push(pos_ids);
|
|
}
|
|
let llm_position = Tensor::cat(&llm_pos_ids_list, 1)?.reshape((3, 1, ()))?;
|
|
position_ids = position_ids
|
|
.slice_assign(&[(0..3), (i..i + 1), (0..input_ids.dim(1)?)], &llm_position)?;
|
|
let position_deltas = llm_position.max_all()?.to_scalar::<u32>()? as i64 + 1
|
|
- input_ids_i.dim(0)? as i64;
|
|
mrope_position_deltas.push(position_deltas);
|
|
}
|
|
|
|
let mut mrope_position_deltas = Tensor::new(mrope_position_deltas, input_ids.device())?;
|
|
if mrope_position_deltas.rank() == 1 {
|
|
mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?;
|
|
}
|
|
Ok((position_ids.contiguous()?, mrope_position_deltas))
|
|
} else if let Some(mask) = mask {
|
|
let mut position_ids = mask
|
|
.to_dtype(candle_core::DType::F64)?
|
|
.cumsum(D::Minus1)?
|
|
.to_dtype(candle_core::DType::U32)?
|
|
.broadcast_sub(&Tensor::new(vec![1_u32], input_ids.device())?)?;
|
|
for i in 0..position_ids.dim(0)? {
|
|
let mut position_ids_i = position_ids.i(i)?;
|
|
let mask_i = mask.i(i)?;
|
|
// 如果有pad, 将填充位置置为1
|
|
// 当bs>1, 可能存在不同序列长度，需要添加pad使seq_len长度一致
|
|
if mask_i.sum_all()?.to_scalar::<u32>()? != mask_i.dim(0)? as u32 {
|
|
let zero_indices = zero_index(&mask_i)?;
|
|
let replace_1 = Tensor::ones(
|
|
zero_indices.dim(0)?,
|
|
candle_core::DType::U32,
|
|
input_ids.device(),
|
|
)?;
|
|
position_ids_i = position_ids_i
|
|
.scatter(&zero_indices, &replace_1, 0)?
|
|
.unsqueeze(0)?;
|
|
position_ids = position_ids
|
|
.slice_assign(&[(i..i + 1), (0..position_ids.dim(1)?)], &position_ids_i)?;
|
|
}
|
|
}
|
|
position_ids = position_ids
|
|
.unsqueeze(0)?
|
|
.broadcast_as((3, input_ids.dim(0)?, input_ids.dim(1)?))?
|
|
.contiguous()?;
|
|
let mut mrope_position_deltas = position_ids
|
|
.max(0)?
|
|
.max(D::Minus1)?
|
|
.broadcast_sub(&Tensor::new(
|
|
vec![mask.dim(D::Minus1)? as u32 - 1],
|
|
input_ids.device(),
|
|
)?)?
|
|
.contiguous()?;
|
|
if mrope_position_deltas.rank() == 1 {
|
|
mrope_position_deltas = mrope_position_deltas.unsqueeze(0)?;
|
|
}
|
|
Ok((position_ids, mrope_position_deltas))
|
|
} else {
|
|
let position_ids =
|
|
Tensor::arange(0_u32, input_ids.dim(D::Minus1)? as u32, input_ids.device())?
|
|
.unsqueeze(0)?
|
|
.unsqueeze(0)?
|
|
.broadcast_as((3, input_ids.dim(0)?, input_ids.dim(D::Minus1)?))?
|
|
.contiguous()?;
|
|
let mrope_position_deltas = Tensor::zeros(
|
|
(input_ids.dim(0)?, 1),
|
|
input_ids.dtype(),
|
|
input_ids.device(),
|
|
)?;
|
|
Ok((position_ids, mrope_position_deltas))
|
|
}
|
|
}
|
|
|
|
pub fn forward(
|
|
&mut self,
|
|
input_ids: &Tensor,
|
|
pixel_values: Option<&Tensor>,
|
|
image_grid_thw: Option<&Tensor>,
|
|
image_mask: &Tensor,
|
|
cache_position: Option<&Tensor>,
|
|
seqlen_offset: usize,
|
|
) -> Result<Tensor> {
|
|
let mut inputs_embeds = self.model.embed_tokens.forward(input_ids)?;
|
|
if let Some(pixel_values) = pixel_values
|
|
&& let Some(image_grid_thw) = image_grid_thw
|
|
{
|
|
let pixel_values = pixel_values.unsqueeze(0)?;
|
|
let mut siglip_position_ids = vec![];
|
|
let mut sample_indices = vec![];
|
|
let mut cu_seqlens = vec![0u32];
|
|
let img_num = image_grid_thw.dim(0)?;
|
|
for i in 0..img_num {
|
|
let [t, h, w] = image_grid_thw.i(i)?.to_vec1::<u32>()?[..] else {
|
|
return Err(anyhow!(format!("grid_thw Expected exactly 3 elements")));
|
|
};
|
|
let numel = h * w;
|
|
let image_position_ids =
|
|
Tensor::arange(0, numel, pixel_values.device())?.repeat(t as usize)?;
|
|
siglip_position_ids.push(image_position_ids);
|
|
let indices =
|
|
Tensor::new(vec![i as u32; (numel * t) as usize], pixel_values.device())?;
|
|
sample_indices.push(indices);
|
|
cu_seqlens.push(cu_seqlens[cu_seqlens.len() - 1] + numel * t);
|
|
}
|
|
let siglip_position_ids = Tensor::cat(&siglip_position_ids, 0)?;
|
|
|
|
let image_embed =
|
|
self.visual
|
|
.forward(&pixel_values, image_grid_thw, &siglip_position_ids, true)?;
|
|
let image_embed = image_embed.squeeze(0)?;
|
|
let image_embed = self.mlp_ar.forward(&image_embed, image_grid_thw)?;
|
|
inputs_embeds = masked_scatter_dim0(&inputs_embeds, &image_embed, image_mask)?;
|
|
}
|
|
let position_ids;
|
|
let rope_deltas;
|
|
if (cache_position.is_some() && cache_position.unwrap().i(0)?.to_scalar::<u32>()? == 0)
|
|
|| self.rope_deltas.is_none()
|
|
{
|
|
(position_ids, rope_deltas) =
|
|
self.get_rope_index(input_ids, image_grid_thw, None, None, None)?;
|
|
self.rope_deltas = Some(rope_deltas);
|
|
} else {
|
|
let (bs, seq_len, _) = inputs_embeds.dims3()?;
|
|
let delta = if let Some(cache_position) = cache_position {
|
|
cache_position
|
|
.i(0)?
|
|
.to_dtype(self.rope_deltas.as_ref().unwrap().dtype())?
|
|
.broadcast_add(self.rope_deltas.as_ref().unwrap())?
|
|
.contiguous()?
|
|
.to_dtype(candle_core::DType::U32)?
|
|
} else {
|
|
Tensor::zeros(1, inputs_embeds.dtype(), inputs_embeds.device())?
|
|
};
|
|
position_ids = Tensor::arange(0u32, seq_len as u32, input_ids.device())?
|
|
.unsqueeze(0)?
|
|
.broadcast_as((bs, seq_len))?
|
|
.broadcast_add(&delta)?
|
|
.unsqueeze(0)?
|
|
.broadcast_as((3, bs, seq_len))?
|
|
.contiguous()?;
|
|
}
|
|
let outputs = self
|
|
.model
|
|
.forward(&inputs_embeds, seqlen_offset, Some(&position_ids))?;
|
|
let seq_len = outputs.dim(1)?;
|
|
let hidden_state = outputs.narrow(1, seq_len - 1, 1)?;
|
|
let logits = self.lm_head.forward(&hidden_state)?;
|
|
Ok(logits)
|
|
}
|
|
|
|
pub fn clear_kv_cache(&mut self) {
|
|
self.model.clear_kv_cache();
|
|
}
|
|
}
|