Skip to content

Commit 3abc58d

Browse files
alvarobarttalvarobartt
authored
Set pad_token_id as nullable & add support for rope_parameters (#832)
1 parent 77ee99d commit 3abc58d

17 files changed

Lines changed: 113 additions & 24 deletions

File tree

‎README.md‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -250,7 +250,7 @@ Options:
250250
[env: DENSE_PATH=]
251251

252252
--hf-token <HF_TOKEN>
253-
Your Hugging Face Hub token. If neither `--hf-token` nor `HF_TOKEN` is set, the token will be read from the `$HF_HOME/token` path, if it exists. This ensures access to private or gated models, and allows for a more permissive rate limiting
253+
Your Hugging Face Hub token. If neither `--hf-token` nor `HF_TOKEN` are set, the token will be read from the `$HF_HOME/token` path, if it exists. This ensures access to private or gated models, and allows for a more permissive rate limiting
254254

255255
[env: HF_TOKEN=]
256256

‎backends/candle/src/layers/mod.rs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,4 +14,4 @@ pub use layer_norm::{LayerNorm, LayerNormNoBias};
1414
pub use linear::{HiddenAct, Linear};
1515
#[allow(unused_imports)]
1616
pub use rms_norm::RMSNorm;
17-
pub use rotary::{apply_rotary, get_cos_sin, get_inv_freqs, RopeScaling};
17+
pub use rotary::{apply_rotary, get_cos_sin, get_inv_freqs, RopeParameters, RopeScaling};

‎backends/candle/src/layers/rotary.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,13 @@
11
use candle::{DType, Device, Result, Tensor, D};
22
use serde::Deserialize;
33

4+
#[derive(Debug, Clone, PartialEq, Deserialize)]
5+
pub struct RopeParameters {
6+
pub rope_theta: f32,
7+
#[allow(unused)]
8+
rope_type: String,
9+
}
10+
411
#[derive(Debug, Clone, PartialEq, Deserialize)]
512
#[serde(untagged)]
613
pub enum RopeScaling {

‎backends/candle/src/lib.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -538,6 +538,7 @@ impl CandleBackend {
538538
rms_norm_eps: config.rms_norm_eps,
539539
model_type: config.model_type.clone(),
540540
rope_theta: config.rope_theta,
541+
rope_parameters: config.rope_parameters,
541542
sliding_window: config.sliding_window,
542543
rope_scaling: config.rope_scaling,
543544
use_bidirectional_attention: config.use_bidirectional_attention,

‎backends/candle/src/models/flash_gte.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -199,9 +199,18 @@ impl FlashGTEModel {
199199
Self::inner_load(vb.pp("new"), config)
200200
.or_else(|_| Self::inner_load(vb.clone(), config))?;
201201

202+
// NOTE: https://github.com/huggingface/transformers/pull/39847
203+
let rope_theta = match config.rope_theta {
204+
Some(rope_theta) => rope_theta,
205+
None => match &config.rope_parameters {
206+
Some(rope_parameters) => rope_parameters.rope_theta,
207+
None => candle::bail!("Neither `rope_theta` nor `rope_parameters.rope_theta` are defined in the `config.json`"),
208+
},
209+
};
210+
202211
let inv_freqs = get_inv_freqs(
203212
layers[0].attention.attention_head_size,
204-
config.rope_theta,
213+
rope_theta,
205214
vb.device(),
206215
config.rope_scaling.as_ref(),
207216
)?;

‎backends/candle/src/models/flash_mistral.rs‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -268,9 +268,18 @@ impl FlashMistralModel {
268268

269269
let norm = RMSNorm::load(vb.pp("norm"), config.hidden_size, config.rms_norm_eps)?;
270270

271+
// NOTE: https://github.com/huggingface/transformers/pull/39847
272+
let rope_theta = match config.rope_theta {
273+
Some(rope_theta) => rope_theta,
274+
None => match &config.rope_parameters {
275+
Some(rope_parameters) => rope_parameters.rope_theta,
276+
None => candle::bail!("Neither `rope_theta` nor `rope_parameters.rope_theta` are defined in the `config.json`"),
277+
},
278+
};
279+
271280
let inv_freqs = get_inv_freqs(
272281
layers[0].attention.attention_head_size,
273-
config.rope_theta,
282+
rope_theta,
274283
vb.device(),
275284
config.rope_scaling.as_ref(),
276285
)?;

‎backends/candle/src/models/flash_qwen2.rs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
use crate::flash_attn::flash_attn_varlen;
22
use crate::layers::{get_cos_sin, get_inv_freqs, index_select, HiddenAct, Linear, RMSNorm};
33
use crate::models::{Model, Qwen2Config};
4+
45
use candle::{DType, Device, IndexOp, Result, Tensor};
56
use candle_nn::{Embedding, Module, VarBuilder};
67
use candle_rotary::apply_rotary_inplace;
8+
79
use text_embeddings_backend_core::{Batch, ModelType, Pool};
810

911
struct Qwen2Attention {
@@ -285,9 +287,18 @@ impl FlashQwen2Model {
285287

286288
let norm = RMSNorm::load(vb.pp("norm"), config.hidden_size, config.rms_norm_eps)?;
287289

290+
// NOTE: https://github.com/huggingface/transformers/pull/39847
291+
let rope_theta = match config.rope_theta {
292+
Some(rope_theta) => rope_theta,
293+
None => match &config.rope_parameters {
294+
Some(rope_parameters) => rope_parameters.rope_theta,
295+
None => candle::bail!("Neither `rope_theta` nor `rope_parameters.rope_theta` are defined in the `config.json`"),
296+
},
297+
};
298+
288299
let inv_freqs = get_inv_freqs(
289300
layers[0].attention.attention_head_size,
290-
config.rope_theta,
301+
rope_theta,
291302
vb.device(),
292303
None,
293304
)?;

‎backends/candle/src/models/flash_qwen3.rs‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
use crate::flash_attn::flash_attn_varlen;
22
use crate::layers::{get_cos_sin, get_inv_freqs, index_select, HiddenAct, Linear, RMSNorm};
33
use crate::models::{Model, Qwen3Config};
4+
45
use candle::{DType, Device, IndexOp, Result, Tensor};
56
use candle_nn::{Embedding, Module, VarBuilder};
67
use candle_rotary::apply_rotary_inplace;
8+
79
use text_embeddings_backend_core::{Batch, ModelType, Pool};
810

911
struct Qwen3Attention {
@@ -353,9 +355,18 @@ impl FlashQwen3Model {
353355
None
354356
};
355357

358+
// NOTE: https://github.com/huggingface/transformers/pull/39847
359+
let rope_theta = match config.rope_theta {
360+
Some(rope_theta) => rope_theta,
361+
None => match &config.rope_parameters {
362+
Some(rope_parameters) => rope_parameters.rope_theta,
363+
None => candle::bail!("Neither `rope_theta` nor `rope_parameters.rope_theta` are defined in the `config.json`"),
364+
},
365+
};
366+
356367
let inv_freqs = get_inv_freqs(
357368
layers[0].attention.attention_head_size,
358-
config.rope_theta,
369+
rope_theta,
359370
vb.device(),
360371
None,
361372
)?;

‎backends/candle/src/models/gemma3.rs‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
use crate::layers::{
22
apply_rotary, get_cos_sin, get_cublas_lt_wrapper, get_inv_freqs, HiddenAct, Linear,
3+
RopeParameters,
34
};
45
use crate::models::Model;
56

@@ -23,9 +24,10 @@ pub struct Gemma3Config {
2324
pub query_pre_attn_scalar: usize,
2425
pub rms_norm_eps: f32,
2526
pub rope_local_base_freq: f32,
26-
pub rope_theta: f32,
27+
pub rope_theta: Option<f32>,
28+
pub rope_parameters: Option<RopeParameters>,
2729
pub sliding_window: Option<usize>,
28-
#[serde(rename(deserialize = "_sliding_window_pattern"))]
30+
#[serde(rename = "_sliding_window_pattern")]
2931
pub sliding_window_pattern: usize,
3032
pub vocab_size: usize,
3133
}
@@ -653,7 +655,16 @@ impl Gemma3Model {
653655
.head_dim
654656
.unwrap_or(config.hidden_size / config.num_attention_heads);
655657

656-
let inv_freqs = get_inv_freqs(rotary_dim, config.rope_theta, vb.device(), None)?;
658+
// NOTE: https://github.com/huggingface/transformers/pull/39847
659+
let rope_theta = match config.rope_theta {
660+
Some(rope_theta) => rope_theta,
661+
None => match &config.rope_parameters {
662+
Some(rope_parameters) => rope_parameters.rope_theta,
663+
None => candle::bail!("Neither `rope_theta` nor `rope_parameters.rope_theta` are defined in the `config.json`"),
664+
},
665+
};
666+
667+
let inv_freqs = get_inv_freqs(rotary_dim, rope_theta, vb.device(), None)?;
657668
let rotary_cache =
658669
get_cos_sin(config.max_position_embeddings, &inv_freqs, vb.dtype(), true)?;
659670

‎backends/candle/src/models/gte.rs‎

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
use crate::layers::{
22
apply_rotary, get_cos_sin, get_cublas_lt_wrapper, get_inv_freqs, HiddenAct, LayerNorm, Linear,
3-
RopeScaling,
3+
RopeParameters, RopeScaling,
44
};
55
use crate::models::{Model, PositionEmbeddingType};
6+
67
use candle::{DType, Device, IndexOp, Result, Tensor, D};
78
use candle_nn::{Embedding, Module, VarBuilder};
89
use serde::Deserialize;
910
use std::collections::HashMap;
11+
1012
use text_embeddings_backend_core::{Batch, ModelType, Pool};
1113

1214
#[derive(Debug, Clone, PartialEq, Deserialize)]
@@ -22,7 +24,8 @@ pub struct GTEConfig {
2224
pub layer_norm_type: String,
2325
pub layer_norm_eps: f32,
2426
pub position_embedding_type: PositionEmbeddingType,
25-
pub rope_theta: f32,
27+
pub rope_theta: Option<f32>,
28+
pub rope_parameters: Option<RopeParameters>,
2629
pub rope_scaling: Option<RopeScaling>,
2730
#[serde(default)]
2831
pub logn_attention_scale: bool,
@@ -412,10 +415,19 @@ impl GTEModel {
412415
Self::inner_load(vb.pp("new"), config)
413416
.or_else(|_| Self::inner_load(vb.clone(), config))?;
414417

418+
// NOTE: https://github.com/huggingface/transformers/pull/39847
419+
let rope_theta = match config.rope_theta {
420+
Some(rope_theta) => rope_theta,
421+
None => match &config.rope_parameters {
422+
Some(rope_parameters) => rope_parameters.rope_theta,
423+
None => candle::bail!("Neither `rope_theta` nor `rope_parameters.rope_theta` are defined in the `config.json`"),
424+
},
425+
};
426+
415427
let rotary_dim = encoder.layers[0].attention.attention_head_size;
416428
let inv_freqs = get_inv_freqs(
417429
rotary_dim,
418-
config.rope_theta,
430+
rope_theta,
419431
vb.device(),
420432
config.rope_scaling.as_ref(),
421433
)?;

0 commit comments

Comments
 (0)