Remove burn-import dep and upgrade to burn 0.15 (#52)

* Remove burn-import dep and upgrade to burn 0.15 * Add missing cubecl runtime feature flag
tracel-ai · Dec 12, 2024 · 3654542 · 3654542
1 parent 473929c
commit 3654542
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 6 deletions.
diff --git a/bert-burn/Cargo.toml b/bert-burn/Cargo.toml
@@ -6,7 +6,7 @@ version = "0.2.0"
 edition = "2021"
 
 [features]
-default = ["wgpu", "fusion"]
+default = []
 f16 = []
 ndarray = ["burn/ndarray"]
 tch-cpu = ["burn/tch"]
@@ -20,14 +20,14 @@ safetensors = ["candle-core/default"]
 
 [dependencies]
 # Burn
-burn = { version = "0.14", default-features = false, features = ["dataset", "std"] }
+burn = { version = "0.15", default-features = false, features = ["dataset", "std"] }
+cubecl-runtime = { version = "0.3.0", features = ["channel-mpsc"] } # missing feature flag when burn default-features are off
 candle-core = { version = "0.3" }
 # Tokenizer
 tokenizers = { version = "0.15.0", default-features = false, features = [
   "onig",
   "http",
 ] }
-burn-import = "0.13"
 derive-new = "0.6.0"
 hf-hub = { version = "0.3.2", features = ["tokio"] }
 

diff --git a/bert-burn/src/embedding.rs b/bert-burn/src/embedding.rs
@@ -56,7 +56,7 @@ impl BertEmbeddingsConfig {
 impl<B: Backend> BertEmbeddings<B> {
     pub fn forward(&self, item: BertInferenceBatch<B>) -> Tensor<B, 3, Float> {
         // Items batch contains the tokenized input and padding mask, each of dim: [batch_size, max_seq_length]
-        let input_shape = &item.tokens.shape();
+        let input_shape = item.tokens.shape();
         let input_ids = item.tokens;
 
         // Embed tokens
@@ -76,7 +76,9 @@ impl<B: Backend> BertEmbeddings<B> {
 
         let seq_length = input_shape.dims[1];
         let mut position_ids_tensor: Tensor<B, 2, Int> =
-            Tensor::arange(0..seq_length as i64, device).reshape([1, seq_length]);
+            Tensor::arange(0..seq_length as i64, device)
+                .reshape([1, seq_length])
+                .expand(input_shape.clone());
 
         if self.max_position_embeddings != 512 {
             // RoBERTa use a different scheme than BERT to create position indexes where padding tokens are given
@@ -87,7 +89,8 @@ impl<B: Backend> BertEmbeddings<B> {
                     ..(seq_length as i64) + (self.pad_token_idx as i64) + 1,
                 device,
             )
-            .reshape([1, seq_length]);
+            .reshape([1, seq_length])
+            .expand(input_shape);
             position_ids_tensor =
                 position_ids.mask_fill(item.mask_pad.clone(), self.pad_token_idx as i32);
         }