Patch DistilBERT variants with different weight keys (#614)

Author: alvarobartt

backends/candle/src/models/distilbert.rs

@@ -475,6 +475,11 @@ impl DistilBertModel {
                     DistilBertEncoder::load(vb.pp("distilbert.transformer"), config),
                 ) {
                     (embeddings, encoder)
+                } else if let (Ok(embeddings), Ok(encoder)) = (
+                    DistilBertEmbeddings::load(vb.pp("embeddings"), config),
+                    DistilBertEncoder::load(vb.pp("transformer"), config),
+                ) {
+                    (embeddings, encoder)
                 } else {
                     return Err(err);
                 }

backends/candle/src/models/flash_distilbert.rs

@@ -213,6 +213,11 @@ impl FlashDistilBertModel {
                     DistilBertEncoder::load(vb.pp("distilbert.transformer"), config),
                 ) {
                     (embeddings, encoder)
+                } else if let (Ok(embeddings), Ok(encoder)) = (
+                    DistilBertEmbeddings::load(vb.pp("embeddings"), config),
+                    DistilBertEncoder::load(vb.pp("transformer"), config),
+                ) {
+                    (embeddings, encoder)
                 } else {
                     return Err(err);
                 }