[rust] Fixes test failure on GPU

extensions/tokenizers/rust/src/layers/cublaslt.rs

@@ -2,8 +2,8 @@
 pub use cudarc::cublaslt::Activation;
 
 use crate::layers::HiddenAct;
-use candle::{CpuStorage, Device, Layout, Result, Shape, Storage, Tensor};
 use candle::backend::BackendStorage;
+use candle::{CpuStorage, Device, Layout, Result, Shape, Storage, Tensor};
 
 #[cfg(feature = "cuda")]
 use candle::cuda_backend::WrapErr;

extensions/tokenizers/rust/src/layers/linear.rs

@@ -20,7 +20,12 @@ pub struct Linear {
 }
 
 impl Linear {
-    pub fn load(vb: VarBuilder, in_dim: usize, out_dim: usize, act: Option<HiddenAct>) -> Result<Self> {
+    pub fn load(
+        vb: VarBuilder,
+        in_dim: usize,
+        out_dim: usize,
+        act: Option<HiddenAct>,
+    ) -> Result<Self> {
         Ok(Self {
             weight: vb.get((out_dim, in_dim), "weight")?,
             bias: Some(vb.get(out_dim, "bias")?),

extensions/tokenizers/rust/src/layers/mod.rs

@@ -2,5 +2,4 @@
 mod cublaslt;
 mod linear;
 
-pub use cublaslt::get_cublas_lt_wrapper;
 pub use linear::{HiddenAct, Linear};

extensions/tokenizers/rust/src/lib.rs

@@ -16,8 +16,8 @@ mod ndarray;
 
 #[cfg(feature = "cuda")]
 mod compute_cap;
-mod models;
 mod layers;
+mod models;
 
 extern crate tokenizers as tk;
 
@@ -30,7 +30,6 @@ use jni::errors::Error;
 use jni::objects::{
     JClass, JLongArray, JMethodID, JObject, JObjectArray, JString, JValue, ReleaseMode,
 };
-#[cfg(feature = "cuda")]
 use jni::sys::JNI_FALSE;
 
 use jni::sys::{jboolean, jint, jlong, jobjectArray, jsize, jvalue, JNI_TRUE};
@@ -811,7 +810,7 @@ pub extern "system" fn Java_ai_djl_engine_rust_RustLibrary_isCudaAvailable<'loca
     }
     #[cfg(not(feature = "cuda"))]
     {
-        JNI_TRUE
+        JNI_FALSE
     }
 }
 

extensions/tokenizers/rust/src/models/bert.rs

@@ -1,5 +1,5 @@
-use crate::models::Model;
 use crate::layers::Linear;
+use crate::models::Model;
 use candle::{Device, Result, Tensor};
 use candle_nn::{embedding, Embedding, Module, VarBuilder};
 use candle_transformers::models::with_tracing::{layer_norm, LayerNorm};
@@ -359,7 +359,12 @@ struct BertIntermediate {
 
 impl BertIntermediate {
     fn load(vb: VarBuilder, config: &BertConfig) -> Result<Self> {
-        let dense = Linear::load(vb.pp("dense"), config.hidden_size, config.intermediate_size, None)?;
+        let dense = Linear::load(
+            vb.pp("dense"),
+            config.hidden_size,
+            config.intermediate_size,
+            None,
+        )?;
         Ok(Self {
             dense,
             intermediate_act: HiddenActLayer::new(config.hidden_act),
@@ -387,7 +392,12 @@ struct BertOutput {
 
 impl BertOutput {
     fn load(vb: VarBuilder, config: &BertConfig) -> Result<Self> {
-        let dense = Linear::load(vb.pp("dense"), config.intermediate_size, config.hidden_size, None)?;
+        let dense = Linear::load(
+            vb.pp("dense"),
+            config.intermediate_size,
+            config.hidden_size,
+            None,
+        )?;
         let layer_norm = layer_norm(
             config.hidden_size,
             config.layer_norm_eps,

extensions/tokenizers/rust/src/models/mod.rs

@@ -4,7 +4,7 @@ mod distilbert;
 use crate::ndarray::as_data_type;
 use crate::{cast_handle, drop_handle, to_handle, to_string_array};
 use bert::{BertConfig, BertModel};
-use candle::{Device, DType, Result, Tensor};
+use candle::{DType, Device, Result, Tensor};
 use candle_nn::VarBuilder;
 use distilbert::{DistilBertConfig, DistilBertModel};
 use jni::objects::{JLongArray, JObject, JString, ReleaseMode};

extensions/tokenizers/src/test/java/ai/djl/huggingface/tokenizers/TextEmbeddingTranslatorTest.java

@@ -12,6 +12,7 @@
  */
 package ai.djl.huggingface.tokenizers;
 
+import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.ModelException;
 import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
@@ -104,6 +105,7 @@ public void testTextEmbeddingTranslator()
                         .optArgument("block_shapes", "(1,7,384)")
                         .optArgument("block_names", "last_hidden_state")
                         .optEngine("PyTorch")
+                        .optDevice(Device.cpu())
                         .optArgument("tokenizer", "bert-base-uncased")
                         .optArgument("pooling", "max")
                         .optOption("hasParameter", "false")