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

huggingface · Mar 3, 2025 · 37db86f · 37db86f
1 parent add3a71
commit 37db86f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/candle-transformers/src/models/modernbert.rs b/candle-transformers/src/models/modernbert.rs
@@ -315,7 +315,7 @@ pub struct ModernBert {
 }
 
 impl ModernBert {
-    fn load(vb: VarBuilder, config: &Config) -> Result<Self> {
+    pub fn load(vb: VarBuilder, config: &Config) -> Result<Self> {
         let word_embeddings = embedding(
             config.vocab_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 global_attention_mask =
             prepare_4d_attention_mask(mask, DType::F32, None)?.to_device(xs.device())?;