fix

Author: kirklandsign

extension/android/jni/jni_layer_llama.cpp

@@ -188,8 +188,7 @@ class ExecuTorchLlamaJni
       facebook::jni::alias_ref<jstring> prompt,
       jlong start_pos,
       jint bos,
-      jint eos,
-      jlong generated_token) {
+      jint eos) {
     facebook::jni::local_ref<jlongArray> tuple_result =
         facebook::jni::make_long_array(3);
     if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) {
@@ -201,9 +200,7 @@ class ExecuTorchLlamaJni
         prompt->toStdString(), start_pos, bos, eos);
     tuple_result->pin()[0] = static_cast<jint>(Error::Ok);
     if (result.ok()) {
-      // TODO(hsz): make  generated_token a reference and update it here
-      generated_token = result.get();
-      tuple_result->pin()[1] = static_cast<jlong>(generated_token);
+      tuple_result->pin()[1] = static_cast<jlong>(result.get());
       tuple_result->pin()[2] = static_cast<jlong>(start_pos);
     }
     return tuple_result;

extension/android/src/main/java/org/pytorch/executorch/LlamaModule.java

@@ -104,8 +104,8 @@ public native int generate(
    * @param eos The number of EOS (end of sequence) token.
    * @return a tuple of (error, token, updated startPos)
    */
-  public static native long[] prefill_prompt(
-      String prompt, long startPos, int bos, int eos, long generatedToken);
+  public static native long[] prefillPrompt(
+      String prompt, long startPos, int bos, int eos);
 
   /**
    * Prefill an LLaVA Module with the given images input.
@@ -117,7 +117,7 @@ public static native long[] prefill_prompt(
    * @param startPos The starting position in KV cache of the input in the LLM.
    * @return a tuple of (error code, updated startPos)
    */
-  public static native long[] prefill_images(
+  public static native long[] prefillImages(
       int[] image, int width, int height, int channels, long startPos);
 
   /**
@@ -129,8 +129,8 @@ public static native long[] prefill_images(
    * @param llamaCallback callback object to receive results.
    * @return The error code.
    */
-  public static native int generate_from_pos(
-      String prompt, int seqLen, long startPos, ExecuTorchLlamaCallback callback);
+  public static native int generateFromPos(
+      String prompt, int seqLen, long startPos, LlamaCallback callback);
 
   /** Stop current generate() before it finishes. */
   @DoNotStrip