1use crate::{
2 models::common::{InferenceModel, modules::NaiveAttnGateUpDownMLPBlock},
3 position_embed::rope::RoPE,
4 utils::tensor_utils::prepare_causal_attention_mask,
5};
6use anyhow::Result;
7use candle_core::Tensor;
8use candle_nn::{
9 Activation, Embedding, Linear, Module, RmsNorm, VarBuilder, embedding, linear_no_bias, rms_norm,
10};
11
12pub struct LlamaModel {
13 pub embed_tokens: Embedding,
14 layers: Vec<NaiveAttnGateUpDownMLPBlock>,
15 norm: RmsNorm,
16 rotary_emb: RoPE,
17}
18
19impl LlamaModel {
20 pub fn new(
21 vb: VarBuilder,
22 vocab_size: usize,
23 hidden_size: usize,
24 num_hidden_layers: usize,
25 num_attention_heads: usize,
26 num_key_value_heads: Option<usize>,
27 head_dim: Option<usize>,
28 attn_bias: bool,
29 attn_pp_name: &str,
30 o_proj_pp_name: Option<&str>,
31 intermediate_size: usize,
32 hidden_act: Activation,
33 mlp_bias: bool,
34 mlp_pp_name: &str,
35 norm_eps: f64,
36 input_norm_pp_name: &str,
37 post_norm_pp_name: &str,
38 rope_theta_base: f32,
39 ) -> Result<Self> {
40 let embed_tokens = embedding(vocab_size, hidden_size, vb.pp("embed_tokens"))?;
41 let mut layers = vec![];
42 let vb_layers = vb.pp("layers");
43 for i in 0..num_hidden_layers {
44 let layers_i = NaiveAttnGateUpDownMLPBlock::new(
45 vb_layers.pp(i),
46 hidden_size,
47 num_attention_heads,
48 num_key_value_heads,
49 head_dim,
50 attn_bias,
51 attn_pp_name,
52 o_proj_pp_name,
53 intermediate_size,
54 hidden_act,
55 mlp_bias,
56 mlp_pp_name,
57 norm_eps,
58 input_norm_pp_name,
59 post_norm_pp_name,
60 )?;
61 layers.push(layers_i);
62 }
63 let norm = rms_norm(hidden_size, norm_eps, vb.pp("norm"))?;
64 let head_dim = head_dim.unwrap_or(hidden_size / num_attention_heads);
65 let rotary_emb = RoPE::new(head_dim, rope_theta_base, vb.device())?;
66 Ok(Self {
67 embed_tokens,
68 layers,
69 norm,
70 rotary_emb,
71 })
72 }
73
74 pub fn forward(&mut self, inputs_embeds: &Tensor, seqlen_offset: usize) -> Result<Tensor> {
75 let (b_size, seq_len, _) = inputs_embeds.dims3()?;
76
77 let (cos, sin) = self
78 .rotary_emb
79 .forward(seqlen_offset, seq_len, inputs_embeds.device())?;
80 let mut xs = inputs_embeds.clone();
81 let attention_mask: Option<Tensor> = {
82 if seq_len <= 1 {
83 None
84 } else {
85 Some(prepare_causal_attention_mask(
86 b_size,
87 seq_len,
88 0,
89 xs.device(),
90 )?)
91 }
92 };
93 for layer in self.layers.iter_mut() {
94 xs = layer.forward(&xs, &cos, &sin, attention_mask.as_ref())?;
95 }
96 let xs = xs.apply(&self.norm)?;
97 Ok(xs)
98 }
99
100 pub fn clear_kv_cache(&mut self) {
101 for layer in self.layers.iter_mut() {
102 layer.clear_kv_cache()
103 }
104 }
105}
106
107pub struct LlamaForCausalLM {
108 pub model: LlamaModel,
109 lm_head: Linear,
110 stop_token_ids: Vec<u32>,
111}
112
113impl LlamaForCausalLM {
114 pub fn new(
115 vb: VarBuilder,
116 vocab_size: usize,
117 hidden_size: usize,
118 num_hidden_layers: usize,
119 num_attention_heads: usize,
120 num_key_value_heads: Option<usize>,
121 head_dim: Option<usize>,
122 attn_bias: bool,
123 attn_pp_name: &str,
124 o_proj_pp_name: Option<&str>,
125 intermediate_size: usize,
126 hidden_act: Activation,
127 mlp_bias: bool,
128 mlp_pp_name: &str,
129 norm_eps: f64,
130 input_norm_pp_name: &str,
131 post_norm_pp_name: &str,
132 rope_theta_base: f32,
133 eos_ids: Vec<u32>,
134 ) -> Result<Self> {
135 let model = LlamaModel::new(
136 vb.pp("model"),
137 vocab_size,
138 hidden_size,
139 num_hidden_layers,
140 num_attention_heads,
141 num_key_value_heads,
142 head_dim,
143 attn_bias,
144 attn_pp_name,
145 o_proj_pp_name,
146 intermediate_size,
147 hidden_act,
148 mlp_bias,
149 mlp_pp_name,
150 norm_eps,
151 input_norm_pp_name,
152 post_norm_pp_name,
153 rope_theta_base,
154 )?;
155 let lm_head = linear_no_bias(hidden_size, vocab_size, vb.pp("lm_head"))?;
156 Ok(Self {
157 model,
158 lm_head,
159 stop_token_ids: eos_ids,
160 })
161 }
162
163 pub fn forward_embeds(
164 &mut self,
165 inputs_embeds: &Tensor,
166 seqlen_offset: usize,
167 ) -> Result<Tensor> {
168 let outputs = self.model.forward(inputs_embeds, seqlen_offset)?;
169 let seq_len = outputs.dim(1)?;
170 let hidden_state = outputs.narrow(1, seq_len - 1, 1)?;
171 let logits = self.lm_head.forward(&hidden_state)?;
172 Ok(logits)
173 }
174 pub fn clear_kv_cache(&mut self) {
175 self.model.clear_kv_cache();
176 }
177}
178
179impl InferenceModel for LlamaForCausalLM {
180 fn forward_step(&mut self, input_ids: &Tensor, seqlen_offset: usize) -> Result<Tensor> {
181 let input_embeds = self.model.embed_tokens.forward(input_ids)?;
182 self.forward_embeds(&input_embeds, seqlen_offset)
183 }
184
185 fn clear_cache(&mut self) {
186 self.clear_kv_cache();
187 }
188
189 fn stop_token_ids(&self) -> Vec<u32> {
190 self.stop_token_ids.clone()
191 }
192}