Repository: deepjavalibrary/djl
Pull request #2850 (2023-11-13 17:58:27)

[tokenizer] Update import script for huggingface_hub api change
## 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/build_triton_and_ft.yml b/.github/workflows/build_triton_and_ft.yml
index 3f7401f94..7d0b9aaf7 100644
--- a/.github/workflows/build_triton_and_ft.yml
+++ b/.github/workflows/build_triton_and_ft.yml
@@ -1,4 +1,4 @@
-name: Build Triton Server and FasterTransformers
+name: Build Triton Server
 
 on:
   workflow_dispatch:
@@ -6,21 +6,7 @@ on:
       triton:
         description: 'triton branch version'
         required: true
-        default: 'r23.04'
-      fastertransformer:
-        description: 'fastertransformer branch/tag version'
-        required: true
-        default: 'main'
-      build_triton_only:
-        description: 'whether to just build triton library'
-        required: false
-        type: boolean
-        default: false
-      is_llama_build:
-        description: 'whether to build custom llama source'
-        required: false
-        type: boolean
-        default: false
+        default: 'r23.10'
 
 jobs:
   build-triton:
@@ -38,6 +24,7 @@ jobs:
       - name: Build Triton Binary
         shell: 'script -q -e -c "bash --noprofile --norc -eo pipefail {0}"'
         run: |
+          pip3 install requests
           python3 build.py --enable-logging --enable-metrics --enable-stats --enable-cpu-metrics --enable-gpu --endpoint http
       - name: Configure AWS Credentials
         uses: aws-actions/configure-aws-credentials@v2
@@ -50,60 +37,3 @@ jobs:
           aws s3 cp build/install/lib/libtritonserver.so s3://djl-ai/publish/tritonserver/${{ github.event.inputs.triton }}/
           aws s3 cp build/install/bin/tritonserver s3://djl-ai/publish/tritonserver/${{ github.event.inputs.triton }}/
           aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/tritonserver/${{ github.event.inputs.triton }}/*"
-
-  create-runner:
-    if: ${{ github.repository == 'deepjavalibrary/djl' && ! github.event.inputs.build_triton_only }}
-    runs-on: [ self-hosted, scheduler ]
-    steps:
-      - name: Create new CPU instance
-        id: create_cpu
-        run: |
-          cd /home/ubuntu/djl_benchmark_script/scripts
-          token=$( curl -X POST -H "Authorization: token ${{ secrets.ACTION_RUNNER_PERSONAL_TOKEN }}" \
-          https://api.github.com/repos/deepjavalibrary/djl/actions/runners/registration-token \
-          --fail \
-          | jq '.token' | tr -d '"' )
-          ./start_instance.sh action_cpu $token djl
-    outputs:
-      cpu_instance_id: ${{ steps.create_cpu.outputs.action_cpu_instance_id }}
-
-
-  build-fastertransformer:
-    if: ${{ github.repository == 'deepjavalibrary/djl' }}
-    runs-on: [ self-hosted, cpu ]
-    container: deepjavalibrary/djl-serving:fastertransformer-nightly
-    timeout-minutes: 60
-    needs: create-runner
-    steps:
-      - uses: actions/checkout@v3
-      - name: Build FasterTransformers
-        run: |
-          tools/scripts/build_ft_deps.sh ${{ github.event.inputs.fastertransformer }} ${{ github.event.inputs.triton }} ${{ github.event.inputs.is_llama_build }}
-      - name: Configure AWS Credentials
-        uses: aws-actions/configure-aws-credentials@v2
-        with:
-          aws-access-key-id: ${{ secrets.AWS_ACCESS_KEY_ID }}
-          aws-secret-access-key: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
-          aws-region: us-east-2
-      - name: Copy files to S3 with the AWS CLI
-        if: github.event.inputs.is_llama_build == 'false'
-        run: |
-          aws s3 sync /tmp/binaries/ s3://djl-ai/publish/fastertransformer/${{ github.event.inputs.fastertransformer }}/
-          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/fastertransformer/${{ github.event.inputs.fastertransformer }}/*"
-      - name: Copy files for llama build to S3 with AWS CLI
-        if: github.event.inputs.is_llama_build == 'true'
-        run: |
-          echo "pushing binaries to ft/llama"
-          aws s3 sync /tmp/binaries/ s3://djl-ai/publish/fastertransformer/llama/
-          aws cloudfront create-invalidation --distribution-id E371VB8JQ6NRVY --paths "/fastertransformer-llama/${{ github.event.inputs.fastertransformer }}/*"
-
-  stop-runner:
-    if: ${{ github.repository == 'deepjavalibrary/djl' && always() }}
-    runs-on: [ self-hosted, scheduler ]
-    needs: [ create-runner, build-fastertransformer]
-    steps:
-      - name: Stop all instances
-        run: |
-          cd /home/ubuntu/djl_benchmark_script/scripts
-          instance_id=${{ needs.create-runner.outputs.cpu_instance_id }}
-          ./stop_instance.sh $instance_id
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index 23d452be3..2b4a9df09 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -3393,6 +3393,48 @@ public interface NDArray extends NDResource, BytesSupplier {
             boolean normalize,
             boolean returnComplex);
 
+    /**
+     * Computes the two-dimensional Discrete Fourier Transform.
+     *
+     * @param sizes Sizes of the transformed axes of the output. Will be zero-padded or trimmed to
+     *     this size.
+     * @param axes Axes over which to compute the 2D-FFT.
+     * @return The truncated or zero-padded input, transformed along the axes.
+     */
+    NDArray fft2(long[] sizes, long[] axes);
+
+    /**
+     * Computes the two-dimensional Discrete Fourier Transform along the last 2 axes.
+     *
+     * @param sizes Sizes of the transformed axes of the output. Will be zero-padded or trimmed to
+     *     this size.
+     * @return The truncated or zero-padded input, transformed along the last two axes
+     */
+    default NDArray fft2(long[] sizes) {
+        return fft2(sizes, new long[] {-2, -1});
+    }
+
+    /**
+     * Computes the two-dimensional inverse Discrete Fourier Transform.
+     *
+     * @param sizes Sizes of the transformed axes of the output. Will be zero-padded or trimmed to
+     *     this size.
+     * @param axes Axes over which to compute the 2D-Inverse-FFT.
+     * @return The truncated or zero-padded input, transformed along the axes.
+     */
+    NDArray ifft2(long[] sizes, long[] axes);
+
+    /**
+     * Computes the two-dimensional inverse Discrete Fourier Transform along the last 2 axes.
+     *
+     * @param sizes Sizes of the transformed axes of the output. Will be zero-padded or trimmed to
+     *     this size.
+     * @return The truncated or zero-padded input, transformed along the axes.
+     */
+    default NDArray ifft2(long[] sizes) {
+        return ifft2(sizes, new long[] {-2, -1});
+    }
+
     /**
      * Reshapes this {@code NDArray} to the given {@link Shape}.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index 855c71830..9a4ad8db9 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -912,6 +912,18 @@ public abstract class NDArrayAdapter implements NDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
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 19eec8372..8b884b399 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
@@ -1160,6 +1160,18 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
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 1e9ac83c1..606aaf24e 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
@@ -1103,6 +1103,18 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
                 this, nFft, hopLength, (PtNDArray) window, center, normalize, returnComplex);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft2(long[] sizes, long[] axes) {
+        return JniUtils.fft2(this, sizes, axes);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft2(long[] sizes, long[] axes) {
+        return JniUtils.ifft2(this, sizes, axes);
+    }
+
     /** {@inheritDoc} */
     @Override
     public PtNDArray reshape(Shape shape) {
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 8e6be7b8d..40a6a0065 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
@@ -1040,6 +1040,18 @@ public final class JniUtils {
         return new PtNDArray(ndArray.getManager(), handle);
     }
 
+    public static PtNDArray fft2(PtNDArray ndArray, long[] sizes, long[] axes) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchFft2(ndArray.getHandle(), sizes, axes));
+    }
+
+    public static PtNDArray ifft2(PtNDArray ndArray, long[] sizes, long[] axes) {
+        return new PtNDArray(
+                ndArray.getManager(),
+                PyTorchLibrary.LIB.torchIfft2(ndArray.getHandle(), sizes, axes));
+    }
+
     public static PtNDArray real(PtNDArray ndArray) {
         long handle = PyTorchLibrary.LIB.torchViewAsReal(ndArray.getHandle());
         if (handle == -1) {
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 a1829306d..54fc54191 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
@@ -273,6 +273,10 @@ final class PyTorchLibrary {
             boolean normalize,
             boolean returnComplex);
 
+    native long torchFft2(long handle, long[] sizes, long[] axes);
+
+    native long torchIfft2(long handle, long[] sizes, long[] axes);
+
     native long torchViewAsReal(long handle);
 
     native long torchViewAsComplex(long handle);
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc
index 5a65e1eca..08932098d 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_torch_other.cc
@@ -34,6 +34,28 @@ JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchFft(
   API_END_RETURN()
 }
 
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchFft2(
+    JNIEnv* env, jobject jthis, jlong jhandle, jlongArray js, jlongArray jaxes) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const std::vector<int64_t> sizes = djl::utils::jni::GetVecFromJLongArray(env, js);
+  const std::vector<int64_t> axes = djl::utils::jni::GetVecFromJLongArray(env, jaxes);
+  const auto* result_ptr = new torch::Tensor(torch::fft_fft2(*tensor_ptr, sizes, axes));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
+}
+
+JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchIfft2(
+    JNIEnv* env, jobject jthis, jlong jhandle, jlongArray js, jlongArray jaxes) {
+  API_BEGIN()
+  const auto* tensor_ptr = reinterpret_cast<torch::Tensor*>(jhandle);
+  const std::vector<int64_t> sizes = djl::utils::jni::GetVecFromJLongArray(env, js);
+  const std::vector<int64_t> axes = djl::utils::jni::GetVecFromJLongArray(env, jaxes);
+  const auto* result_ptr = new torch::Tensor(torch::fft_ifft2(*tensor_ptr, sizes, axes));
+  return reinterpret_cast<uintptr_t>(result_ptr);
+  API_END_RETURN()
+}
+
 JNIEXPORT jlong JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_torchStft(JNIEnv* env, jobject jthis, jlong jhandle,
     jlong jn_fft, jlong jhop_length, jlong jwindow, jboolean jcenter, jboolean jnormalize, jboolean jreturn_complex) {
 #ifdef V1_11_X
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 444885377..419be4c09 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
@@ -1184,6 +1184,18 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
         throw new UnsupportedOperationException("Not implemented yet.");
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public NDArray fft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray ifft2(long[] sizes, long[] axes) {
+        throw new UnsupportedOperationException("Not implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray reshape(Shape shape) {
diff --git a/extensions/tokenizers/src/main/python/huggingface_models.py b/extensions/tokenizers/src/main/python/huggingface_models.py
index 7e5ada55e..549db3788 100644
--- a/extensions/tokenizers/src/main/python/huggingface_models.py
+++ b/extensions/tokenizers/src/main/python/huggingface_models.py
@@ -16,7 +16,7 @@ import os
 from argparse import Namespace
 from typing import List
 
-from huggingface_hub import HfApi, ModelSearchArguments
+from huggingface_hub import HfApi
 from huggingface_hub import hf_hub_download
 from huggingface_hub.hf_api import ModelInfo
 
@@ -27,7 +27,7 @@ ARCHITECTURES_2_TASK = {
     "ForMultipleChoice": "text-classification",
     "ForMaskedLM": "fill-mask",
 }
-LANGUAGES = ModelSearchArguments().language
+LANGUAGES = HfApi().get_model_tags()["language"]
 
 
 def get_lang_tags(model_info):
diff --git a/extensions/tokenizers/src/main/python/requirements.txt b/extensions/tokenizers/src/main/python/requirements.txt
index bf197b644..05ce0bc48 100644
--- a/extensions/tokenizers/src/main/python/requirements.txt
+++ b/extensions/tokenizers/src/main/python/requirements.txt
@@ -1,4 +1,4 @@
 huggingface_hub
 transformers
-torch==1.11.0
+torch
 protobuf==3.20.2
diff --git a/tools/scripts/build_ft_deps.sh b/tools/scripts/build_ft_deps.sh
deleted file mode 100755
index 110bedc2d..000000000
--- a/tools/scripts/build_ft_deps.sh
+++ /dev/null
@@ -1,49 +0,0 @@
-#!/usr/bin/env bash
-
-set -ex
-
-FT_VERSION=$1
-NVIDIA_TRITON_SERVER_VERSION=$2
-IS_LLAMA_BUILD=$3
-
-apt-get update && apt-get install -y rapidjson-dev
-
-pushd /tmp
-
-git clone https://github.com/NVIDIA/FasterTransformer.git -b ${FT_VERSION}
-
-export FT_DIR=/tmp/FasterTransformer
-mkdir -p /tmp/binaries
-
-# Build FasterTransformer Triton library
-if [ "$IS_LLAMA_BUILD" = "false" ] ; then
-      git clone https://github.com/triton-inference-server/fastertransformer_backend.git
-else
-      echo "cloning forked FT backend repo with llama support"
-      git clone https://github.com/rohithkrn/fastertransformer_backend.git -b llama_void_main
-fi
-mkdir -p fastertransformer_backend/build
-cd fastertransformer_backend/build
-cmake \
-      -D CMAKE_EXPORT_COMPILE_COMMANDS=1 \
-      -D CMAKE_BUILD_TYPE=Release \
-      -D ENABLE_FP8=OFF \
-      -D CMAKE_INSTALL_PREFIX=/opt/tritonserver \
-      -D TRITON_COMMON_REPO_TAG="${NVIDIA_TRITON_SERVER_VERSION}" \
-      -D TRITON_CORE_REPO_TAG="${NVIDIA_TRITON_SERVER_VERSION}" \
-      -D TRITON_BACKEND_REPO_TAG="${NVIDIA_TRITON_SERVER_VERSION}" \
-      ..
-make -j$(nproc) install
-cp /opt/tritonserver/backends/fastertransformer/*.so /tmp/binaries/
-cd ../../
-
-# Build FasterTransformer TH Ops library
-mkdir -p FasterTransformer/build
-cd FasterTransformer/build
-git submodule init && git submodule update
-cmake -DCMAKE_BUILD_TYPE=Release -DSM=70,75,80,86,90 -DBUILD_PYT=ON -DBUILD_MULTI_GPU=ON ..
-make -j$(nproc)
-cp lib/libth_transformer.so /tmp/binaries/
-cd ../../
-
-popd


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java
index 00e7465f7..66bb136ab 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayOtherOpTest.java
@@ -1087,4 +1087,58 @@ public class NDArrayOtherOpTest {
             Assertions.assertAlmostEquals(result.real().flatten(), expected);
         }
     }
+
+    @Test
+    public void testFft2() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray array =
+                    manager.create(
+                            new float[][] {
+                                {1f, 6.6f, 4.315f, 2.0f},
+                                {16.9f, 6.697f, 2.399f, 67.9f},
+                                {0f, 5f, 67.09f, 9.87f}
+                            });
+            NDArray result = array.fft2(new long[] {3, 4}, new long[] {0, 1});
+            result = result.real().flatten(1, 2); // flatten complex numbers
+            NDArray expected =
+                    manager.create(
+                            new float[][] {
+                                {189.771f, 0f, -55.904f, 61.473f, -6.363f, 0f, -55.904f, -61.473f},
+                                {
+                                    -74.013f,
+                                    -10.3369f,
+                                    71.7653f,
+                                    -108.2964f,
+                                    -1.746f,
+                                    93.1133f,
+                                    -25.8063f,
+                                    -33.0234f
+                                },
+                                {
+                                    -74.013f, 10.3369f, -25.8063f, 33.0234f, -1.746f, -93.1133f,
+                                    71.7653f, 108.2964f
+                                }
+                            });
+            Assertions.assertAlmostEquals(result, expected);
+        }
+    }
+
+    @Test
+    public void testIfft2() {
+        try (NDManager manager = NDManager.newBaseManager(TestUtils.getEngine())) {
+            NDArray array =
+                    manager.create(
+                            new float[][] {
+                                {1f, 6.6f, 4.315f, 2.0f},
+                                {16.9f, 6.697f, 2.399f, 67.9f},
+                                {0f, 5f, 67.09f, 9.87f}
+                            });
+            long[] sizes = {3, 4};
+            long[] axes = {0, 1};
+            NDArray fft2 = array.fft2(sizes, axes);
+            NDArray actual = fft2.ifft2(sizes, axes).real();
+            NDArray expected = array.toType(DataType.COMPLEX64, true).real();
+            Assertions.assertAlmostEquals(expected, actual);
+        }
+    }
 }