diff --git a/extension/android/jni/jni_layer_llama.cpp b/extension/android/jni/jni_layer_llama.cpp index 5c58c89ee91..e71323da201 100644 --- a/extension/android/jni/jni_layer_llama.cpp +++ b/extension/android/jni/jni_layer_llama.cpp @@ -132,16 +132,22 @@ class ExecuTorchLlamaJni jint model_type_category, facebook::jni::alias_ref model_path, facebook::jni::alias_ref tokenizer_path, - jfloat temperature) { + jfloat temperature, + facebook::jni::alias_ref data_path) { return makeCxxInstance( - model_type_category, model_path, tokenizer_path, temperature); + model_type_category, + model_path, + tokenizer_path, + temperature, + data_path); } ExecuTorchLlamaJni( jint model_type_category, facebook::jni::alias_ref model_path, facebook::jni::alias_ref tokenizer_path, - jfloat temperature) { + jfloat temperature, + facebook::jni::alias_ref data_path = nullptr) { #if defined(ET_USE_THREADPOOL) // Reserve 1 thread for the main thread. uint32_t num_performant_cores = @@ -160,10 +166,18 @@ class ExecuTorchLlamaJni tokenizer_path->toStdString().c_str(), temperature); } else if (model_type_category == MODEL_TYPE_CATEGORY_LLM) { - runner_ = std::make_unique( - model_path->toStdString().c_str(), - tokenizer_path->toStdString().c_str(), - temperature); + if (data_path != nullptr) { + runner_ = std::make_unique( + model_path->toStdString().c_str(), + tokenizer_path->toStdString().c_str(), + temperature, + data_path->toStdString().c_str()); + } else { + runner_ = std::make_unique( + model_path->toStdString().c_str(), + tokenizer_path->toStdString().c_str(), + temperature); + } #if defined(EXECUTORCH_BUILD_MEDIATEK) } else if (model_type_category == MODEL_TYPE_MEDIATEK_LLAMA) { runner_ = std::make_unique( diff --git a/extension/android/src/main/java/org/pytorch/executorch/LlamaModule.java b/extension/android/src/main/java/org/pytorch/executorch/LlamaModule.java index 6de26bc7fe8..d64c7d0bf89 100644 --- a/extension/android/src/main/java/org/pytorch/executorch/LlamaModule.java +++ b/extension/android/src/main/java/org/pytorch/executorch/LlamaModule.java @@ -39,16 +39,24 @@ public class LlamaModule { @DoNotStrip private static native HybridData initHybrid( - int modelType, String modulePath, String tokenizerPath, float temperature); + int modelType, String modulePath, String tokenizerPath, float temperature, String dataPath); - /** Constructs a LLAMA Module for a model with given path, tokenizer, and temperature. */ + /** Constructs a LLAMA Module for a model with given model path, tokenizer, temperature. */ public LlamaModule(String modulePath, String tokenizerPath, float temperature) { - mHybridData = initHybrid(MODEL_TYPE_TEXT, modulePath, tokenizerPath, temperature); + mHybridData = initHybrid(MODEL_TYPE_TEXT, modulePath, tokenizerPath, temperature, null); + } + + /** + * Constructs a LLAMA Module for a model with given model path, tokenizer, temperature and data + * path. + */ + public LlamaModule(String modulePath, String tokenizerPath, float temperature, String dataPath) { + mHybridData = initHybrid(MODEL_TYPE_TEXT, modulePath, tokenizerPath, temperature, dataPath); } /** Constructs a LLM Module for a model with given path, tokenizer, and temperature. */ public LlamaModule(int modelType, String modulePath, String tokenizerPath, float temperature) { - mHybridData = initHybrid(modelType, modulePath, tokenizerPath, temperature); + mHybridData = initHybrid(modelType, modulePath, tokenizerPath, temperature, null); } public void resetNative() {