Repository: deepjavalibrary/djl
Pull request #3408 (2024-08-12 18:13:27)

[docs] Fixes broken links in markdown files.
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



Patch:
diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml
index b862f6752..64ccf132e 100644
--- a/docs/mkdocs.yml
+++ b/docs/mkdocs.yml
@@ -52,7 +52,7 @@ plugins:
   - mike
   - exclude:
       regex:
-        - '.*\.(symbol|params|pt|jar|gz|zip|txt|json|sh|bat|cmd|puml|g4|java|py|cc|h|pro|properties|xml|html|scala|sbt|kt|npy|npz|onnx|rs|rst|safetensors|tar|toml|trt|vue|scss|yml|yaml|Dockerfile|cmake|RepositoryFactory|EngineProvider|ZooProvider)$|LICENSE|NOTICE|.*gradle.*|.*/mvnw|.*/Dockerfile.*|.*(hdpi|\.git|\.mvn|docker|buildSrc|gradle|integration|testing|jacoco|tools|javadoc|META-INF|/webapp|apache-spark/notebook|visualization-vue|malicious-url-detector/docs|frontend|interactive-console)/.*'
+        - '.*\.(symbol|params|pt|jar|gz|zip|txt|json|sh|bat|cmd|puml|g4|java|py|cc|h|pro|properties|xml|html|scala|sbt|kt|npy|npz|onnx|rs|rst|safetensors|tar|toml|trt|vue|scss|yml|yaml|Dockerfile|cmake|RepositoryFactory|EngineProvider|ZooProvider)$|LICENSE|NOTICE|.*gradle.*|.*/mvnw|.*/Dockerfile.*|.*(hdpi|\.git|\.mvn|docker|buildSrc|gradle|integration|testing|jacoco|tools|javadoc|META-INF|/webapp|apache-spark/notebook|frontend)/.*'
 
 
 nav:
@@ -170,7 +170,6 @@ nav:
       - User Guides:
           - 'docs/serving/serving/docs/lmi/user_guides/README.md'
           - 'docs/serving/serving/docs/lmi/user_guides/starting-guide.md'
-          - 'docs/serving/serving/docs/lmi/user_guides/deepspeed_user_guide.md'
           - 'docs/serving/serving/docs/lmi/user_guides/lmi-dist_user_guide.md'
           - 'docs/serving/serving/docs/lmi/user_guides/vllm_user_guide.md'
           - 'docs/serving/serving/docs/lmi/user_guides/tnx_user_guide.md'
@@ -218,7 +217,6 @@ nav:
           - Doodledraw (PyTorch): 'docs/demos/android/pytorch_android/quickdraw_recognition/README.md'
           - Style Transfer (PyTorch): 'docs/demos/android/pytorch_android/style_transfer_cyclegan/README.md'
           - Face Detection (PyTorch): 'docs/demos/android/pytorch_android/face_detection/README.md'
-          - MXNet Android Template: 'docs/demos/android/mxnet-android/README.md'
       - EcoSystem:
           - Java Integrations:
               - DJL Component in Apache Camel: 'docs/demos/camel-djl/README.md'
diff --git a/examples/docs/mask_detection.md b/examples/docs/mask_detection.md
index 55bd56b4a..d1ea0465d 100644
--- a/examples/docs/mask_detection.md
+++ b/examples/docs/mask_detection.md
@@ -5,7 +5,7 @@ YOLOv5 is a powerful model for object detection tasks. With the transfer learnin
 In this example, we apply it on the [Face Mask Detection dataset](https://www.kaggle.com/datasets/andrewmvd/face-mask-detection?select=images). We first train the YOLOv5s model in Python, with the help of [ATLearn](https://github.com/awslabs/atlearn/blob/main/examples/docs/face_mask_detection.md), a python transfer learning toolkit.
 Then, the model is saved as an ONNX model, which is then imported into DJL for inference. We apply it on the mask wearing detection task. 
 
-The source code can be found at [MaskDetectionOnnx.java](../src/main/java/ai/djl/examples/inference/cv/MaskDetection.java)
+The source code can be found at [MaskDetectionOnnx.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/MaskDetection.java)
 
 ## The training part in ATLearn
 
@@ -13,7 +13,7 @@ We initially attempted to import a pretrained YOLOv5 into DJL, and fine-tune it
 
 The retraining of YOLOv5 can be found in an example in ATLearn: `examples/docs/face_mask_detection.md`. In this example, the YOLOv5 layers near the input are frozen while those near the output are fine-tuned with the customized data. This follows the transfer learning idea.
 
-In this example, the trained model is first exported to ONNX file, eg. `mask.onnx` and then imported in  [MaskDetectionOnnx.java](../src/main/java/ai/djl/examples/inference/cv/MaskDetection.java).
+In this example, the trained model is first exported to ONNX file, eg. `mask.onnx` and then imported in  [MaskDetectionOnnx.java](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/cv/MaskDetection.java).
 This ONNX model file can also be used for inference in python, which will serve as a benchmark. See the [tutorial doc](https://github.com/awslabs/atlearn/blob/main/examples/docs/face_mask_detection.md). 
 
 ## Setup guide
diff --git a/examples/docs/train_transfer_fresh_fruit.md b/examples/docs/train_transfer_fresh_fruit.md
index 2f1feb0cf..2a8dc9b12 100644
--- a/examples/docs/train_transfer_fresh_fruit.md
+++ b/examples/docs/train_transfer_fresh_fruit.md
@@ -26,7 +26,7 @@ The example is structured as follows:
 3. The demonstration of transfer learning in Java
 4. Experiment on the reduction of the training data size
 
-The full source code is available [here](../src/main/java/ai/djl/examples/training/transferlearning/TransferFreshFruit.java).
+The full source code is available [here](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/training/transferlearning/TransferFreshFruit.java).
 
 ## The data set and the problem formulation
 In this example, we demonstrate with the [fruit fresh and rotten dataset](https://www.kaggle.com/datasets/sriramr/fruits-fresh-and-rotten-for-classification),
diff --git a/extensions/timeseries/README.md b/extensions/timeseries/README.md
index 0cd38b8b4..e5798263e 100644
--- a/extensions/timeseries/README.md
+++ b/extensions/timeseries/README.md
@@ -116,7 +116,7 @@ The following pseudocode demonstrates how to create a `DeepARTranslator` with `a
 If you want to customize your own time series model translator, you can easily use the transform
 package for your data preprocess.
 
-See [examples](src/test/java/ai/djl/timeseries/translator/DeepARTranslatorTest.java) for more details.
+See [examples](https://github.com/deepjavalibrary/djl/blob/master/extensions/timeseries/src/test/java/ai/djl/timeseries/translator/DeepARTranslatorTest.java) for more details.
 
 ## Simple Example
 
@@ -223,7 +223,7 @@ Since djl doesn't support drawing yet, you can find our script for visualization
 [here](https://gist.github.com/Carkham/a5162c9298bc51fec648a458a3437008).
 
 The **full source code** for this example is available
-[here](../../examples/src/main/java/ai/djl/examples/inference/timeseries/AirPassengersDeepAR.java)
+[here](https://github.com/deepjavalibrary/djl/blob/master/examples/src/main/java/ai/djl/examples/inference/timeseries/AirPassengersDeepAR.java)
 
 ## Documentation
 
diff --git a/extensions/tokenizers/rust/src/models/bert.rs b/extensions/tokenizers/rust/src/models/bert.rs
index 00bcd27ad..0e7b1722f 100644
--- a/extensions/tokenizers/rust/src/models/bert.rs
+++ b/extensions/tokenizers/rust/src/models/bert.rs
@@ -487,8 +487,8 @@ impl BertClassificationHead {
             vb.pp("bert.pooler.dense"),
             config.hidden_size,
             config.hidden_size,
-            None)
-        {
+            None,
+        ) {
             Ok(layer) => Some(layer),
             Err(_) => None,
         };
@@ -555,7 +555,6 @@ impl BertModel {
 }
 
 impl Model for BertModel {
-
     fn get_input_names(&self) -> Vec<String> {
         return vec![
             "input_ids".to_string(),
@@ -590,9 +589,8 @@ pub struct BertForSequenceClassification {
 impl BertForSequenceClassification {
     pub fn load(vb: VarBuilder, config: &BertConfig) -> Result<Self> {
         let bert = Box::new(BertModel::load(vb.clone(), &config)?);
-        let classifier: Box<dyn ClassificationHead + Send> = Box::new(
-            BertClassificationHead::load(vb.pp("classifier"), config)?,
-        );
+        let classifier: Box<dyn ClassificationHead + Send> =
+            Box::new(BertClassificationHead::load(vb.pp("classifier"), config)?);
         Ok(Self {
             bert,
             classifier,
@@ -603,7 +601,6 @@ impl BertForSequenceClassification {
 }
 
 impl Model for BertForSequenceClassification {
-
     fn get_input_names(&self) -> Vec<String> {
         return vec![
             "input_ids".to_string(),
@@ -626,4 +623,4 @@ impl Model for BertForSequenceClassification {
         let logits = self.classifier.forward(&sequence_output)?;
         Ok(logits)
     }
-}
\ No newline at end of file
+}
diff --git a/extensions/tokenizers/rust/src/models/camembert.rs b/extensions/tokenizers/rust/src/models/camembert.rs
index 0aa119914..da42a4f29 100644
--- a/extensions/tokenizers/rust/src/models/camembert.rs
+++ b/extensions/tokenizers/rust/src/models/camembert.rs
@@ -166,7 +166,8 @@ impl BertEmbeddings {
         let mut embeddings = (&input_embeddings + token_type_embeddings)?;
         if let Some(position_embeddings) = &self.position_embeddings {
             let position_offset = self.config.pad_token_id as u32 + 1;
-            let position_ids = (position_offset..(seq_len as u32 + position_offset)).collect::<Vec<_>>();
+            let position_ids =
+                (position_offset..(seq_len as u32 + position_offset)).collect::<Vec<_>>();
             let position_ids = Tensor::new(&position_ids[..], input_ids.device())?;
             embeddings = embeddings.broadcast_add(&position_embeddings.forward(&position_ids)?)?
         }
@@ -504,7 +505,6 @@ impl CamembertModel {
 }
 
 impl Model for CamembertModel {
-
     fn get_input_names(&self) -> Vec<String> {
         return vec![
             "input_ids".to_string(),
@@ -526,4 +526,4 @@ impl Model for CamembertModel {
         let sequence_output = self.encoder.forward(&embedding_output)?;
         Ok(sequence_output)
     }
-}
\ No newline at end of file
+}
diff --git a/extensions/tokenizers/rust/src/models/distilbert.rs b/extensions/tokenizers/rust/src/models/distilbert.rs
index fabf11c01..345be40d1 100644
--- a/extensions/tokenizers/rust/src/models/distilbert.rs
+++ b/extensions/tokenizers/rust/src/models/distilbert.rs
@@ -366,7 +366,6 @@ impl DistilBertModel {
 }
 
 impl Model for DistilBertModel {
-
     fn get_input_names(&self) -> Vec<String> {
         return vec!["input_ids".to_string(), "attention_mask".to_string()];
     }
diff --git a/extensions/tokenizers/rust/src/models/mod.rs b/extensions/tokenizers/rust/src/models/mod.rs
index 707cfc58f..c11ec585a 100644
--- a/extensions/tokenizers/rust/src/models/mod.rs
+++ b/extensions/tokenizers/rust/src/models/mod.rs
@@ -6,18 +6,18 @@ mod xlm_roberta;
 
 use crate::ndarray::as_data_type;
 use crate::{cast_handle, drop_handle, to_handle, to_string_array};
-use candle::{DType, Device, Result, Tensor, Error};
-use candle_nn::VarBuilder;
-use bert::{BertConfig, BertModel, BertForSequenceClassification};
+use bert::{BertConfig, BertForSequenceClassification, BertModel};
 use camembert::{CamembertConfig, CamembertModel};
+use candle::{DType, Device, Error, Result, Tensor};
+use candle_nn::VarBuilder;
 use distilbert::{DistilBertConfig, DistilBertModel};
-use roberta::{RobertaConfig, RobertaModel, RobertaForSequenceClassification};
-use xlm_roberta::{XLMRobertaConfig, XLMRobertaModel, XLMRobertaForSequenceClassification};
 use jni::objects::{JLongArray, JObject, JString, ReleaseMode};
 use jni::sys::{jint, jlong, jobjectArray};
 use jni::JNIEnv;
+use roberta::{RobertaConfig, RobertaForSequenceClassification, RobertaModel};
 use serde::Deserialize;
 use std::path::PathBuf;
+use xlm_roberta::{XLMRobertaConfig, XLMRobertaForSequenceClassification, XLMRobertaModel};
 
 #[derive(Debug, PartialEq, Clone)]
 #[allow(dead_code, unused)]
@@ -40,7 +40,6 @@ enum Config {
 }
 
 pub(crate) trait Model {
-
     fn get_input_names(&self) -> Vec<String>;
 
     fn forward(
@@ -109,7 +108,9 @@ fn load_model<'local>(
             config.use_flash_attn = Some(use_flash_attn);
             match config.architectures.first() {
                 Some(arch) => match arch.as_str() {
-                    "BertForSequenceClassification" => Ok(Box::new(BertForSequenceClassification::load(vb, &config)?)),
+                    "BertForSequenceClassification" => {
+                        Ok(Box::new(BertForSequenceClassification::load(vb, &config)?))
+                    }
                     _ => Ok(Box::new(BertModel::load(vb, &config)?)),
                 },
                 None => Ok(Box::new(BertModel::load(vb, &config)?)),
@@ -125,7 +126,9 @@ fn load_model<'local>(
             config.use_flash_attn = Some(use_flash_attn);
             match config.architectures.first() {
                 Some(arch) => match arch.as_str() {
-                    "RobertaForSequenceClassification" => Ok(Box::new(RobertaForSequenceClassification::load(vb, &config)?)),
+                    "RobertaForSequenceClassification" => Ok(Box::new(
+                        RobertaForSequenceClassification::load(vb, &config)?,
+                    )),
                     _ => Ok(Box::new(RobertaModel::load(vb, &config)?)),
                 },
                 None => Ok(Box::new(RobertaModel::load(vb, &config)?)),
@@ -136,7 +139,9 @@ fn load_model<'local>(
             config.use_flash_attn = Some(use_flash_attn);
             match config.architectures.first() {
                 Some(arch) => match arch.as_str() {
-                    "XLMRobertaForSequenceClassification" => Ok(Box::new(XLMRobertaForSequenceClassification::load(vb, &config)?)),
+                    "XLMRobertaForSequenceClassification" => Ok(Box::new(
+                        XLMRobertaForSequenceClassification::load(vb, &config)?,
+                    )),
                     _ => Ok(Box::new(XLMRobertaModel::load(vb, &config)?)),
                 },
                 None => Ok(Box::new(XLMRobertaModel::load(vb, &config)?)),
diff --git a/extensions/tokenizers/rust/src/models/roberta.rs b/extensions/tokenizers/rust/src/models/roberta.rs
index 4198927f4..27fd65995 100644
--- a/extensions/tokenizers/rust/src/models/roberta.rs
+++ b/extensions/tokenizers/rust/src/models/roberta.rs
@@ -166,7 +166,8 @@ impl BertEmbeddings {
         let mut embeddings = (&input_embeddings + token_type_embeddings)?;
         if let Some(position_embeddings) = &self.position_embeddings {
             let position_offset = self.config.pad_token_id as u32 + 1;
-            let position_ids = (position_offset..(seq_len as u32 + position_offset)).collect::<Vec<_>>();
+            let position_ids =
+                (position_offset..(seq_len as u32 + position_offset)).collect::<Vec<_>>();
             let position_ids = Tensor::new(&position_ids[..], input_ids.device())?;
             embeddings = embeddings.broadcast_add(&position_embeddings.forward(&position_ids)?)?
         }
@@ -484,7 +485,8 @@ impl RobertaClassificationHead {
             None => candle::bail!("`id2label` must be set for classifier models"),
             Some(id2label) => id2label.len(),
         };
-        let intermediate = Linear::load(vb.pp("dense"), config.hidden_size, config.hidden_size, None)?;
+        let intermediate =
+            Linear::load(vb.pp("dense"), config.hidden_size, config.hidden_size, None)?;
         let output = Linear::load(vb.pp("out_proj"), config.hidden_size, n_classes, None)?;
         Ok(Self {
             intermediate,
@@ -544,7 +546,6 @@ impl RobertaModel {
 }
 
 impl Model for RobertaModel {
-
     fn get_input_names(&self) -> Vec<String> {
         return vec![
             "input_ids".to_string(),
@@ -592,7 +593,6 @@ impl RobertaForSequenceClassification {
 }
 
 impl Model for RobertaForSequenceClassification {
-
     fn get_input_names(&self) -> Vec<String> {
         return vec![
             "input_ids".to_string(),
@@ -615,4 +615,4 @@ impl Model for RobertaForSequenceClassification {
         let logits = self.classifier.forward(&sequence_output)?;
         Ok(logits)
     }
-}
\ No newline at end of file
+}
diff --git a/extensions/tokenizers/rust/src/models/xlm_roberta.rs b/extensions/tokenizers/rust/src/models/xlm_roberta.rs
index f3b6ebc9d..4d7eef4ee 100644
--- a/extensions/tokenizers/rust/src/models/xlm_roberta.rs
+++ b/extensions/tokenizers/rust/src/models/xlm_roberta.rs
@@ -166,7 +166,8 @@ impl BertEmbeddings {
         let mut embeddings = (&input_embeddings + token_type_embeddings)?;
         if let Some(position_embeddings) = &self.position_embeddings {
             let position_offset = self.config.pad_token_id as u32 + 1;
-            let position_ids = (position_offset..(seq_len as u32 + position_offset)).collect::<Vec<_>>();
+            let position_ids =
+                (position_offset..(seq_len as u32 + position_offset)).collect::<Vec<_>>();
             let position_ids = Tensor::new(&position_ids[..], input_ids.device())?;
             embeddings = embeddings.broadcast_add(&position_embeddings.forward(&position_ids)?)?
         }
@@ -484,7 +485,8 @@ impl XLMRobertaClassificationHead {
             None => candle::bail!("`id2label` must be set for classifier models"),
             Some(id2label) => id2label.len(),
         };
-        let intermediate = Linear::load(vb.pp("dense"), config.hidden_size, config.hidden_size, None)?;
+        let intermediate =
+            Linear::load(vb.pp("dense"), config.hidden_size, config.hidden_size, None)?;
         let output = Linear::load(vb.pp("out_proj"), config.hidden_size, n_classes, None)?;
         Ok(Self {
             intermediate,
@@ -549,7 +551,6 @@ impl XLMRobertaModel {
 }
 
 impl Model for XLMRobertaModel {
-
     fn get_input_names(&self) -> Vec<String> {
         return vec![
             "input_ids".to_string(),
@@ -597,7 +598,6 @@ impl XLMRobertaForSequenceClassification {
 }
 
 impl Model for XLMRobertaForSequenceClassification {
-
     fn get_input_names(&self) -> Vec<String> {
         return vec![
             "input_ids".to_string(),
@@ -620,4 +620,4 @@ impl Model for XLMRobertaForSequenceClassification {
         let logits = self.classifier.forward(&sequence_output)?;
         Ok(logits)
     }
-}
\ No newline at end of file
+}
diff --git a/extensions/tokenizers/rust/src/ndarray/mod.rs b/extensions/tokenizers/rust/src/ndarray/mod.rs
index f2bd7f56c..781c459ec 100644
--- a/extensions/tokenizers/rust/src/ndarray/mod.rs
+++ b/extensions/tokenizers/rust/src/ndarray/mod.rs
@@ -9,6 +9,7 @@ use crate::{cast_handle, drop_handle, to_handle};
 mod binary;
 mod cmp;
 mod creation;
+mod nn;
 mod other;
 mod reduce;
 mod unary;
diff --git a/extensions/tokenizers/rust/src/ndarray/nn.rs b/extensions/tokenizers/rust/src/ndarray/nn.rs
new file mode 100644
index 000000000..f577310a3
--- /dev/null
+++ b/extensions/tokenizers/rust/src/ndarray/nn.rs
@@ -0,0 +1,53 @@
+use crate::cast_handle;
+use crate::ndarray::return_handle;
+use candle::Tensor;
+use jni::objects::JObject;
+use jni::sys::{jfloat, jint, jlong};
+use jni::JNIEnv;
+
+#[no_mangle]
+pub extern "system" fn Java_ai_djl_engine_rust_RustLibrary_sigmoid<'local>(
+    mut env: JNIEnv,
+    _: JObject,
+    handle: jlong,
+) -> jlong {
+    let tensor = cast_handle::<Tensor>(handle);
+    let ret = candle_nn::ops::sigmoid(&tensor);
+    return_handle(&mut env, ret)
+}
+
+#[no_mangle]
+pub extern "system" fn Java_ai_djl_engine_rust_RustLibrary_leakyRelu<'local>(
+    mut env: JNIEnv,
+    _: JObject,
+    handle: jlong,
+    alpha: jfloat,
+) -> jlong {
+    let tensor = cast_handle::<Tensor>(handle);
+    let ret = candle_nn::ops::leaky_relu(&tensor, alpha as f64);
+    return_handle(&mut env, ret)
+}
+
+#[no_mangle]
+pub extern "system" fn Java_ai_djl_engine_rust_RustLibrary_softmax<'local>(
+    mut env: JNIEnv,
+    _: JObject,
+    handle: jlong,
+    axis: jint,
+) -> jlong {
+    let tensor = cast_handle::<Tensor>(handle);
+    let ret = candle_nn::ops::softmax(&tensor, axis as usize);
+    return_handle(&mut env, ret)
+}
+
+#[no_mangle]
+pub extern "system" fn Java_ai_djl_engine_rust_RustLibrary_logSoftmax<'local>(
+    mut env: JNIEnv,
+    _: JObject,
+    handle: jlong,
+    axis: jint,
+) -> jlong {
+    let tensor = cast_handle::<Tensor>(handle);
+    let ret = candle_nn::ops::log_softmax(&tensor, axis as usize);
+    return_handle(&mut env, ret)
+}
diff --git a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java
index ba526b098..97e1148bf 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RsNDArray.java
@@ -1229,6 +1229,9 @@ public class RsNDArray extends NativeResource<Long> implements NDArray {
     /** {@inheritDoc} */
     @Override
     public RsNDArray softmax(int axis) {
+        if (getShape().isScalar() || shape.size() == 0) {
+            return (RsNDArray) duplicate();
+        }
         return toArray(RustLibrary.softmax(getHandle(), axis));
     }
 
diff --git a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RustLibrary.java b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RustLibrary.java
index 131e1f7ca..852c392f1 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RustLibrary.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/engine/rust/RustLibrary.java
@@ -322,13 +322,9 @@ public final class RustLibrary {
         throw new UnsupportedOperationException("Not implemented");
     }
 
-    public static long softmax(long handle, int axis) {
-        throw new UnsupportedOperationException("Not implemented");
-    }
+    public static native long softmax(long handle, int axis);
 
-    public static long logSoftmax(long handle, int axis) {
-        throw new UnsupportedOperationException("Not implemented");
-    }
+    public static native long logSoftmax(long handle, int axis);
 
     public static native long cumSum(long handle, int axis);
 
@@ -400,9 +396,7 @@ public final class RustLibrary {
         throw new UnsupportedOperationException("Not implemented");
     }
 
-    public static long sigmoid(long handle) {
-        throw new UnsupportedOperationException("Not implemented");
-    }
+    public static native long sigmoid(long handle);
 
     public static long softPlus(long handle) {
         throw new UnsupportedOperationException("Not implemented");
@@ -412,9 +406,7 @@ public final class RustLibrary {
         throw new UnsupportedOperationException("Not implemented");
     }
 
-    public static long leakyRelu(long handle, float alpha) {
-        throw new UnsupportedOperationException("Not implemented");
-    }
+    public static native long leakyRelu(long handle, float alpha);
 
     public static long elu(long handle, float alpha) {
         throw new UnsupportedOperationException("Not implemented");
diff --git a/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py b/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py
index 6776f8d1f..c616d9948 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py
@@ -127,7 +127,7 @@ class HuggingfaceConverter:
             os.makedirs(temp_dir)
 
         tokenizer = AutoTokenizer.from_pretrained(model_id)
-        include_types = "token_type_ids" in tokenizer.model_input_names
+        include_types = config.model_type != "distilbert"
         hf_pipeline = PipelineHolder(tokenizer, ModelHolder(config))
         try:
             # Save tokenizer.json to temp dir
diff --git a/extensions/tokenizers/src/main/python/djl_converter/model_converter.py b/extensions/tokenizers/src/main/python/djl_converter/model_converter.py
index 2d0f99827..a9f7b6d90 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/model_converter.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/model_converter.py
@@ -43,7 +43,7 @@ def main():
 
     if os.listdir(output_dir):
         logging.error(f"output directory: {output_dir} is not empty.")
-        return False
+        return 1
 
     if os.path.exists(args.model_id):
         logging.info(f"converting local model: {args.model_id}")
@@ -66,7 +66,7 @@ def main():
     if not task:
         logging.error(
             f"Unsupported model architecture: {arch} for {args.model_id}.")
-        return False
+        return 1
 
     converter = SUPPORTED_TASKS[task]
 
@@ -75,16 +75,15 @@ def main():
                                                  output_dir, False)
         if result:
             logging.info(f"Convert model {model_info.modelId} finished.")
-            return True
+            return 0
         else:
             logging.error(f"{model_info.modelId}: {reason}")
     except Exception as e:
         logging.warning(f"Failed to convert model: {model_info.modelId}.")
         logging.warning(e, exc_info=True)
 
-    return False
+    return 1
 
 
 if __name__ == "__main__":
-    if not main():
-        sys.exit(1)
+    sys.exit(main())


Test patch:
diff --git a/extensions/tokenizers/src/test/java/ai/djl/engine/rust/NDArrayTests.java b/extensions/tokenizers/src/test/java/ai/djl/engine/rust/NDArrayTests.java
index 8f9523c7f..fd56c436d 100644
--- a/extensions/tokenizers/src/test/java/ai/djl/engine/rust/NDArrayTests.java
+++ b/extensions/tokenizers/src/test/java/ai/djl/engine/rust/NDArrayTests.java
@@ -18,6 +18,7 @@ import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Activation;
 import ai.djl.testing.Assertions;
 import ai.djl.testing.TestRequirements;
 
@@ -215,6 +216,70 @@ public class NDArrayTests {
         }
     }
 
+    @Test
+    public void testSoftmax() {
+        try (NDManager manager = NDManager.newBaseManager("Rust")) {
+            NDArray array = manager.ones(new Shape(10));
+            NDArray expected = manager.zeros(new Shape(10)).add(0.1f);
+            Assertions.assertAlmostEquals(array.softmax(0), expected);
+            // test multi-dim
+            array = manager.ones(new Shape(2, 3, 1, 3));
+            expected = manager.zeros(new Shape(2, 3, 1, 3)).add(0.5f);
+            Assertions.assertAlmostEquals(array.softmax(0), expected);
+            expected = manager.zeros(new Shape(2, 3, 1, 3)).add(0.33333334f);
+            Assertions.assertAlmostEquals(array.softmax(1), expected);
+            expected = manager.ones(new Shape(2, 3, 1, 3));
+            Assertions.assertAlmostEquals(array.softmax(2), expected);
+            expected = manager.zeros(new Shape(2, 3, 1, 3)).add(0.33333334f);
+            Assertions.assertAlmostEquals(array.softmax(3), expected);
+            // test scalar
+            array = manager.create(1f);
+            Assertions.assertAlmostEquals(array.softmax(0), array);
+            // test zero
+            array = manager.create(new Shape(2, 0, 1));
+            Assertions.assertAlmostEquals(array.softmax(0), array);
+        }
+    }
+
+    @Test
+    public void testLogSoftmax() {
+        try (NDManager manager = NDManager.newBaseManager("Rust")) {
+            NDArray array = manager.ones(new Shape(10));
+            NDArray expected = manager.zeros(new Shape(10)).add(-2.3025851f);
+            Assertions.assertAlmostEquals(array.logSoftmax(0), expected);
+            // test multi-dim
+            array = manager.ones(new Shape(2, 3, 1, 3));
+            expected = manager.zeros(new Shape(2, 3, 1, 3)).add(-0.6931472f);
+            Assertions.assertAlmostEquals(array.logSoftmax(0), expected);
+            expected = manager.zeros(new Shape(2, 3, 1, 3)).add(-1.0986123f);
+            Assertions.assertAlmostEquals(array.logSoftmax(1), expected);
+            // test scalar
+            array = manager.create(1f);
+            Assertions.assertAlmostEquals(array.softmax(0), array);
+            // test zero
+            array = manager.create(new Shape(2, 0, 1));
+            Assertions.assertAlmostEquals(array.softmax(0), array);
+        }
+    }
+
+    @Test
+    public void testSigmoid() {
+        try (NDManager manager = NDManager.newBaseManager("Rust")) {
+            NDArray data = manager.create(new float[] {0});
+            NDArray expected = manager.create(new float[] {0.5f});
+            Assertions.assertAlmostEquals(Activation.sigmoid(data), expected);
+        }
+    }
+
+    @Test
+    public void testLeakyRelu() {
+        try (NDManager manager = NDManager.newBaseManager("Rust")) {
+            NDArray data = manager.create(new float[] {-1, 0, 2});
+            NDArray expected = manager.create(new float[] {-1, 0, 2});
+            Assert.assertEquals(Activation.leakyRelu(data, 1.0f), expected);
+        }
+    }
+
     @Test
     public void testExpandDim() {
         try (NDManager manager = NDManager.newBaseManager("Rust")) {