Update with generate_from_pos

Summary:

Test Plan:

Reviewers:

Subscribers:

Tasks:

Tags:

Author: larryliu0820

examples/models/llava/runner/llava_runner.cpp

@@ -93,6 +93,32 @@ Result<uint64_t> LlavaRunner::prefill_prompt(
   return text_prefiller_->prefill(prompt_tokens, start_pos);
 }
 
+Error LlavaRunner::generate_from_pos(
+    const std::string& prompt,
+    int32_t seq_len,
+    int64_t start_pos,
+    std::function<void(const std::string&)> token_callback,
+    std::function<void(const ::executorch::extension::llm::Stats&)>
+        stats_callback) {
+  // prefill user prompt. No BOS because preset prompt already has it.
+  token_callback(prompt);
+
+  uint64_t prefill_next_token =
+      ET_UNWRAP(prefill_prompt(prompt, start_pos, /*bos=*/0, /*eos*/ 0));
+  stats_.num_prompt_tokens = start_pos;
+
+  // Generate tokens
+  int64_t num_generated_tokens = ET_UNWRAP(text_token_generator_->generate(
+      {prefill_next_token}, start_pos, seq_len, token_callback));
+
+  // Bookkeeping
+  stats_.num_generated_tokens = num_generated_tokens;
+  ::executorch::llm::print_report(stats_);
+  if (stats_callback) {
+    stats_callback(stats_);
+  }
+}
+
 Error LlavaRunner::generate(
     std::vector<Image> images,
     const std::string& prompt,
@@ -122,25 +148,9 @@ Error LlavaRunner::generate(
   // prefill images
   prefill_images(images, pos);
 
-  // prefill user prompt. No BOS because preset prompt already has it.
-  wrapped_callback(prompt);
-
-  uint64_t prefill_next_token =
-      ET_UNWRAP(prefill_prompt(prompt, pos, /*bos=*/0, /*eos*/ 0));
-  stats_.num_prompt_tokens = pos;
-
   // Generate tokens
-  int64_t num_generated_tokens = ET_UNWRAP(text_token_generator_->generate(
-      {prefill_next_token}, pos, seq_len, wrapped_callback));
-
-  // Bookkeeping
-  stats_.num_generated_tokens = num_generated_tokens;
-  ::executorch::llm::print_report(stats_);
-  if (stats_callback) {
-    stats_callback(stats_);
-  }
-
-  return Error::Ok;
+  return generate_from_pos(
+      prompt, seq_len, pos, wrapped_callback, stats_callback);
 }
 
 } // namespace torch::executor

examples/models/llava/runner/llava_runner.h

@@ -62,6 +62,24 @@ class LlavaRunner : public MultimodalRunner {
       int8_t bos = 0,
       int8_t eos = 0);
 
+  /**
+   * Generate tokens from the given prompt, starting from the given position.
+   * @param prompt The text prompt to LLaVA.
+   * @param seq_len The total sequence length, including the prompt tokens and
+   * new tokens.
+   * @param start_pos The starting position in KV cache of the input in the LLM.
+   * @param token_callback What to do after a token is generated.
+   * @param stats_callback What to do with Stats.
+   * @return The error code.
+   */
+  Error generate_from_pos(
+      const std::string& prompt,
+      int32_t seq_len = 1024,
+      int64_t start_pos = 0,
+      std::function<void(const std::string&)> token_callback = {},
+      std::function<void(const ::executorch::extension::llm::Stats&)>
+          stats_callback = {});
+
  private:
   inline static const std::string kPresetPrompt =
       "A chat between a curious human and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the human's questions. USER: ";