Update setup.sh for tokenizer selection

Summary: For LLAMA3, users need to use tiktoken. Add a option to load from env var.

Differential Revision: D56374673

Author: kirklandsign

examples/demo-apps/android/LlamaDemo/README.md

@@ -38,7 +38,15 @@ Note: `<path_to_android_ndk>` is the root for the NDK, which is usually under
 `~/Library/Android/sdk/ndk/XX.Y.ZZZZZ` for macOS, and contains NOTICE and README.md.
 We use `<path_to_android_ndk>/build/cmake/android.toolchain.cmake` for CMake to cross-compile.
 
-3. Run the following command set up the required JNI library:
+3. (Optional) If you need to use tiktoken as the tokenizer (for LLaMA3), set
+`EXECUTORCH_USE_TIKTOKEN=ON` and later CMake will use it as the tokenizer.
+If you need to run other models like LLaMA2, skip this skip.
+
+```bash
+export EXECUTORCH_USE_TIKTOKEN=ON # Only for LLaMA3
+```
+
+4. Run the following command set up the required JNI library:
 ```bash
 pushd examples/demo-apps/android/LlamaDemo
 ./gradlew :app:setup

examples/demo-apps/android/LlamaDemo/setup.sh

@@ -8,6 +8,7 @@
 set -eu
 
 CMAKE_OUT="${CMAKE_OUT:-cmake-out-android}"
+EXECUTORCH_USE_TIKTOKEN="${EXECUTORCH_USE_TIKTOKEN:-OFF}"
 # Note: Set up ANDROID_NDK and ANDROID_ABI
 cmake . -DCMAKE_INSTALL_PREFIX="${CMAKE_OUT}" \
   -DCMAKE_TOOLCHAIN_FILE="${ANDROID_NDK}/build/cmake/android.toolchain.cmake" \
@@ -32,6 +33,7 @@ cmake examples/models/llama2 \
          -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \
          -DANDROID_ABI="$ANDROID_ABI" \
          -DCMAKE_INSTALL_PREFIX="${CMAKE_OUT}" \
+         -DEXECUTORCH_USE_TIKTOKEN="${EXECUTORCH_USE_TIKTOKEN}" \
          -DCMAKE_BUILD_TYPE=Release \
          -B"${CMAKE_OUT}"/examples/models/llama2
 
@@ -42,6 +44,7 @@ cmake extension/android \
   -DANDROID_ABI="${ANDROID_ABI}" \
   -DCMAKE_INSTALL_PREFIX="${CMAKE_OUT}" \
   -DEXECUTORCH_BUILD_LLAMA_JNI=ON \
+  -DEXECUTORCH_USE_TIKTOKEN="${EXECUTORCH_USE_TIKTOKEN}" \
   -DCMAKE_BUILD_TYPE=Release \
   -B"${CMAKE_OUT}"/extension/android
 

extension/android/CMakeLists.txt

@@ -80,4 +80,14 @@ if(EXECUTORCH_BUILD_LLAMA_JNI)
   target_link_libraries(executorch_llama_jni ${link_libraries} llama_runner
                         custom_ops cpublas eigen_blas)
   target_compile_options(executorch_llama_jni PUBLIC ${_common_compile_options})
+  if(EXECUTORCH_USE_TIKTOKEN)
+    set(ABSL_ENABLE_INSTALL ON)
+    set(_pic_flag
+      ${CMAKE_POSITION_INDEPENDENT_CODE})
+    set(CMAKE_POSITION_INDEPENDENT_CODE ON)
+    add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../examples/models/llama2/third-party/abseil-cpp ${CMAKE_CURRENT_BINARY_DIR}/abseil-cpp)
+    add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/../../examples/models/llama2/third-party/re2 ${CMAKE_CURRENT_BINARY_DIR}/re2)
+    set(CMAKE_POSITION_INDEPENDENT_CODE ${_pic_flag})
+    target_link_libraries(executorch_llama_jni re2::re2)
+  endif()
 endif()