From 37db86ff79629f46d45ae3f4f2faddea0785e934 Mon Sep 17 00:00:00 2001
From: Andrew Wason <rectalogic@rectalogic.com>
Date: Mon, 3 Mar 2025 06:39:04 -0500
Subject: [PATCH] Allow ModernBert to be used to generate embeddings. (#2791)

---
 candle-transformers/src/models/modernbert.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/candle-transformers/src/models/modernbert.rs b/candle-transformers/src/models/modernbert.rs
index b0ba9b4695..268ebc3346 100644
--- 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())?;