Make seq_len param available in JNI layer generate() (#4745)

Summary:
Pull Request resolved: #4745

- Previously, on JNI side, we fixed the seq_len to 128.
- In this diff, we expose seq_len as a paramter in generate(), so devs can customize this as needed

Reviewed By: kirklandsign

Differential Revision: D61343892

Author: Riandy

extension/android/jni/jni_layer_llama.cpp

@@ -127,10 +127,11 @@ class ExecuTorchLlamaJni
 
   jint generate(
       facebook::jni::alias_ref<jstring> prompt,
+      jint seq_len,
       facebook::jni::alias_ref<ExecuTorchLlamaCallbackJni> callback) {
     runner_->generate(
         prompt->toStdString(),
-        128,
+        seq_len,
         [callback](std::string result) { callback->onResult(result); },
         [callback](const Stats& result) { callback->onStats(result); });
     return 0;

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

@@ -22,6 +22,7 @@ public class LlamaModule {
   }
 
   private final HybridData mHybridData;
+  private static final int DEFAULT_SEQ_LEN = 128;
 
   @DoNotStrip
   private static native HybridData initHybrid(
@@ -42,8 +43,19 @@ public void resetNative() {
    * @param prompt Input prompt
    * @param llamaCallback callback object to receive results.
    */
+  public int generate(String prompt, LlamaCallback llamaCallback) {
+    return generate(prompt, DEFAULT_SEQ_LEN, llamaCallback);
+  }
+
+  /**
+   * Start generating tokens from the module.
+   *
+   * @param prompt Input prompt
+   * @param seqLen sequence length
+   * @param llamaCallback callback object to receive results.
+   */
   @DoNotStrip
-  public native int generate(String prompt, LlamaCallback llamaCallback);
+  public native int generate(String prompt, int seqLen, LlamaCallback llamaCallback);
 
   /** Stop current generate() before it finishes. */
   @DoNotStrip