[rust] Fix camembert and distilbert model loading

deepjavalibrary · Aug 15, 2024 · 5e96385 · 5e96385
1 parent e0519ac
commit 5e96385
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 14 deletions.
diff --git a/extensions/tokenizers/rust/src/models/camembert.rs b/extensions/tokenizers/rust/src/models/camembert.rs
@@ -480,14 +480,14 @@ pub struct CamembertModel {
 impl CamembertModel {
     pub fn load(vb: VarBuilder, config: &CamembertConfig) -> Result<Self> {
         let (embeddings, encoder) = match (
-            BertEmbeddings::load(vb.pp("embeddings"), config),
-            BertEncoder::load(vb.pp("encoder"), config),
+            BertEmbeddings::load(vb.pp("roberta.embeddings"), config),
+            BertEncoder::load(vb.pp("roberta.encoder"), config),
         ) {
             (Ok(embeddings), Ok(encoder)) => (embeddings, encoder),
             (Err(err), _) | (_, Err(err)) => {
                 if let (Ok(embeddings), Ok(encoder)) = (
-                    BertEmbeddings::load(vb.pp("camembert.embeddings".to_string()), config),
-                    BertEncoder::load(vb.pp("camembert.encoder".to_string()), config),
+                    BertEmbeddings::load(vb.pp("deberta.embeddings".to_string()), config),
+                    BertEncoder::load(vb.pp("deberta.encoder".to_string()), config),
                 ) {
                     (embeddings, encoder)
                 } else {

diff --git a/extensions/tokenizers/rust/src/models/distilbert.rs b/extensions/tokenizers/rust/src/models/distilbert.rs
@@ -340,17 +340,13 @@ impl DistilBertModel {
             Embeddings::load(vb.pp("embeddings"), config),
             Transformer::load(vb.pp("transformer"), config),
         ) {
-            (Ok(embeddings), Ok(encoder)) => (embeddings, encoder),
+            (Ok(embeddings), Ok(transformer)) => (embeddings, transformer),
             (Err(err), _) | (_, Err(err)) => {
-                if let Some(model_type) = &config.model_type {
-                    if let (Ok(embeddings), Ok(encoder)) = (
-                        Embeddings::load(vb.pp(&format!("{model_type}.embeddings")), config),
-                        Transformer::load(vb.pp(&format!("{model_type}.transformer")), config),
-                    ) {
-                        (embeddings, encoder)
-                    } else {
-                        return Err(err);
-                    }
+                if let (Ok(embeddings), Ok(transformer)) = (
+                    Embeddings::load(vb.pp("distilbert.embeddings".to_string()), config),
+                    Transformer::load(vb.pp("distilbert.transformer".to_string()), config),
+                ) {
+                    (embeddings, transformer)
                 } else {
                     return Err(err);
                 }