Address comments about what to expose

Author: kirklandsign

extension/android/jni/jni_layer_llama.cpp

@@ -180,17 +180,16 @@ class ExecuTorchLlamaJni
     return 0;
   }
 
-  // Returns a tuple of (error, token, start_pos)
+  // Returns a tuple of (error, start_pos)
   // Contract is valid within an AAR (JNI + corresponding Java code)
-  // If the first element is not Error::Ok, the other two elements are
-  // undefined.
+  // If the first element is not Error::Ok, the other element is undefined.
   facebook::jni::local_ref<jlongArray> prefill_prompt(
       facebook::jni::alias_ref<jstring> prompt,
       jlong start_pos,
       jint bos,
       jint eos) {
     facebook::jni::local_ref<jlongArray> tuple_result =
-        facebook::jni::make_long_array(3);
+        facebook::jni::make_long_array(2);
     if (model_type_category_ != MODEL_TYPE_CATEGORY_MULTIMODAL) {
       tuple_result->pin()[0] = static_cast<jint>(Error::NotSupported);
       return tuple_result;
@@ -200,8 +199,7 @@ class ExecuTorchLlamaJni
         prompt->toStdString(), start_pos, bos, eos);
     tuple_result->pin()[0] = static_cast<jint>(Error::Ok);
     if (result.ok()) {
-      tuple_result->pin()[1] = static_cast<jlong>(result.get());
-      tuple_result->pin()[2] = static_cast<jlong>(start_pos);
+      tuple_result->pin()[1] = static_cast<jlong>(start_pos);
     }
     return tuple_result;
   }

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

@@ -113,7 +113,7 @@ public long prefillImages(int[] image, int width, int height, int channels, long
     return nativeResult[1];
   }
 
-  // returns a tuple of (error code, updated startPos)
+  // returns a tuple of (status, updated startPos)
   private native long[] prefillImagesNative(
       int[] image, int width, int height, int channels, long startPos);
 
@@ -133,10 +133,10 @@ public long prefillPrompt(String prompt, long startPos, int bos, int eos) {
     if (nativeResult[0] != 0) {
       throw new RuntimeException("Prefill failed with error code: " + nativeResult[0]);
     }
-    return nativeResult[2];
+    return nativeResult[1];
   }
 
-  // returns a tuple of (error, token, updated startPos)
+  // returns a tuple of (status, updated startPos)
   private native long[] prefillPromptNative(String prompt, long startPos, int bos, int eos);
 
   /**