huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 18:48:51 +00:00
Code Issues Packages Projects Releases Wiki Activity

Allow ModernBert to be used to generate embeddings. (#2791)

Browse Source
This commit is contained in:
Andrew Wason
2025-03-03 06:39:04 -05:00
committed by GitHub
parent add3a714aa
commit 37db86ff79
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
candle-transformers/src/models/modernbert.rs
View File

@ -315,7 +315,7 @@ pub struct ModernBert {
} }
impl ModernBert { impl ModernBert {
fn load(vb: VarBuilder, config: &Config) -> Result<Self> { pub fn load(vb: VarBuilder, config: &Config) -> Result<Self> {
let word_embeddings = embedding( let word_embeddings = embedding(
config.vocab_size, config.vocab_size,
config.hidden_size, config.hidden_size,
@ -371,7 +371,7 @@ impl ModernBert {
}) })
} }
fn forward(&self, xs: &Tensor, mask: &Tensor) -> Result<Tensor> { pub fn forward(&self, xs: &Tensor, mask: &Tensor) -> Result<Tensor> {
let seq_len = xs.shape().dims()[1]; let seq_len = xs.shape().dims()[1];
let global_attention_mask = let global_attention_mask =
prepare_4d_attention_mask(mask, DType::F32, None)?.to_device(xs.device())?; prepare_4d_attention_mask(mask, DType::F32, None)?.to_device(xs.device())?;