Repository: deepjavalibrary/djl
Pull request #3177 (2024-05-11 19:50:31)

[examples] Fixes TextGeneration EOS bug
## 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/.github/workflows/continuous.yml b/.github/workflows/continuous.yml
index 2abbcd5da..3b91b8168 100644
--- a/.github/workflows/continuous.yml
+++ b/.github/workflows/continuous.yml
@@ -93,7 +93,7 @@ jobs:
         uses: actions/upload-artifact@v3
         if: always()
         with:
-          name: reports
+          name: reports-${{ matrix.operating-system }}
           path: |
             ${{ github.workspace }}/**/build/reports/**/*
             !${{ github.workspace }}/**/build/reports/jacoco/*
diff --git a/.github/workflows/native_jni_s3_pytorch.yml b/.github/workflows/native_jni_s3_pytorch.yml
index b31bf97ea..18aab7584 100644
--- a/.github/workflows/native_jni_s3_pytorch.yml
+++ b/.github/workflows/native_jni_s3_pytorch.yml
@@ -6,6 +6,10 @@ on:
       pt_version:
         description: 'pytorch version'
         required: false
+      cuda:
+        description: 'CUDA version'
+        required: true
+        default: 'cu121'
   schedule:
     - cron: '0 5 * * *'
 
@@ -83,7 +87,8 @@ jobs:
           ./gradlew -Pjni -Ppt_version=$PYTORCH_VERSION :integration:test "-Dai.djl.default_engine=PyTorch"
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
           export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
-          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcuda=cu121 -Ppt_version=$PYTORCH_VERSION
+          CUDA_VERSION=${{ github.event.inputs.cuda }}
+          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcuda=$CUDA_VERSION -Ppt_version=$PYTORCH_VERSION
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v3
@@ -134,7 +139,8 @@ jobs:
           ./gradlew :engines:pytorch:pytorch-native:cleanJNI
           rm -rf ~/.djl.ai
           export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.9 9.0"
-          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcuda=cu121 -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
+          CUDA_VERSION=${{ github.event.inputs.cuda }}
+          ./gradlew :engines:pytorch:pytorch-native:compileJNI -Pcuda=$CUDA_VERSION -Pprecxx11 -Ppt_version=$PYTORCH_VERSION
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v2
         with:
@@ -185,7 +191,7 @@ jobs:
           call "C:\Program Files\Microsoft Visual Studio\2022\Enterprise\VC\Auxiliary\Build\vcvarsall.bat" amd64
           set "CUDA_PATH=%ProgramFiles%\NVIDIA GPU Computing Toolkit\CUDA\v11.7"
           set "PATH=%CUDA_PATH%\bin;%CUDA_PATH%\libnvvp;%PATH%"
-          gradlew :engines:pytorch:pytorch-native:cleanJNI :engines:pytorch:pytorch-native:compileJNI -Pcuda=cu121 -Ppt_version=${{ github.event.inputs.pt_version }}
+          gradlew :engines:pytorch:pytorch-native:cleanJNI :engines:pytorch:pytorch-native:compileJNI -Pcuda=${{ github.event.inputs.cuda }} -Ppt_version=${{ github.event.inputs.pt_version }}
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v4
         with:
diff --git a/.github/workflows/serving_publish.yml b/.github/workflows/serving_publish.yml
index 380666487..73fe45a56 100644
--- a/.github/workflows/serving_publish.yml
+++ b/.github/workflows/serving_publish.yml
@@ -79,13 +79,13 @@ jobs:
         run: |
           ./gradlew :awscurl:jar
           aws s3 cp awscurl/build/awscurl s3://djl-ai/publish/awscurl/
-          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/awscurl/*"
+          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/awscurl/awscurl"
       - name: Copy awscurl staging artifacts to S3
         if: ${{ github.event.inputs.mode == 'staging' }}
         run: |
           ./gradlew :awscurl:jar
           DJL_VERSION=$(cat gradle.properties | awk -F '=' '/djl_version/ {print $2}')
-          aws s3 cp awscurl/build/awscurl s3://djl-ai/publish/${DJL_VERSION}/awscurl/
+          aws s3 cp awscurl/build/awscurl s3://djl-ai/publish/awscurl/${DJL_VERSION}/
           aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/awscurl/${DJL_VERSION}/*"
       - name: Publish to snapshot repository
         if: ${{ github.event.inputs.mode == '' || github.event.inputs.mode == 'snapshot' }}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java b/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java
index 7590244e4..7664b4f7e 100644
--- a/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java
@@ -118,9 +118,12 @@ public class TextGenerator {
 
             // Termination Criteria
             long[] outputIdsArray = searchState.getNextInputIds().toLongArray();
-            for (int i = 0; i < outputIdsArray.length; i++) {
-                if (outputIdsArray[i] == config.getEosTokenId()) {
-                    endPosition[i] = searchState.getPastOutputIds().getShape().get(1) + 1;
+            for (int i = 0; i < endPosition.length; ++i) {
+                for (long tokenId : outputIdsArray) {
+                    if (tokenId == config.getEosTokenId()) {
+                        endPosition[i] = searchState.getPastOutputIds().getShape().get(1) + 1;
+                        break;
+                    }
                 }
             }
             if (searchState.getPastOutputIds().getShape().get(1) + 1 >= config.getMaxSeqLength()) {
@@ -240,9 +243,12 @@ public class TextGenerator {
 
             // Termination Criteria
             long[] outputIdsArray = searchState.getNextInputIds().toLongArray();
-            for (int i = 0; i < outputIdsArray.length; i++) {
-                if (outputIdsArray[i] == config.getEosTokenId()) {
-                    endPosition[i] = searchState.getPastOutputIds().getShape().get(1) + 1;
+            for (int i = 0; i < endPosition.length; ++i) {
+                for (long tokenId : outputIdsArray) {
+                    if (tokenId == config.getEosTokenId()) {
+                        endPosition[i] = searchState.getPastOutputIds().getShape().get(1) + 1;
+                        break;
+                    }
                 }
             }
             if (searchState.getPastOutputIds().getShape().getLastDimension() + 1
@@ -366,9 +372,12 @@ public class TextGenerator {
 
             // Termination Criteria
             long[] outputIdsArray = searchState.getPastOutputIds().toLongArray();
-            for (int i = 0; i < outputIdsArray.length; i++) {
-                if (outputIdsArray[i] == config.getEosTokenId()) {
-                    endPosition[i] = searchState.getPastOutputIds().getShape().get(1);
+            for (int i = 0; i < endPosition.length; ++i) {
+                for (long tokenId : outputIdsArray) {
+                    if (tokenId == config.getEosTokenId()) {
+                        endPosition[i] = searchState.getPastOutputIds().getShape().get(1);
+                        break;
+                    }
                 }
             }
             if (searchState.getPastOutputIds().getShape().get(1) >= config.getMaxSeqLength()) {
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index f2b993fc4..abbe7df97 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -225,6 +225,25 @@ public interface NDArray extends NDResource, BytesSupplier {
         return getShape().size();
     }
 
+    /** {@inheritDoc} */
+    @Override
+    default ByteBuffer toByteBuffer() {
+        return toByteBuffer(false);
+    }
+
+    /**
+     * Returns the {@code ByteBuffer} presentation of the object.
+     *
+     * <p>If returned ByteBuffer is a DirectByteBuffer, it shared the same native memory as the
+     * NDArray. The native memory will be deleted when NDArray is closed.
+     *
+     * <p>Not all the engine support return DirectByteBuffer.
+     *
+     * @param tryDirect use DirectBuffer if possible
+     * @return the {@code ByteBuffer} presentation of the object
+     */
+    ByteBuffer toByteBuffer(boolean tryDirect);
+
     /**
      * Converts this {@code NDArray} to a double array.
      *
@@ -236,7 +255,7 @@ public interface NDArray extends NDResource, BytesSupplier {
             throw new IllegalStateException(
                     "DataType mismatch, Required double" + " Actual " + getDataType());
         }
-        DoubleBuffer db = toByteBuffer().asDoubleBuffer();
+        DoubleBuffer db = toByteBuffer(true).asDoubleBuffer();
         double[] ret = new double[db.remaining()];
         db.get(ret);
         return ret;
@@ -255,7 +274,7 @@ public interface NDArray extends NDResource, BytesSupplier {
             throw new IllegalStateException(
                     "DataType mismatch, Required float, Actual " + getDataType());
         }
-        FloatBuffer fb = toByteBuffer().asFloatBuffer();
+        FloatBuffer fb = toByteBuffer(true).asFloatBuffer();
         float[] ret = new float[fb.remaining()];
         fb.get(ret);
         return ret;
@@ -272,7 +291,7 @@ public interface NDArray extends NDResource, BytesSupplier {
             throw new IllegalStateException(
                     "DataType mismatch, Required int" + " Actual " + getDataType());
         }
-        ShortBuffer ib = toByteBuffer().asShortBuffer();
+        ShortBuffer ib = toByteBuffer(true).asShortBuffer();
         short[] ret = new short[ib.remaining()];
         ib.get(ret);
         return ret;
@@ -289,7 +308,7 @@ public interface NDArray extends NDResource, BytesSupplier {
             throw new IllegalStateException(
                     "DataType mismatch, Required int" + " Actual " + getDataType());
         }
-        ShortBuffer ib = toByteBuffer().asShortBuffer();
+        ShortBuffer ib = toByteBuffer(true).asShortBuffer();
         int[] ret = new int[ib.remaining()];
         for (int i = 0; i < ret.length; ++i) {
             ret[i] = ib.get() & 0xffff;
@@ -309,7 +328,7 @@ public interface NDArray extends NDResource, BytesSupplier {
             throw new IllegalStateException(
                     "DataType mismatch, Required int" + " Actual " + getDataType());
         }
-        IntBuffer ib = toByteBuffer().asIntBuffer();
+        IntBuffer ib = toByteBuffer(true).asIntBuffer();
         int[] ret = new int[ib.remaining()];
         ib.get(ret);
         return ret;
@@ -326,7 +345,7 @@ public interface NDArray extends NDResource, BytesSupplier {
             throw new IllegalStateException(
                     "DataType mismatch, Required int" + " Actual " + getDataType());
         }
-        IntBuffer ib = toByteBuffer().asIntBuffer();
+        IntBuffer ib = toByteBuffer(true).asIntBuffer();
         long[] ret = new long[ib.remaining()];
         for (int i = 0; i < ret.length; ++i) {
             ret[i] = ib.get() & 0X00000000FFFFFFFFL;
@@ -345,7 +364,7 @@ public interface NDArray extends NDResource, BytesSupplier {
             throw new IllegalStateException(
                     "DataType mismatch, Required long" + " Actual " + getDataType());
         }
-        LongBuffer lb = toByteBuffer().asLongBuffer();
+        LongBuffer lb = toByteBuffer(true).asLongBuffer();
         long[] ret = new long[lb.remaining()];
         lb.get(ret);
         return ret;
@@ -358,7 +377,7 @@ public interface NDArray extends NDResource, BytesSupplier {
      * @throws IllegalStateException when {@link DataType} of this {@code NDArray} mismatches
      */
     default byte[] toByteArray() {
-        ByteBuffer bb = toByteBuffer();
+        ByteBuffer bb = toByteBuffer(true);
         if (bb.hasArray() && bb.remaining() == bb.array().length) {
             return bb.array();
         }
@@ -374,7 +393,7 @@ public interface NDArray extends NDResource, BytesSupplier {
      * @throws IllegalStateException when {@link DataType} of this {@code NDArray} mismatches
      */
     default int[] toUint8Array() {
-        ByteBuffer bb = toByteBuffer();
+        ByteBuffer bb = toByteBuffer(true);
         int[] buf = new int[bb.remaining()];
         for (int i = 0; i < buf.length; ++i) {
             buf[i] = bb.get() & 0xff;
@@ -393,7 +412,7 @@ public interface NDArray extends NDResource, BytesSupplier {
             throw new IllegalStateException(
                     "DataType mismatch, Required boolean" + " Actual " + getDataType());
         }
-        ByteBuffer bb = toByteBuffer();
+        ByteBuffer bb = toByteBuffer(true);
         boolean[] ret = new boolean[bb.remaining()];
         for (int i = 0; i < ret.length; ++i) {
             ret[i] = bb.get() != 0;
diff --git a/api/src/main/java/ai/djl/util/passthrough/PassthroughNDArray.java b/api/src/main/java/ai/djl/util/passthrough/PassthroughNDArray.java
index 140012205..3654f2974 100644
--- a/api/src/main/java/ai/djl/util/passthrough/PassthroughNDArray.java
+++ b/api/src/main/java/ai/djl/util/passthrough/PassthroughNDArray.java
@@ -74,17 +74,17 @@ public class PassthroughNDArray extends NDArrayAdapter {
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         if (object instanceof ByteBuffer) {
             return (ByteBuffer) object;
         }
-        throw new UnsupportedOperationException("Operation not supported for FastText");
+        throw new UnsupportedOperationException("Operation not supported for PassthroughNDArray");
     }
 
     /** {@inheritDoc} */
     @Override
     public void intern(NDArray replaced) {
-        throw new UnsupportedOperationException("Operation not supported for FastText");
+        throw new UnsupportedOperationException("Operation not supported for PassthroughNDArray");
     }
 
     /** {@inheritDoc} */
diff --git a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmDataset.java b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmDataset.java
index 869718fa9..69ac433f5 100644
--- a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmDataset.java
+++ b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmDataset.java
@@ -148,7 +148,7 @@ public class LgbmDataset extends NDArrayAdapter {
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         throw new UnsupportedOperationException("Not supported by the LgbmDataset yet");
     }
 
diff --git a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDArray.java b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDArray.java
index 5203457b6..103866d7a 100644
--- a/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDArray.java
+++ b/engines/ml/lightgbm/src/main/java/ai/djl/ml/lightgbm/LgbmNDArray.java
@@ -39,8 +39,8 @@ public class LgbmNDArray extends NDArrayAdapter {
 
     private AtomicReference<SWIGTYPE_p_void> handle;
     private int typeConstant;
-    private SWIGTYPE_p_float floatData;
-    private SWIGTYPE_p_double doubleData;
+    private AtomicReference<SWIGTYPE_p_float> floatDataRef;
+    private AtomicReference<SWIGTYPE_p_double> doubleDataRef;
 
     LgbmNDArray(
             NDManager manager,
@@ -53,6 +53,8 @@ public class LgbmNDArray extends NDArrayAdapter {
         this.format = SparseFormat.DENSE;
         manager.attachInternal(uid, this);
         handle = new AtomicReference<>();
+        floatDataRef = new AtomicReference<>();
+        doubleDataRef = new AtomicReference<>();
     }
 
     /** {@inheritDoc} */
@@ -82,19 +84,19 @@ public class LgbmNDArray extends NDArrayAdapter {
             if (getDataType() == DataType.FLOAT32) {
                 typeConstant = lightgbmlibConstants.C_API_DTYPE_FLOAT32;
                 FloatBuffer d1 = toByteBuffer().asFloatBuffer();
-                floatData = lightgbmlib.new_floatArray(size);
+                floatDataRef.set(lightgbmlib.new_floatArray(size));
                 for (int i = 0; i < size; i++) {
-                    lightgbmlib.floatArray_setitem(floatData, i, d1.get(i));
+                    lightgbmlib.floatArray_setitem(floatDataRef.get(), i, d1.get(i));
                 }
-                handle.set(lightgbmlib.float_to_voidp_ptr(floatData));
+                handle.set(lightgbmlib.float_to_voidp_ptr(floatDataRef.get()));
             } else if (getDataType() == DataType.FLOAT64) {
                 typeConstant = lightgbmlibConstants.C_API_DTYPE_FLOAT64;
                 DoubleBuffer d1 = toByteBuffer().asDoubleBuffer();
-                doubleData = lightgbmlib.new_doubleArray(size);
+                doubleDataRef.set(lightgbmlib.new_doubleArray(size));
                 for (int i = 0; i < size; i++) {
-                    lightgbmlib.doubleArray_setitem(doubleData, i, d1.get(i));
+                    lightgbmlib.doubleArray_setitem(doubleDataRef.get(), i, d1.get(i));
                 }
-                handle.set(lightgbmlib.double_to_voidp_ptr(doubleData));
+                handle.set(lightgbmlib.double_to_voidp_ptr(doubleDataRef.get()));
             } else {
                 throw new IllegalArgumentException(
                         "The LightGBM operation can only be performed with a Float32 or Float64"
@@ -140,7 +142,7 @@ public class LgbmNDArray extends NDArrayAdapter {
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         if (data == null) {
             throw new UnsupportedOperationException("Cannot obtain value from DMatrix");
         }
@@ -151,18 +153,19 @@ public class LgbmNDArray extends NDArrayAdapter {
     /** {@inheritDoc} */
     @Override
     public void intern(NDArray replaced) {
+        LgbmNDArray array = (LgbmNDArray) replaced;
+
+        SWIGTYPE_p_float floatData = floatDataRef.getAndSet(array.floatDataRef.getAndSet(null));
         if (floatData != null) {
             lightgbmlib.delete_floatArray(floatData);
         }
+        SWIGTYPE_p_double doubleData = doubleDataRef.getAndSet(array.doubleDataRef.getAndSet(null));
         if (doubleData != null) {
             lightgbmlib.delete_doubleArray(doubleData);
         }
-        LgbmNDArray array = (LgbmNDArray) replaced;
+        handle.set(array.handle.getAndSet(null));
         data = array.data;
-        handle = array.handle;
         format = array.format;
-        floatData = array.floatData;
-        doubleData = array.doubleData;
         typeConstant = array.typeConstant;
         shape = array.shape;
         dataType = array.dataType;
@@ -180,11 +183,15 @@ public class LgbmNDArray extends NDArrayAdapter {
     @Override
     public void close() {
         super.close();
+        SWIGTYPE_p_float floatData = floatDataRef.getAndSet(null);
         if (floatData != null) {
             lightgbmlib.delete_floatArray(floatData);
         }
+        SWIGTYPE_p_double doubleData = doubleDataRef.getAndSet(null);
         if (doubleData != null) {
             lightgbmlib.delete_doubleArray(doubleData);
         }
+        handle.set(null);
+        data = null;
     }
 }
diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java
index 129e3de09..69eb6914b 100644
--- a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java
+++ b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDArray.java
@@ -77,7 +77,7 @@ public class XgbNDArray extends NDArrayAdapter {
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         if (data == null) {
             throw new UnsupportedOperationException("Cannot obtain value from DMatrix");
         }
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
index fafb5f97d..52fac203d 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
@@ -288,7 +288,7 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         if (getSparseFormat() != SparseFormat.DENSE) {
             throw new IllegalStateException("Require Dense NDArray, actual " + getSparseFormat());
         }
diff --git a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
index 0ba330980..0beaee96a 100644
--- a/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
+++ b/engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtNDArray.java
@@ -117,7 +117,7 @@ public class OrtNDArray extends NDArrayAdapter {
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         if (getDataType() == DataType.STRING) {
             throw new IllegalArgumentException("Please use toStringArray() for String NDArray.");
         }
diff --git a/engines/pytorch/pytorch-engine/build.gradle b/engines/pytorch/pytorch-engine/build.gradle
index 7eeb7c710..73107ab45 100644
--- a/engines/pytorch/pytorch-engine/build.gradle
+++ b/engines/pytorch/pytorch-engine/build.gradle
@@ -39,8 +39,4 @@ publishing {
     }
 }
 
-clean.doFirst {
-    delete fileTree(System.getProperty("user.home") + "/.djl.ai/pytorch/") {
-        include '**/*djl_torch.*'
-    }
-}
+clean.dependsOn ":engines:pytorch:pytorch-jni:clean"
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index 551a16d03..8b37ccf66 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -225,12 +225,12 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         if (getDataType() == DataType.STRING) {
             throw new UnsupportedOperationException(
                     "toByteBuffer is not supported for String tensor.");
         }
-        return JniUtils.getByteBuffer(this);
+        return JniUtils.getByteBuffer(this, tryDirect);
     }
 
     /** {@inheritDoc} */
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index 40a6a0065..aa35100db 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -1671,11 +1671,25 @@ public final class JniUtils {
         return new Shape(PyTorchLibrary.LIB.torchSizes(ndArray.getHandle()));
     }
 
-    public static ByteBuffer getByteBuffer(PtNDArray ndArray) {
+    public static ByteBuffer getByteBuffer(PtNDArray ndArray, boolean tryDirect) {
         // Operation is CPU only
         if (!ndArray.getDevice().equals(Device.cpu())) {
             ndArray = ndArray.toDevice(Device.cpu(), false);
         }
+        if (tryDirect) {
+            if (ndArray.isSparse()
+                    || getLayout(ndArray) == 2
+                    || !PyTorchLibrary.LIB.torchIsContiguous(ndArray.getHandle())) {
+                // keep the same lifecycle as origin NDArray
+                ndArray =
+                        new PtNDArray(
+                                ndArray.getManager(),
+                                PyTorchLibrary.LIB.torchToContiguous(ndArray.getHandle()));
+            }
+            return PyTorchLibrary.LIB
+                    .torchDirectByteBuffer(ndArray.getHandle())
+                    .order(ByteOrder.nativeOrder());
+        }
         return ByteBuffer.wrap(PyTorchLibrary.LIB.torchDataPtr(ndArray.getHandle()))
                 .order(ByteOrder.nativeOrder());
     }
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
index 54fc54191..9ba75c9a6 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
@@ -50,6 +50,12 @@ final class PyTorchLibrary {
 
     native byte[] torchDataPtr(long handle);
 
+    native ByteBuffer torchDirectByteBuffer(long handle);
+
+    native boolean torchIsContiguous(long handle);
+
+    native long torchToContiguous(long handle);
+
     native int torchDType(long handle);
 
     native int[] torchDevice(long handle);
diff --git a/engines/pytorch/pytorch-native/build.cmd b/engines/pytorch/pytorch-native/build.cmd
index 256636c61..adb3e6374 100644
--- a/engines/pytorch/pytorch-native/build.cmd
+++ b/engines/pytorch/pytorch-native/build.cmd
@@ -19,8 +19,17 @@ if exist %FILEPATH% (
     echo Finished downloading libtorch
 )
 
-if "%VERSION%" == "1.11.0" (
-    set PT_VERSION=V1_11_X
+if "%VERSION%" == "1.13.1" (
+    set PT_VERSION=V1_13_X
+)
+if "%VERSION%" == "2.0.1" (
+    set PT_VERSION=V1_13_X
+)
+if "%VERSION%" == "2.1.1" (
+    set PT_VERSION=V1_13_X
+)
+if "%VERSION%" == "2.1.2" (
+    set PT_VERSION=V1_13_X
 )
 
 if /i "%2:~0,2%" == "cu" (
diff --git a/engines/pytorch/pytorch-native/build.sh b/engines/pytorch/pytorch-native/build.sh
index 214c38a25..3e4f03d48 100755
--- a/engines/pytorch/pytorch-native/build.sh
+++ b/engines/pytorch/pytorch-native/build.sh
@@ -34,7 +34,7 @@ if [[ ! -d "libtorch" ]]; then
       curl -s https://download.pytorch.org/libtorch/${FLAVOR}/libtorch${CXX11ABI}-shared-with-deps-${VERSION}%2B${FLAVOR}.zip | jar xv >/dev/null
     fi
   elif [[ $PLATFORM == 'darwin' ]]; then
-    if [[ "$VERSION" =~ ^(2.2.)* ]]; then
+    if [[ "$VERSION" =~ ^(2.[2-9].*)$ ]]; then
       if [[ $ARCH == 'aarch64' ]]; then
         curl -s https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-${VERSION}.zip | jar xv >/dev/null
       else
@@ -53,6 +53,10 @@ if [[ ! -d "libtorch" ]]; then
   fi
 fi
 
+if [[ "$VERSION" == "1.13.1" || "$VERSION" == "2.0.1" || "$VERSION" =~ ^(2.1.*)$ ]]; then
+  PT_VERSION=V1_13_X
+fi
+
 if [[ "$FLAVOR" = cu* ]]; then
   USE_CUDA=1
 fi
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
index e6aa2c145..9ac419894 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_tensor.cc
@@ -319,6 +319,53 @@ JNIEXPORT jbyteArray JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchDataPtr
   API_END_RETURN()
 }
 
+JNIEXPORT jobject JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchDirectByteBuffer(
+    JNIEnv* env, jobject jthis, jlong jhandle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  // Check if the tensor is sparse or optimized by MKL-DNN, if so, throw an exception
+  if (tensor_ptr->is_sparse() || tensor_ptr->is_mkldnn()) {
+    env->ThrowNew(ENGINE_EXCEPTION_CLASS, "torchDirectByteBuffer() is not supported for sparse or MKL-DNN tensors");
+    return nullptr;
+  }
+  // Check if the tensor is contiguous, if not, throw an exception
+  if (!tensor_ptr->is_contiguous()) {
+    env->ThrowNew(ENGINE_EXCEPTION_CLASS, "torchDirectByteBuffer() requires the tensor to be contiguous");
+    return nullptr;
+  }
+  size_t nbytes = tensor_ptr->nbytes();
+  // todo indeed, we can remove it in future!
+  if (nbytes > 0x7fffffff) {
+    env->ThrowNew(ENGINE_EXCEPTION_CLASS, "torchDirectByteBuffer() is not supported for large tensor");
+    return nullptr;
+  }
+  // Use tensor.data_ptr() to obtain the data pointer, and create a direct ByteBuffer using NewDirectByteBuffer
+  void* data_ptr = tensor_ptr->data_ptr();
+  return env->NewDirectByteBuffer(data_ptr, nbytes);
+  API_END_RETURN()
+}
+
+JNIEXPORT jboolean JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIsContiguous(
+    JNIEnv* env, jobject jthis, jlong jhandle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  bool is_contiguous = tensor_ptr->is_contiguous();
+  return static_cast<jboolean>(is_contiguous);
+  API_END_RETURN()
+}
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchToContiguous(
+    JNIEnv* env, jobject jthis, jlong jhandle) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  // sparse and mkldnn are required to be converted to dense to access data ptr
+  auto tensor = (tensor_ptr->is_sparse() || tensor_ptr->is_mkldnn()) ? tensor_ptr->to_dense() : *tensor_ptr;
+  tensor = (tensor.is_contiguous()) ? tensor : tensor.contiguous();
+  const auto* result_ptr = new torch::Tensor(tensor);
+  return reinterpret_cast<jlong>(result_ptr);
+  API_END_RETURN()
+}
+
 JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchDeleteTensor(
     JNIEnv* env, jobject jthis, jlong jhandle) {
   API_BEGIN()
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc
index ccf2616dc..251966570 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_pointwise.cc
@@ -356,7 +356,7 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchAtan(JNIEnv*
 }
 
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchAtan2(
-JNIEnv* env, jobject jthis, jlong jself, jlong jother) {
+    JNIEnv* env, jobject jthis, jlong jself, jlong jother) {
   API_BEGIN()
   const auto* self_ptr = reinterpret_cast<torch::Tensor*>(jself);
   const auto* other_ptr = reinterpret_cast<torch::Tensor*>(jother);
diff --git a/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h b/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
index 968368180..4951ea885 100644
--- a/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
+++ b/engines/pytorch/pytorch-native/src/main/native/djl_pytorch_utils.h
@@ -18,8 +18,13 @@
 #include <jni.h>
 #include <torch/csrc/api/include/torch/enum.h>
 #include <torch/script.h>
-#include <variant>
+
 #include <iostream>
+#ifdef V1_13_X
+#include <c10/util/variant.h>
+#else
+#include <variant>
+#endif
 
 #include "djl_pytorch_jni_log.h"
 
@@ -29,10 +34,16 @@ namespace utils {
 
 #if !defined(__ANDROID__)
 // for image interpolation
+#ifdef V1_13_X
+typedef torch::variant<torch::enumtype::kNearest, torch::enumtype::kLinear, torch::enumtype::kBilinear,
+    torch::enumtype::kBicubic, torch::enumtype::kTrilinear, torch::enumtype::kArea, torch::enumtype::kNearestExact>
+    mode_t;
+#else
 typedef std::variant<torch::enumtype::kNearest, torch::enumtype::kLinear, torch::enumtype::kBilinear,
     torch::enumtype::kBicubic, torch::enumtype::kTrilinear, torch::enumtype::kArea, torch::enumtype::kNearestExact>
     mode_t;
 #endif
+#endif
 
 inline jint GetDTypeFromScalarType(const torch::ScalarType& type) {
   if (torch::kFloat32 == type) {
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
index 106fe450f..363d18ca4 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
@@ -179,7 +179,7 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         if (getDataType() == DataType.STRING) {
             throw new IllegalArgumentException("Please use toStringArray() for String NDArray.");
         }
diff --git a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDArray.java b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDArray.java
index 3461e35ef..46210fb0f 100644
--- a/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDArray.java
+++ b/engines/tensorrt/src/main/java/ai/djl/tensorrt/engine/TrtNDArray.java
@@ -55,7 +55,7 @@ public class TrtNDArray extends NDArrayAdapter {
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         data.rewind();
         return data;
     }
diff --git a/extensions/tokenizers/.gitignore b/extensions/tokenizers/.gitignore
index 41e087b17..171bfe581 100644
--- a/extensions/tokenizers/.gitignore
+++ b/extensions/tokenizers/.gitignore
@@ -1,3 +1,6 @@
 Cargo.lock
 /tokenizers
 /jnilib
+model/
+tmp/
+models.json
diff --git a/extensions/tokenizers/README.md b/extensions/tokenizers/README.md
index fc700007b..a7d08306a 100644
--- a/extensions/tokenizers/README.md
+++ b/extensions/tokenizers/README.md
@@ -34,6 +34,10 @@ You can pull the module from the central Maven repository by including the follo
 If you are trying to convert a complete HuggingFace (transformers) model,
 you can try to use our all-in-one conversion solution to convert to Java:
 
+
+
+#### Convert Huggingface model to torchscript
+
 ```bash
 python3 -m pip install -r src/main/python/requirements.txt
 python3 src/main/python/model_zoo_importer.py -m deepset/bert-base-cased-squad2
@@ -45,6 +49,32 @@ This will generate a zip file into your local folder:
 model/nlp/question_answer/ai/djl/huggingface/pytorch/deepset/bert-base-cased-squad2/0.0.1/bert-base-cased-squad2.zip
 ```
 
+#### Convert Huggingface model to OnnxRuntime
+
+```bash
+python3 -m pip install -r src/main/python/requirements.txt
+python3 src/main/python/model_zoo_importer.py -m deepset/bert-base-cased-squad2 -f OnnxRuntime
+```
+
+#### Convert Huggingface model to Rust
+
+```bash
+python3 -m pip install -r src/main/python/requirements.txt
+python3 src/main/python/model_zoo_importer.py -m deepset/bert-base-cased-squad2 -f Rust
+```
+#### Use command line
+You can also use install djl-converter from source and use it as commandline:
+
+```bash
+cd djl/extensions/tokenizers/src/main/python
+python3 -m pip install -e .
+# convert a single model to TorchScript, Onnxruntime or Rust
+djl-convert --help
+
+# import models as DJL Model Zoo
+djl-import --help
+```
+
 Then, all you need to do, is to load this model in DJL:
 
 ```java
@@ -76,7 +106,7 @@ tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/msmarco-distilb
 Java
 
 ```java
-HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance("sentence-transformers/msmarco-distilbert-dot-v5")
+HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance("sentence-transformers/msmarco-distilbert-dot-v5");
 ```
 
 This way requires network connection to huggingface repo.
@@ -101,7 +131,7 @@ From your local directory, you will find a `tokenizer.json` file.
 Java
 
 ```java
-HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance(Paths.get("./tokenizer.json"))
+HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance(Paths.get("./tokenizer.json"));
 ```
 
 ### From pretrained json file
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 443dd1392..8799c389a 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
@@ -202,7 +202,7 @@ public class RsNDArray extends NativeResource<Long> implements NDArray {
 
     /** {@inheritDoc} */
     @Override
-    public ByteBuffer toByteBuffer() {
+    public ByteBuffer toByteBuffer(boolean tryDirect) {
         ByteBuffer bb = RustLibrary.getByteBuffer(getHandle());
         bb.order(ByteOrder.nativeOrder());
         return bb;
diff --git a/extensions/tokenizers/src/main/python/.gitignore b/extensions/tokenizers/src/main/python/.gitignore
index a96c5d3f5..567994ceb 100644
--- a/extensions/tokenizers/src/main/python/.gitignore
+++ b/extensions/tokenizers/src/main/python/.gitignore
@@ -1,4 +1,2 @@
 __pycache__
-model/
-tmp/
-models.json
+*.egg-info/
diff --git a/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py b/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py
index 4561cbc02..ebc97b979 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/arg_parser.py
@@ -16,6 +16,50 @@ import os
 
 
 def converter_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-o", "--output-dir", help="Model output directory")
+    parser.add_argument("-f",
+                        "--output-format",
+                        default="PyTorch",
+                        choices=["PyTorch", "OnnxRuntime", "Rust"],
+                        help="Model output format")
+    parser.add_argument("-u",
+                        "--cpu-only",
+                        action='store_true',
+                        help="Only validate jit traced model on CPU")
+    parser.add_argument("-m",
+                        "--model-id",
+                        help="Huggingface model_id to convert")
+    parser.add_argument("-r", "--revision", help="Huggingface model revision")
+    parser.add_argument("-t", "--token", help="Huggingface token")
+    parser.add_argument("--optimize",
+                        choices=["O1", "O2", "O3", "O4"],
+                        help="Optimization option for ONNX models")
+    parser.add_argument(
+        "--device",
+        help='The device to use to do the export. Defaults to "cpu".')
+    parser.add_argument(
+        "--dtype",
+        choices=["fp32", "fp16", "bf16"],
+        help="The floating point precision to use for the export.")
+    parser.add_argument(
+        "--trust-remote-code",
+        action="store_true",
+        help=
+        "Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository."
+    )
+
+    args = parser.parse_args()
+    if args.output_dir is None:
+        args.output_dir = "."
+
+    if not os.path.exists(args.output_dir):
+        raise ValueError(f"Invalid output directory: {args.output_dir}.")
+
+    return args
+
+
+def importer_args():
     parser = argparse.ArgumentParser()
     parser.add_argument("-l",
                         "--limit",
diff --git a/extensions/tokenizers/src/main/python/djl_converter/fill_mask_converter.py b/extensions/tokenizers/src/main/python/djl_converter/fill_mask_converter.py
index ff9de4bea..5c225293f 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/fill_mask_converter.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/fill_mask_converter.py
@@ -14,7 +14,7 @@ import logging
 
 import torch
 
-from huggingface_converter import HuggingfaceConverter
+from djl_converter.huggingface_converter import HuggingfaceConverter
 
 
 class FillMaskConverter(HuggingfaceConverter):
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 167b4ad81..c76656a2a 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/huggingface_converter.py
@@ -17,14 +17,14 @@ import sys
 from argparse import Namespace
 
 import onnx
-import safetensors_convert
+from djl_converter.safetensors_convert import convert_file
 import torch
 from huggingface_hub import hf_hub_download, HfApi
 from transformers import pipeline, AutoTokenizer, AutoConfig
 
-from metadata import HuggingfaceMetadata
-from shasum import sha1_sum
-from zip_utils import zip_dir
+from djl_converter.metadata import HuggingfaceMetadata
+from djl_converter.shasum import sha1_sum
+from djl_converter.zip_utils import zip_dir
 
 
 class PipelineHolder(object):
@@ -52,15 +52,18 @@ class HuggingfaceConverter:
         self.outputs = None
         self.api = HfApi()
 
-    def save_model(self, model_info, args: Namespace, temp_dir: str):
+    def save_model(self, model_info, args: Namespace, temp_dir: str,
+                   model_zoo: bool):
         if args.output_format == "OnnxRuntime":
-            return self.save_onnx_model(model_info, args, temp_dir)
+            return self.save_onnx_model(model_info, args, temp_dir, model_zoo)
         elif args.output_format == "Rust":
-            return self.save_rust_model(model_info, args, temp_dir)
+            return self.save_rust_model(model_info, args, temp_dir, model_zoo)
         else:
-            return self.save_pytorch_model(model_info, args, temp_dir)
+            return self.save_pytorch_model(model_info, args, temp_dir,
+                                           model_zoo)
 
-    def save_onnx_model(self, model_info, args: Namespace, temp_dir: str):
+    def save_onnx_model(self, model_info, args: Namespace, temp_dir: str,
+                        model_zoo: bool):
         model_id = model_info.modelId
 
         if not os.path.exists(temp_dir):
@@ -70,9 +73,17 @@ class HuggingfaceConverter:
 
         from optimum.commands.optimum_cli import main
 
-        sys.argv = [
-            "model_zoo_importer.py", "export", "onnx", "-m", model_id, temp_dir
-        ]
+        sys.argv = ["model_zoo_importer.py", "export", "onnx", "-m", model_id]
+        if args.optimize:
+            sys.argv.extend(["--optimize", args.optimize])
+        if args.device:
+            sys.argv.extend(["--device", args.device])
+        if args.dtype:
+            sys.argv.extend(["--dtype", args.dtype])
+        if args.trust_remote_code:
+            sys.argv.append("--trust-remote-code")
+        sys.argv.append(temp_dir)
+
         main()
 
         model = onnx.load_model(os.path.join(temp_dir, "model.onnx"),
@@ -83,13 +94,19 @@ class HuggingfaceConverter:
         tokenizer = AutoTokenizer.from_pretrained(model_id)
         config = AutoConfig.from_pretrained(model_id)
         hf_pipeline = PipelineHolder(tokenizer, ModelHolder(config))
-        size = self.save_to_model_zoo(model_info, args.output_dir,
-                                      "OnnxRuntime", temp_dir, hf_pipeline,
-                                      include_types)
+        arguments = self.save_serving_properties(model_info, "OnnxRuntime",
+                                                 temp_dir, hf_pipeline,
+                                                 include_types)
+        if model_zoo:
+            size = self.save_to_model_zoo(model_info, args.output_dir,
+                                          "OnnxRuntime", temp_dir, arguments)
+        else:
+            size = -1
 
         return True, None, size
 
-    def save_rust_model(self, model_info, args: Namespace, temp_dir: str):
+    def save_rust_model(self, model_info, args: Namespace, temp_dir: str,
+                        model_zoo: bool):
         model_id = model_info.modelId
 
         config = AutoConfig.from_pretrained(model_id)
@@ -139,16 +156,22 @@ class HuggingfaceConverter:
         elif has_pt_file:
             file = hf_hub_download(repo_id=model_id,
                                    filename="pytorch_model.bin")
-            safetensors_convert.convert_file(file, target)
+            convert_file(file, target)
         else:
             return False, f"No model file found for: {model_id}", -1
 
-        size = self.save_to_model_zoo(model_info, args.output_dir, "Rust",
-                                      temp_dir, hf_pipeline, include_types)
+        arguments = self.save_serving_properties(model_info, "Rust", temp_dir,
+                                                 hf_pipeline, include_types)
+        if model_zoo:
+            size = self.save_to_model_zoo(model_info, args.output_dir, "Rust",
+                                          temp_dir, arguments)
+        else:
+            size = -1
 
         return True, None, size
 
-    def save_pytorch_model(self, model_info, args: Namespace, temp_dir: str):
+    def save_pytorch_model(self, model_info, args: Namespace, temp_dir: str,
+                           model_zoo: bool):
         model_id = model_info.modelId
         if not os.path.exists(temp_dir):
             os.makedirs(temp_dir)
@@ -194,8 +217,14 @@ class HuggingfaceConverter:
             if not result:
                 return False, reason, -1
 
-        size = self.save_to_model_zoo(model_info, args.output_dir, "PyTorch",
-                                      temp_dir, hf_pipeline, include_types)
+        arguments = self.save_serving_properties(model_info, "PyTorch",
+                                                 temp_dir, hf_pipeline,
+                                                 include_types)
+        if model_zoo:
+            size = self.save_to_model_zoo(model_info, args.output_dir,
+                                          "PyTorch", temp_dir, arguments)
+        else:
+            size = -1
 
         return True, None, size
 
@@ -245,17 +274,11 @@ class HuggingfaceConverter:
 
         return model_file
 
-    def save_to_model_zoo(self, model_info, output_dir: str, engine: str,
-                          temp_dir: str, hf_pipeline, include_types: bool):
+    def save_serving_properties(self, model_info, engine: str, temp_dir: str,
+                                hf_pipeline, include_types: bool) -> dict:
         model_id = model_info.modelId
         model_name = model_id.split("/")[-1]
-        group_id = f"ai/djl/huggingface/{engine.lower()}"
-        repo_dir = f"{output_dir}/model/{self.application}/{group_id}/{model_id}"
-        model_dir = f"{repo_dir}/0.0.1"
-        if not os.path.exists(model_dir):
-            os.makedirs(model_dir)
 
-        # Save serving.properties
         serving_file = os.path.join(temp_dir, "serving.properties")
         arguments = self.get_extra_arguments(hf_pipeline, model_id, temp_dir)
         if include_types:
@@ -271,13 +294,25 @@ class HuggingfaceConverter:
             for k, v in arguments.items():
                 f.write(f"{k}={v}\n")
 
+        arguments["engine"] = engine
+        return arguments
+
+    def save_to_model_zoo(self, model_info, output_dir: str, engine: str,
+                          temp_dir: str, arguments: dict):
+        model_id = model_info.modelId
+        model_name = model_id.split("/")[-1]
+        group_id = f"ai/djl/huggingface/{engine.lower()}"
+        repo_dir = f"{output_dir}/model/{self.application}/{group_id}/{model_id}"
+        model_dir = f"{repo_dir}/0.0.1"
+        if not os.path.exists(model_dir):
+            os.makedirs(model_dir)
+
         # Save model as .zip file
         logging.info(f"Saving DJL model as zip: {model_name}.zip ...")
         zip_file = os.path.join(model_dir, f"{model_name}.zip")
         zip_dir(temp_dir, zip_file)
 
         # Save metadata.json
-        arguments["engine"] = engine
         sha1 = sha1_sum(zip_file)
         file_size = os.path.getsize(zip_file)
         metadata = HuggingfaceMetadata(model_info, engine, self.application,
diff --git a/extensions/tokenizers/src/main/python/djl_converter/huggingface_models.py b/extensions/tokenizers/src/main/python/djl_converter/huggingface_models.py
index cba19d477..e856871d1 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/huggingface_models.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/huggingface_models.py
@@ -19,6 +19,12 @@ from typing import List
 from huggingface_hub import HfApi
 from huggingface_hub import hf_hub_download
 from huggingface_hub.hf_api import ModelInfo
+from djl_converter.fill_mask_converter import FillMaskConverter
+from djl_converter.metadata import get_lang_tags
+from djl_converter.question_answering_converter import QuestionAnsweringConverter
+from djl_converter.sentence_similarity_converter import SentenceSimilarityConverter
+from djl_converter.text_classification_converter import TextClassificationConverter
+from djl_converter.token_classification_converter import TokenClassificationConverter
 
 ARCHITECTURES_2_TASK = {
     "ForQuestionAnswering": "question-answering",
@@ -27,19 +33,13 @@ ARCHITECTURES_2_TASK = {
     "ForMultipleChoice": "text-classification",
     "ForMaskedLM": "fill-mask",
 }
-LANGUAGES = HfApi().get_model_tags()["language"]
-
-
-def get_lang_tags(model_info):
-    tags = {}
-    for tag in model_info.tags:
-        if tag in LANGUAGES:
-            tags[tag] = "true"
-
-    if not tags:
-        tags["en"] = "true"
-
-    return tags
+SUPPORTED_TASKS = {
+    "fill-mask": FillMaskConverter(),
+    "question-answering": QuestionAnsweringConverter(),
+    "sentence-similarity": SentenceSimilarityConverter(),
+    "text-classification": TextClassificationConverter(),
+    "token-classification": TokenClassificationConverter(),
+}
 
 
 class HuggingfaceModels:
diff --git a/extensions/tokenizers/src/main/python/djl_converter/metadata.py b/extensions/tokenizers/src/main/python/djl_converter/metadata.py
index 7e835eb94..3b61e9667 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/metadata.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/metadata.py
@@ -11,7 +11,22 @@
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
 import json
-from huggingface_models import get_lang_tags
+
+from huggingface_hub import HfApi
+
+LANGUAGES = HfApi().get_model_tags()["language"]
+
+
+def get_lang_tags(model_info):
+    tags = {}
+    for tag in model_info.tags:
+        if tag in LANGUAGES:
+            tags[tag] = "true"
+
+    if not tags:
+        tags["en"] = "true"
+
+    return tags
 
 
 class HuggingfaceMetadata:
diff --git a/extensions/tokenizers/src/main/python/djl_converter/model_converter.py b/extensions/tokenizers/src/main/python/djl_converter/model_converter.py
new file mode 100644
index 000000000..01009a2d6
--- /dev/null
+++ b/extensions/tokenizers/src/main/python/djl_converter/model_converter.py
@@ -0,0 +1,64 @@
+#!/usr/bin/env python
+#
+# Copyright 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file
+# except in compliance with the License. A copy of the License is located at
+#
+# http://aws.amazon.com/apache2.0/
+#
+# or in the "LICENSE.txt" file accompanying this file. This file is distributed on an "AS IS"
+# BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
+# the specific language governing permissions and limitations under the License.
+import logging
+import os
+import sys
+
+from huggingface_hub import HfApi
+
+sys.path.append(os.path.dirname(os.path.realpath(__file__)))
+
+from djl_converter.arg_parser import converter_args
+
+
+def main():
+    logging.basicConfig(stream=sys.stdout,
+                        format="%(message)s",
+                        level=logging.INFO)
+    args = converter_args()
+
+    output_dir = args.output_dir
+    if output_dir == ".":
+        output_dir = f"model/{args.model_id.split('/')[-1]}"
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    if os.listdir(output_dir):
+        logging.error(f"output directory: {output_dir} is not empty.")
+        return
+
+    api = HfApi()
+    model_info = api.model_info(args.model_id,
+                                revision=args.revision,
+                                token=args.token)
+
+    from djl_converter.huggingface_models import HuggingfaceModels, SUPPORTED_TASKS
+
+    task, arch = HuggingfaceModels.to_supported_task(model_info.config)
+    converter = SUPPORTED_TASKS[task]
+
+    try:
+        result, reason, _ = converter.save_model(model_info, args, output_dir,
+                                                 False)
+        if result:
+            logging.info(f"Convert model {model_info.modelId} finished.")
+        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)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/extensions/tokenizers/src/main/python/djl_converter/model_zoo_importer.py b/extensions/tokenizers/src/main/python/djl_converter/model_zoo_importer.py
index de7599a5a..7e569e832 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/model_zoo_importer.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/model_zoo_importer.py
@@ -11,32 +11,23 @@
 # BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, express or implied. See the License for
 # the specific language governing permissions and limitations under the License.
 import logging
-import os.path
+import os
 import shutil
 import sys
 
-from arg_parser import converter_args
-from fill_mask_converter import FillMaskConverter
-from huggingface_models import HuggingfaceModels
-from question_answering_converter import QuestionAnsweringConverter
-from sentence_similarity_converter import SentenceSimilarityConverter
-from text_classification_converter import TextClassificationConverter
-from token_classification_converter import TokenClassificationConverter
+sys.path.append(os.path.dirname(os.path.realpath(__file__)))
 
-SUPPORTED_TASK = {
-    "fill-mask": FillMaskConverter(),
-    "question-answering": QuestionAnsweringConverter(),
-    "sentence-similarity": SentenceSimilarityConverter(),
-    "text-classification": TextClassificationConverter(),
-    "token-classification": TokenClassificationConverter(),
-}
+from djl_converter.arg_parser import importer_args
 
 
 def main():
     logging.basicConfig(stream=sys.stdout,
                         format="%(message)s",
                         level=logging.INFO)
-    args = converter_args()
+    args = importer_args()
+
+    # import transformer takes a long time
+    from djl_converter.huggingface_models import HuggingfaceModels, SUPPORTED_TASKS
 
     huggingface_models = HuggingfaceModels(args.output_dir)
     temp_dir = f"{args.output_dir}/tmp"
@@ -48,11 +39,11 @@ def main():
     for model in models:
         task = model["task"]
         model_info = model["model_info"]
-        converter = SUPPORTED_TASK[task]
+        converter = SUPPORTED_TASKS[task]
 
         try:
             result, reason, size = converter.save_model(
-                model_info, args, temp_dir)
+                model_info, args, temp_dir, True)
             if not result:
                 logging.error(f"{model_info.modelId}: {reason}")
         except Exception as e:
diff --git a/extensions/tokenizers/src/main/python/djl_converter/question_answering_converter.py b/extensions/tokenizers/src/main/python/djl_converter/question_answering_converter.py
index 935c21186..dd0f73920 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/question_answering_converter.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/question_answering_converter.py
@@ -14,7 +14,7 @@ import logging
 
 import torch
 
-from huggingface_converter import HuggingfaceConverter
+from djl_converter.huggingface_converter import HuggingfaceConverter
 
 
 class QuestionAnsweringConverter(HuggingfaceConverter):
diff --git a/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py b/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py
index 7d28db57e..68d6e422a 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/sentence_similarity_converter.py
@@ -19,7 +19,7 @@ import requests
 import torch
 from transformers import AutoTokenizer, AutoModel, AutoConfig
 
-from huggingface_converter import HuggingfaceConverter, PipelineHolder
+from djl_converter.huggingface_converter import HuggingfaceConverter, PipelineHolder
 from huggingface_hub import hf_hub_download
 
 
diff --git a/extensions/tokenizers/src/main/python/djl_converter/text_classification_converter.py b/extensions/tokenizers/src/main/python/djl_converter/text_classification_converter.py
index 8b1dba4ef..56552d111 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/text_classification_converter.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/text_classification_converter.py
@@ -15,7 +15,7 @@ import math
 
 import torch
 
-from huggingface_converter import HuggingfaceConverter
+from djl_converter.huggingface_converter import HuggingfaceConverter
 
 
 class TextClassificationConverter(HuggingfaceConverter):
diff --git a/extensions/tokenizers/src/main/python/djl_converter/token_classification_converter.py b/extensions/tokenizers/src/main/python/djl_converter/token_classification_converter.py
index fead8d247..33f910d3c 100644
--- a/extensions/tokenizers/src/main/python/djl_converter/token_classification_converter.py
+++ b/extensions/tokenizers/src/main/python/djl_converter/token_classification_converter.py
@@ -14,7 +14,7 @@ import logging
 
 import torch
 
-from huggingface_converter import HuggingfaceConverter
+from djl_converter.huggingface_converter import HuggingfaceConverter
 
 
 class TokenClassificationConverter(HuggingfaceConverter):
diff --git a/extensions/tokenizers/src/main/python/requirements.txt b/extensions/tokenizers/src/main/python/requirements.txt
index 28adc0488..231427da0 100644
--- a/extensions/tokenizers/src/main/python/requirements.txt
+++ b/extensions/tokenizers/src/main/python/requirements.txt
@@ -1,5 +1,6 @@
 huggingface_hub
 transformers
+sentence_transformers
 torch
 protobuf==3.20.2
 optimum[exporters,onnxruntime]
diff --git a/extensions/tokenizers/src/main/python/setup.py b/extensions/tokenizers/src/main/python/setup.py
index 68acd5fd1..50f5dd623 100755
--- a/extensions/tokenizers/src/main/python/setup.py
+++ b/extensions/tokenizers/src/main/python/setup.py
@@ -35,7 +35,7 @@ if __name__ == '__main__':
     version = detect_version()
 
     requirements = [
-        'huggingface_hub', 'transformers', 'onnx', 'sentence_transformers',
+        'huggingface_hub', 'transformers', 'sentence_transformers', 'torch',
         'optimum[exporters,onnxruntime]', 'safetensors'
     ]
 
@@ -52,7 +52,8 @@ if __name__ == '__main__':
           install_requires=requirements,
           entry_points={
               'console_scripts': [
-                  'djl-converter=djl_converter.model_zoo_importer:main',
+                  'djl-import=djl_converter.model_zoo_importer:main',
+                  'djl-convert=djl_converter.model_converter:main',
               ]
           },
           include_package_data=True,


Test patch:
diff --git a/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java b/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java
index 295571745..8ef4b7ade 100644
--- a/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java
+++ b/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java
@@ -28,19 +28,6 @@ public class TextGenerationTest {
         TestRequirements.linux();
         TestRequirements.weekly();
 
-        // Beam with Ort
-        String[] output0 = TextGeneration.generateTextWithOnnxRuntimeBeam();
-        Assert.assertEquals(
-                output0[0],
-                "DeepMind Company is a global leader in the field of artificial intelligence and"
-                    + " artificial intelligence research and development.\n"
-                    + "\n"
-                    + "Our mission is to provide the world with the best and brightest minds in the"
-                    + " field of artificial intelligence and artificial intelligence research and"
-                    + " development.\n"
-                    + "\n"
-                    + "Our mission is to provide the world with the best");
-
         // Greedy
         String expected =
                 "DeepMind Company is a global leader in the field of artificial"
@@ -82,11 +69,10 @@ public class TextGenerationTest {
                     + " development.\n"
                     + "\n"
                     + "Our mission is to");
-        Assert.assertEquals(
-                output2[3],
-                "Memories follow me left and right. I can't tell you how many times I've been told"
-                    + " that I'm not a good person. I'm not a good person. I'm not a good person."
-                    + " I'm not a good person. I'm not a good person. I'm not a");
+        Assert.assertTrue(
+                output2[3].startsWith(
+                        "Memories follow me left and right. I can't tell you how many times I've"
+                                + " been told that I'm not a good person."));
     }
 
     @Test
@@ -119,4 +105,25 @@ public class TextGenerationTest {
                         + "But if you're lucky, you can escape from prison and live happily ever"
                         + " after.\n");
     }
+
+    @Test
+    public void testTextGenerationWithOnnx()
+            throws TranslateException, ModelException, IOException {
+        TestRequirements.linux();
+        TestRequirements.weekly();
+        TestRequirements.engine("PyTorch");
+
+        // Beam with Ort
+        String[] output0 = TextGeneration.generateTextWithOnnxRuntimeBeam();
+        Assert.assertEquals(
+                output0[0],
+                "DeepMind Company is a global leader in the field of artificial intelligence and"
+                    + " artificial intelligence research and development.\n"
+                    + "\n"
+                    + "Our mission is to provide the world with the best and brightest minds in the"
+                    + " field of artificial intelligence and artificial intelligence research and"
+                    + " development.\n"
+                    + "\n"
+                    + "Our mission is to provide the world with the best");
+    }
 }