tckm.yan/aha
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add hunyuan_ocr

Browse Source
This commit is contained in:
jhqxxx
2025-12-03 17:21:01 +08:00
parent 7d72cb3baf
commit 697484cf23
29 changed files with 1756 additions and 190 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/models/common/mod.rs
+119 -41
View File
@@ -1,27 +1,41 @@
use anyhow::Result;
use candle_core::{D, Tensor};
use candle_nn::{Activation, Linear, Module, VarBuilder, linear, linear_no_bias};
use candle_nn::{
Activation, Conv2d, Conv2dConfig, LayerNorm, LayerNormConfig, Linear, Module, VarBuilder,
conv2d, conv2d_no_bias, layer_norm, linear, linear_no_bias,
};
use crate::{position_embed::rope::apply_rotary_pos_emb, utils::tensor_utils::repeat_kv};
#[derive(Debug, Clone)]
pub struct MLPWithBias {
pub struct GateUpDownMLP {
gate_proj: Linear,
up_proj: Linear,
down_proj: Linear,
act_fn: Activation,
}
impl MLPWithBias {
impl GateUpDownMLP {
pub fn new(
vb: VarBuilder,
hidden_size: usize,
intermediate_size: usize,
act_fn: Activation,
bias: bool,
) -> Result<Self> {
let gate_proj = linear(hidden_size, intermediate_size, vb.pp("gate_proj"))?;
let up_proj = linear(hidden_size, intermediate_size, vb.pp("up_proj"))?;
let down_proj = linear(intermediate_size, hidden_size, vb.pp("down_proj"))?;
let (gate_proj, up_proj, down_proj) = if bias {
(
linear(hidden_size, intermediate_size, vb.pp("gate_proj"))?,
linear(hidden_size, intermediate_size, vb.pp("up_proj"))?,
linear(intermediate_size, hidden_size, vb.pp("down_proj"))?,
)
} else {
(
linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?,
linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?,
linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?,
)
};
Ok(Self {
gate_proj,
up_proj,
@@ -31,7 +45,7 @@ impl MLPWithBias {
}
}
impl Module for MLPWithBias {
impl Module for GateUpDownMLP {
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?;
let rhs = xs.apply(&self.up_proj)?;
@@ -39,43 +53,51 @@ impl Module for MLPWithBias {
}
}
#[derive(Debug, Clone)]
pub struct MLPNoBias {
gate_proj: Linear,
up_proj: Linear,
down_proj: Linear,
act_fn: Activation,
pub struct TwoLinearMLP {
linear1: Linear,
linear2: Linear,
act: Activation,
}
impl MLPNoBias {
impl TwoLinearMLP {
pub fn new(
vb: VarBuilder,
hidden_size: usize,
intermediate_size: usize,
act_fn: Activation,
embedding_dim: usize,
mlp_dim: usize,
act: Activation,
bias: bool,
linear1_pp_name: &str,
linear2_pp_name: &str,
) -> Result<Self> {
let gate_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?;
let up_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?;
let down_proj = linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?;
let (linear1, linear2) = if bias {
(
linear(embedding_dim, mlp_dim, vb.pp(linear1_pp_name))?,
linear(mlp_dim, embedding_dim, vb.pp(linear2_pp_name))?,
)
} else {
(
linear_no_bias(embedding_dim, mlp_dim, vb.pp(linear1_pp_name))?,
linear_no_bias(mlp_dim, embedding_dim, vb.pp(linear2_pp_name))?,
)
};
Ok(Self {
gate_proj,
up_proj,
down_proj,
act_fn,
linear1,
linear2,
act,
})
}
}
impl Module for MLPNoBias {
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?;
let rhs = xs.apply(&self.up_proj)?;
(lhs * rhs)?.apply(&self.down_proj)
pub fn forward(&self, xs: &Tensor) -> Result<Tensor> {
let xs = xs
.apply(&self.linear1)?
.apply(&self.act)?
.apply(&self.linear2)?;
Ok(xs)
}
}
#[derive(Debug, Clone)]
pub struct AttentionNobias {
// pub struct AttentionNobias {
pub struct NaiveAttention {
q_proj: Linear,
k_proj: Linear,
v_proj: Linear,
@@ -88,19 +110,33 @@ pub struct AttentionNobias {
kv_cache: Option<(Tensor, Tensor)>,
}
impl AttentionNobias {
// impl AttentionNobias {
impl NaiveAttention {
pub fn new(
vb: VarBuilder,
hidden_size: usize,
num_attention_heads: usize,
num_key_value_heads: usize,
bias: bool,
) -> Result<Self> {
let num_kv_groups = num_attention_heads / num_key_value_heads;
let head_dim = hidden_size / num_attention_heads;
let q_proj = linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?;
let k_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?;
let v_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?;
let o_proj = linear_no_bias(hidden_size, hidden_size, vb.pp("o_proj"))?;
let (q_proj, k_proj, v_proj, o_proj) = if bias {
(
linear(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?,
linear(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?,
linear(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?,
linear(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?,
)
} else {
(
linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?,
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?,
linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?,
linear_no_bias(num_attention_heads * head_dim, hidden_size, vb.pp("o_proj"))?,
)
};
Ok(Self {
q_proj,
k_proj,
@@ -118,8 +154,8 @@ impl AttentionNobias {
pub fn forward(
&self,
xs: &Tensor,
cos: &Tensor,
sin: &Tensor,
cos: Option<&Tensor>,
sin: Option<&Tensor>,
attention_mask: Option<&Tensor>,
tof32: bool,
) -> Result<Tensor> {
@@ -136,8 +172,14 @@ impl AttentionNobias {
let value_states = value_states
.reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
.transpose(1, 2)?;
let (query_states, key_states) =
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?;
let (query_states, key_states) = if let Some(cos) = cos
&& let Some(sin) = sin
{
apply_rotary_pos_emb(&query_states, &key_states, cos, sin, tof32)?
} else {
(query_states, key_states)
};
let scale = 1f64 / f64::sqrt(self.head_dim as f64);
let attn_output = eager_attention_forward(
&query_states,
@@ -260,3 +302,39 @@ pub fn eager_attention_forward(
Ok(attn_output)
}
pub fn get_conv2d(
vb: VarBuilder,
in_c: usize,
out_c: usize,
kernel_size: usize,
padding: usize,
stride: usize,
dilation: usize,
groups: usize,
bias: bool,
) -> Result<Conv2d> {
let cfg = Conv2dConfig {
padding,
stride,
dilation,
groups,
cudnn_fwd_algo: None,
};
let conv2d = if bias {
conv2d(in_c, out_c, kernel_size, cfg, vb)?
} else {
conv2d_no_bias(in_c, out_c, kernel_size, cfg, vb)?
};
Ok(conv2d)
}
pub fn get_layer_norm(vb: VarBuilder, eps: f64, dim: usize) -> Result<LayerNorm> {
let ln_config = LayerNormConfig {
eps,
remove_mean: true, // true for layernorm, false for RMSNorm
affine: true, // true for with bias, false for without bias
};
let norm = layer_norm(dim, ln_config, vb)?;
Ok(norm)
}