[Android] New config API for Llm init and generate

extension/android/BUCK

@@ -25,7 +25,9 @@ non_fbcode_target(_kind = fb_android_library,
     name = "executorch_llama",
     srcs = [
         "executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmCallback.java",
+        "executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmGenerationConfig.java",
         "executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java",
+        "executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModuleConfig.java",
     ],
     autoglob = False,
     language = "JAVA",

extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmCallback.java

@@ -31,8 +31,7 @@ public interface LlmCallback {
   /**
    * Called when the statistics for the generate() is available.
    *
-   * The result will be a JSON string. See extension/llm/stats.h for the field
-   * definitions.
+   * <p>The result will be a JSON string. See extension/llm/stats.h for the field definitions.
    *
    * @param stats JSON string containing the statistics for the generate()
    */

extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmGenerationConfig.java

@@ -0,0 +1,156 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+package org.pytorch.executorch.extension.llm;
+
+/**
+ * Configuration class for controlling text generation parameters in LLM operations.
+ *
+ * <p>This class provides settings for text generation behavior including output formatting,
+ * generation limits, and sampling parameters. Instances should be created using the {@link
+ * #create()} method and the fluent builder pattern.
+ */
+public class LlmGenerationConfig {
+  private final boolean echo;
+  private final int maxNewTokens;
+  private final boolean warming;
+  private final int seqLen;
+  private final float temperature;
+
+  private LlmGenerationConfig(Builder builder) {
+    this.echo = builder.echo;
+    this.maxNewTokens = builder.maxNewTokens;
+    this.warming = builder.warming;
+    this.seqLen = builder.seqLen;
+    this.temperature = builder.temperature;
+  }
+
+  /**
+   * Creates a new Builder instance for constructing generation configurations.
+   *
+   * @return a new Builder with default configuration values
+   */
+  public static Builder create() {
+    return new Builder();
+  }
+
+  /**
+   * @return true if input prompt should be included in the output
+   */
+  public boolean isEcho() {
+    return echo;
+  }
+
+  /**
+   * @return maximum number of tokens to generate (-1 for unlimited)
+   */
+  public int getMaxNewTokens() {
+    return maxNewTokens;
+  }
+
+  /**
+   * @return true if model warming is enabled
+   */
+  public boolean isWarming() {
+    return warming;
+  }
+
+  /**
+   * @return maximum sequence length for generation (-1 for default)
+   */
+  public int getSeqLen() {
+    return seqLen;
+  }
+
+  /**
+   * @return temperature value for sampling (higher = more random)
+   */
+  public float getTemperature() {
+    return temperature;
+  }
+
+  /**
+   * Builder class for constructing LlmGenerationConfig instances.
+   *
+   * <p>Provides a fluent interface for configuring generation parameters with sensible defaults.
+   * All methods return the builder instance to enable method chaining.
+   */
+  public static class Builder {
+    private boolean echo = true;
+    private int maxNewTokens = -1;
+    private boolean warming = false;
+    private int seqLen = -1;
+    private float temperature = 0.8f;
+
+    Builder() {}
+
+    /**
+     * Sets whether to include the input prompt in the generated output.
+     *
+     * @param echo true to include input prompt, false to return only new tokens
+     * @return this builder instance
+     */
+    public Builder echo(boolean echo) {
+      this.echo = echo;
+      return this;
+    }
+
+    /**
+     * Sets the maximum number of new tokens to generate.
+     *
+     * @param maxNewTokens the token limit (-1 for unlimited generation)
+     * @return this builder instance
+     */
+    public Builder maxNewTokens(int maxNewTokens) {
+      this.maxNewTokens = maxNewTokens;
+      return this;
+    }
+
+    /**
+     * Enables or disables model warming.
+     *
+     * @param warming true to generate initial tokens for model warmup
+     * @return this builder instance
+     */
+    public Builder warming(boolean warming) {
+      this.warming = warming;
+      return this;
+    }
+
+    /**
+     * Sets the maximum sequence length for generation.
+     *
+     * @param seqLen maximum sequence length (-1 for default behavior)
+     * @return this builder instance
+     */
+    public Builder seqLen(int seqLen) {
+      this.seqLen = seqLen;
+      return this;
+    }
+
+    /**
+     * Sets the temperature for random sampling.
+     *
+     * @param temperature sampling temperature (typical range 0.0-1.0)
+     * @return this builder instance
+     */
+    public Builder temperature(float temperature) {
+      this.temperature = temperature;
+      return this;
+    }
+
+    /**
+     * Constructs the LlmGenerationConfig instance with the configured parameters.
+     *
+     * @return new LlmGenerationConfig instance with current builder settings
+     */
+    public LlmGenerationConfig build() {
+      return new LlmGenerationConfig(this);
+    }
+  }
+}

extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModule.java

@@ -15,8 +15,8 @@
 import org.pytorch.executorch.annotations.Experimental;
 
 /**
- * LlmModule is a wrapper around the Executorch LLM. It provides a simple interface to
- * generate text from the model.
+ * LlmModule is a wrapper around the Executorch LLM. It provides a simple interface to generate text
+ * from the model.
  *
  * <p>Warning: These APIs are experimental and subject to change without notice
  */

extension/android/executorch_android/src/main/java/org/pytorch/executorch/extension/llm/LlmModuleConfig.java

@@ -0,0 +1,165 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+package org.pytorch.executorch.extension.llm;
+
+/**
+ * Configuration class for initializing a LlmModule.
+ *
+ * <p>{@link #create()} method and the fluent builder pattern.
+ */
+public class LlmModuleConfig {
+  private final String modulePath;
+  private final String tokenizerPath;
+  private final float temperature;
+  private final String dataPath;
+  private final int modelType;
+
+  private LlmModuleConfig(Builder builder) {
+    this.modulePath = builder.modulePath;
+    this.tokenizerPath = builder.tokenizerPath;
+    this.temperature = builder.temperature;
+    this.dataPath = builder.dataPath;
+    this.modelType = builder.modelType;
+  }
+
+  /** Model type constant for text-only models. */
+  public static final int MODEL_TYPE_TEXT = 1;
+
+  /** Model type constant for text-and-vision multimodal models. */
+  public static final int MODEL_TYPE_TEXT_VISION = 2;
+
+  /**
+   * Creates a new Builder instance for constructing LlmModuleConfig objects.
+   *
+   * @return a new Builder instance with default configuration values
+   */
+  public static Builder create() {
+    return new Builder();
+  }
+
+  // Getters with documentation
+  /**
+   * @return Path to the compiled model module (.pte file)
+   */
+  public String getModulePath() {
+    return modulePath;
+  }
+
+  /**
+   * @return Path to the tokenizer file or directory
+   */
+  public String getTokenizerPath() {
+    return tokenizerPath;
+  }
+
+  /**
+   * @return Temperature value for sampling (higher = more random)
+   */
+  public float getTemperature() {
+    return temperature;
+  }
+
+  /**
+   * @return Optional path to additional data files
+   */
+  public String getDataPath() {
+    return dataPath;
+  }
+
+  /**
+   * @return Type of model (text-only or text-vision)
+   */
+  public int getModelType() {
+    return modelType;
+  }
+
+  /**
+   * Builder class for constructing LlmModuleConfig instances with optional parameters.
+   *
+   * <p>The builder provides a fluent interface for configuring model parameters and validates
+   * required fields before construction.
+   */
+  public static class Builder {
+    private String modulePath;
+    private String tokenizerPath;
+    private float temperature = 0.8f;
+    private String dataPath = "";
+    private int modelType = MODEL_TYPE_TEXT;
+
+    Builder() {}
+
+    /**
+     * Sets the path to the module.
+     *
+     * @param modulePath Path to module
+     * @return This builder instance for method chaining
+     */
+    public Builder modulePath(String modulePath) {
+      this.modulePath = modulePath;
+      return this;
+    }
+
+    /**
+     * Sets the path to the tokenizer.
+     *
+     * @param tokenizerPath Path to tokenizer
+     * @return This builder instance for method chaining
+     */
+    public Builder tokenizerPath(String tokenizerPath) {
+      this.tokenizerPath = tokenizerPath;
+      return this;
+    }
+
+    /**
+     * Sets the temperature for sampling generation.
+     *
+     * @param temperature Temperature value (typical range 0.0-1.0)
+     * @return This builder instance for method chaining
+     */
+    public Builder temperature(float temperature) {
+      this.temperature = temperature;
+      return this;
+    }
+
+    /**
+     * Sets the path to optional additional data files.
+     *
+     * @param dataPath Path to supplementary data resources
+     * @return This builder instance for method chaining
+     */
+    public Builder dataPath(String dataPath) {
+      this.dataPath = dataPath;
+      return this;
+    }
+
+    /**
+     * Sets the model type (text-only or multimodal).
+     *
+     * @param modelType One of MODEL_TYPE_TEXT or MODEL_TYPE_TEXT_VISION
+     * @return This builder instance for method chaining
+     */
+    public Builder modelType(int modelType) {
+      this.modelType = modelType;
+      return this;
+    }
+
+    /**
+     * Constructs the LlmModuleConfig instance with validated parameters.
+     *
+     * @return New LlmModuleConfig instance with configured values
+     * @throws IllegalArgumentException if required fields are missing
+     */
+    public LlmModuleConfig build() {
+      if (modulePath == null || tokenizerPath == null) {
+        throw new IllegalArgumentException("Module path and tokenizer path are required");
+      }
+      return new LlmModuleConfig(this);
+    }
+  }
+}