Prefill

Author: kirklandsign

extension/android/jni/jni_layer_llama.cpp

@@ -121,7 +121,9 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
   float temperature_ = 0.0f;
   int model_type_category_;
   std::unique_ptr<llm::IRunner> runner_;
-  std::unique_ptr<executorch::extension::llm::MultimodalRunner> multi_modal_runner_;
+  std::unique_ptr<executorch::extension::llm::MultimodalRunner>
+      multi_modal_runner_;
+  std::vector<llm::MultimodalInput> prefill_inputs_;
 
  public:
   constexpr static auto kJavaDescriptor =
@@ -215,6 +217,9 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       facebook::jni::alias_ref<ExecuTorchLlmCallbackJni> callback,
       jboolean echo) {
     if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) {
+      std::vector<llm::MultimodalInput> inputs = prefill_inputs_;
+      prefill_inputs_.clear();
+      inputs.emplace_back(llm::MultimodalInput{prompt->toStdString()});
       auto image_size = image->size();
       std::vector<llm::Image> images;
       if (image_size != 0) {
@@ -225,15 +230,18 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
           image_data[i] = image_data_jint[i];
         }
         llm::Image image_runner{image_data, width, height, channels};
-        images.push_back(image_runner);
+        inputs.emplace_back(llm::MultimodalInput{std::move(image_runner)});
       }
+      executorch::extension::llm::GenerationConfig config{
+          .echo = static_cast<bool>(echo),
+          .seq_len = seq_len,
+          .temperature = temperature_,
+      };
       multi_modal_runner_->generate(
-          std::move(images),
-          prompt->toStdString(),
-          seq_len,
-          [callback](std::string result) { callback->onResult(result); },
-          [callback](const llm::Stats& result) { callback->onStats(result); },
-          echo);
+          std::move(inputs),
+          config,
+          [callback](const std::string& result) { callback->onResult(result); },
+          [callback](const llm::Stats& result) { callback->onStats(result); });
     } else if (model_type_category_ == MODEL_TYPE_CATEGORY_LLM) {
       executorch::extension::llm::GenerationConfig config{
           .echo = static_cast<bool>(echo),
@@ -257,9 +265,10 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       jlong start_pos,
       jint bos,
       jint eos) {
+    prefill_inputs_.emplace_back(llm::MultimodalInput{prompt->toStdString()});
     facebook::jni::local_ref<jlongArray> tuple_result =
         facebook::jni::make_long_array(2);
-    tuple_result->pin()[0] = static_cast<jint>(Error::NotSupported);
+    tuple_result->pin()[0] = static_cast<jint>(Error::Ok);
     return tuple_result;
   }
 
@@ -273,10 +282,24 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       jint height,
       jint channels,
       jlong start_pos) {
+    std::vector<llm::Image> images;
+    auto image_size = image->size();
+    if (image_size != 0) {
+      std::vector<jint> image_data_jint(image_size);
+      std::vector<uint8_t> image_data(image_size);
+      image->getRegion(0, image_size, image_data_jint.data());
+      for (int i = 0; i < image_size; i++) {
+        image_data[i] = image_data_jint[i];
+      }
+      llm::Image image_runner{image_data, width, height, channels};
+      prefill_inputs_.emplace_back(
+          llm::MultimodalInput{std::move(image_runner)});
+    }
+
     facebook::jni::local_ref<jlongArray> tuple_result =
         facebook::jni::make_long_array(2);
 
-    tuple_result->pin()[0] = static_cast<jint>(Error::NotSupported);
+    tuple_result->pin()[0] = static_cast<jint>(Error::Ok);
     return tuple_result;
   }
 
@@ -287,10 +310,13 @@ class ExecuTorchLlmJni : public facebook::jni::HybridClass<ExecuTorchLlmJni> {
       facebook::jni::alias_ref<ExecuTorchLlmCallbackJni> callback,
       jboolean echo) {
     if (model_type_category_ == MODEL_TYPE_CATEGORY_MULTIMODAL) {
-
+      std::vector<llm::MultimodalInput> inputs = prefill_inputs_;
+      prefill_inputs_.clear();
+      inputs.emplace_back(llm::MultimodalInput{prompt->toStdString()});
       return static_cast<jint>(multi_modal_runner_->generate(
-          std::vector<llm::MultimodalInput>{llm::MultimodalInput{prompt->toStdString()}},
-          llm::GenerationConfig {.echo = static_cast<bool>(echo), .seq_len = seq_len},
+          inputs,
+          llm::GenerationConfig{
+              .echo = static_cast<bool>(echo), .seq_len = seq_len},
           [callback](const std::string& result) { callback->onResult(result); },
           [callback](const llm::Stats& stats) { callback->onStats(stats); }));
     } else if (model_type_category_ == MODEL_TYPE_CATEGORY_LLM) {

extension/llm/runner/multimodal_runner.cpp

@@ -65,8 +65,8 @@ Error MultimodalRunner::load() {
 Error MultimodalRunner::generate(
     const std::vector<MultimodalInput>& inputs,
     const GenerationConfig& config,
-    std::function<void(const std::string&)>& token_callback,
-    std::function<void(const Stats&)>& stats_callback) {
+    std::function<void(const std::string&)> token_callback,
+    std::function<void(const Stats&)> stats_callback) {
   if (inputs.empty()) {
     ET_LOG(Error, "MultimodalInput vector cannot be empty");
     return Error::InvalidArgument;

extension/llm/runner/multimodal_runner.h

@@ -116,8 +116,8 @@ class ET_EXPERIMENTAL MultimodalRunner {
   virtual ::executorch::runtime::Error generate(
       const std::vector<MultimodalInput>& inputs,
       const GenerationConfig& config,
-      std::function<void(const std::string&)>& token_callback,
-      std::function<void(const Stats&)>& stats_callback);
+      std::function<void(const std::string&)> token_callback,
+      std::function<void(const Stats&)> stats_callback);
 
   inline void stop() {
     text_token_generator_->stop();