Repository: deepjavalibrary/djl
Pull request #2637 (2023-06-06 19:26:28)

[api] implements text-generation search algorithm
This PR succeeds PR #2547 #2509 #2557, #2572 which contains the benchmark outputs of the searching results. Model tracing is in #2509 's PR description.

This PR contains only the features of LMSearch. 

`djl/examples/src/main/java/ai/djl/examples/inference/GPTInference.java` contains the front_end design.

jawaff: I'm just leaving my 2 cents since I'm interested in your work. It's great to see it get added to DJL. My only real fear is that there might need to be a lot of refactoring to get support for other models. Flan-T5 is one of the most powerful open source models (that supports commercial use) we have available and it has a variety of sizes available. I'd be most interested in seeing it be supported.

I'm not super familiar with GPT2 aside from it being a decoder-only model. There's a chance that T5 and GPT2 share some similarities in the decoder aspect, but T5 has an initial encoder pass on the initial inputs. The hidden state of the encoded inputs are then used for each pass of the decoder alongside the ids that have been currently selected for generation.
jawaff: That's all I've got to add, good work. I just want to see this turn into a bigger feature beyond what you're working on.
frankfliu: > That's all I've got to add, good work. I just want to see this turn into a bigger feature beyond what you're working on.

We are planning to add T5 model. This is just a starting point to add textgeneration support.
KexinFeng: @jawaff Thanks for pointing out the encoder-decoder model T5 to us and reminding us of the possible refactoring. 

I think to implement encoder-decoder model, the major edition will be in the search algorithms, where we will need an if (encoderDecoder is true) block, which computes the encoding). The rest part of the code will basically be shared. This structure is seen in huggingface transformer.


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Patch coverage: **`54.85`**% and project coverage change: **`-0.05`** :warning:
> Comparison is base [(`bb5073f`)](https://app.codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.08% compared to head [(`c9c6fd4`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.04%.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2637      +/-   ##
============================================
- Coverage     72.08%   72.04%   -0.05%     
- Complexity     5126     7018    +1892     
============================================
  Files           473      698     +225     
  Lines         21970    31253    +9283     
  Branches       2351     3224     +873     
============================================
+ Hits          15838    22517    +6679     
- Misses         4925     7204    +2279     
- Partials       1207     1532     +325     
```


| [Impacted Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...djl/modality/cv/transform/RandomFlipLeftRight.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zZm9ybS9SYW5kb21GbGlwTGVmdFJpZ2h0LmphdmE=) | `25.00% <0.00%> (-25.00%)` | :arrow_down: |
| [...djl/modality/cv/transform/RandomFlipTopBottom.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zZm9ybS9SYW5kb21GbGlwVG9wQm90dG9tLmphdmE=) | `25.00% <0.00%> (-25.00%)` | :arrow_down: |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [.../modality/cv/translator/ImageFeatureExtractor.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW1hZ2VGZWF0dXJlRXh0cmFjdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `27.77% <0.00%> (+18.95%)` | :arrow_up: |
| [...ain/java/ai/djl/modality/cv/util/NDImageUtils.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3V0aWwvTkRJbWFnZVV0aWxzLmphdmE=) | `67.10% <0.00%> (+7.89%)` | :arrow_up: |
| [api/src/main/java/ai/djl/modality/nlp/Decoder.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9EZWNvZGVyLmphdmE=) | `63.63% <ø> (ø)` | |
| ... and [227 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

... and [368 files with indirect coverage changes](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2637?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/Application.java b/api/src/main/java/ai/djl/Application.java
index be245eca7..a0352e0d3 100644
--- a/api/src/main/java/ai/djl/Application.java
+++ b/api/src/main/java/ai/djl/Application.java
@@ -88,6 +88,8 @@ public class Application {
                 return NLP.TOKEN_CLASSIFICATION;
             case "nlp/word_embedding":
                 return NLP.WORD_EMBEDDING;
+            case "nlp/text_generation":
+                return NLP.TEXT_GENERATION;
             case "tabular":
                 return Tabular.ANY;
             case "tabular/linear_regression":
@@ -261,6 +263,8 @@ public class Application {
          */
         Application WORD_EMBEDDING = new Application("nlp/word_embedding");
 
+        Application TEXT_GENERATION = new Application("nlp/text_generation");
+
         /**
          * An application that translates text from one language to another.
          *
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/BatchTensorList.java b/api/src/main/java/ai/djl/modality/nlp/generate/BatchTensorList.java
new file mode 100644
index 000000000..6e0074fed
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/BatchTensorList.java
@@ -0,0 +1,174 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+/**
+ * BatchTensorList represents a search state, and the NDArrays inside are updated in each iteration
+ * of the autoregressive loop.
+ *
+ * <p>It is a struct consisting of NDArrays, whose first dimension is batch, and also contains
+ * sequence dimension (whose position in tensor's shape is specified by seqDimOrder). The SeqBatcher
+ * batch operations will operate on these two dimensions.
+ */
+public abstract class BatchTensorList {
+    // [batch, seq_past]. seq-dim-size == |past_seq| + |inputIds|. Will grow.
+    private NDArray pastOutputIds;
+
+    // [batch, seq_past]
+    // The cache of past attentionMask. seq-dim-size == |past_seq| + |inputIds|. Will grow.
+    private NDArray pastAttentionMask;
+
+    // (k, v) * numLayer,
+    // kv: [batch, heads, seq_past, kvfeature]
+    // The cache of past sequence. seq-dim-size == |past_seq| + |inputIds|. Will grow.
+    private NDList pastKeyValues;
+
+    // Sequence dimension order among all dimensions for each element in the batch list.
+    private long[] seqDimOrder;
+
+    BatchTensorList() {}
+
+    /**
+     * Constructs a new {@code BatchTensorList} instance.
+     *
+     * @param list the NDList that contains the serialized version of the batch tensors
+     * @param seqDimOrder the sequence dimension order that specifies where the sequence dimension
+     *     is in a tensor's shape
+     */
+    BatchTensorList(NDList list, long[] seqDimOrder) {
+        this.seqDimOrder = seqDimOrder;
+        pastOutputIds = list.get(0);
+        pastAttentionMask = list.get(1);
+        pastKeyValues = list.subNDList(2);
+    }
+
+    /**
+     * Constructs a new {@code BatchTensorList} instance.
+     *
+     * @param pastOutputIds past output token ids
+     * @param pastAttentionMask past attention mask
+     * @param pastKeyValues past kv cache
+     * @param seqDimOrder the sequence dimension order that specifies where the sequence dimension
+     *     is in a tensor's shape
+     */
+    BatchTensorList(
+            NDArray pastOutputIds,
+            NDArray pastAttentionMask,
+            NDList pastKeyValues,
+            long[] seqDimOrder) {
+        this.pastKeyValues = pastKeyValues;
+        this.pastOutputIds = pastOutputIds;
+        this.pastAttentionMask = pastAttentionMask;
+        this.seqDimOrder = seqDimOrder;
+    }
+
+    /**
+     * Constructs a new {@code BatchTensorList} instance from the serialized version of the batch
+     * tensors.
+     *
+     * <p>The pastOutputIds has to be the first in the output list.
+     *
+     * @param inputList the serialized version of the batch tensors
+     * @param seqDimOrder the sequence dimension order that specifies where the sequence dimension
+     *     is in a tensor's shape
+     * @return BatchTensorList
+     */
+    public abstract BatchTensorList fromList(NDList inputList, long[] seqDimOrder);
+
+    /**
+     * Returns the serialized version of the BatchTensorList. The pastOutputIds has to be the first
+     * in the output list.
+     *
+     * @return the NDList that contains the serialized BatchTensorList
+     */
+    public abstract NDList getList();
+
+    /**
+     * Returns the sequence dimension order which specifies where the sequence dimension is in a
+     * tensor's shape.
+     *
+     * @return the sequence dimension order which specifies where the sequence dimension is in a
+     *     tensor's shape
+     */
+    public long[] getSeqDimOrder() {
+        return seqDimOrder;
+    }
+
+    /**
+     * Returns the value of the pastOutputIds.
+     *
+     * @return the value of pastOutputIds
+     */
+    public NDArray getPastOutputIds() {
+        return pastOutputIds;
+    }
+
+    /**
+     * Sets the past output token ids.
+     *
+     * @param pastOutputIds the past output token ids
+     */
+    public void setPastOutputIds(NDArray pastOutputIds) {
+        this.pastOutputIds = pastOutputIds;
+    }
+
+    /**
+     * Returns the value of the pastAttentionMask.
+     *
+     * @return the value of pastAttentionMask
+     */
+    public NDArray getPastAttentionMask() {
+        return pastAttentionMask;
+    }
+
+    /**
+     * Sets the attention mask.
+     *
+     * @param pastAttentionMask the attention mask
+     */
+    public void setPastAttentionMask(NDArray pastAttentionMask) {
+        this.pastAttentionMask = pastAttentionMask;
+    }
+
+    /**
+     * Returns the value of the pastKeyValues.
+     *
+     * @return the value of pastKeyValues
+     */
+    public NDList getPastKeyValues() {
+        return pastKeyValues;
+    }
+
+    /**
+     * Sets the kv cache.
+     *
+     * @param pastKeyValues the kv cache
+     */
+    public void setPastKeyValues(NDList pastKeyValues) {
+        this.pastKeyValues = pastKeyValues;
+    }
+
+    /**
+     * Sets the sequence dimension order which specifies where the sequence dimension is in a
+     * tensor's shape.
+     *
+     * @param seqDimOrder the sequence dimension order which specifies where the sequence dimension
+     *     is in a tensor's shape
+     */
+    public void setSeqDimOrder(long[] seqDimOrder) {
+        this.seqDimOrder = seqDimOrder;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/BeamBatchTensorList.java b/api/src/main/java/ai/djl/modality/nlp/generate/BeamBatchTensorList.java
new file mode 100644
index 000000000..e22b430c7
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/BeamBatchTensorList.java
@@ -0,0 +1,128 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+class BeamBatchTensorList extends BatchTensorList {
+
+    // [batch, beam, seq=1]
+    private NDArray nextInputIds;
+
+    // [batch, beam]
+    private NDArray lastProbs;
+
+    // [batch, beam, seq_past + new_seq]
+    // The cache of past attentionMask. seq-dim-size == |past_seq| + |inputIds|. Will grow.
+    private NDArray pastAttentionMask;
+
+    /* Variables below are one time step behind the above state variables. Ie, they contain all the past sequence but excludes the time step that corresponds to the above input. */
+
+    // [batch, beam, seq_past]. seq-dim-size == |past_seq| + |inputIds|. Will grow.
+    private NDArray pastOutputIds;
+
+    // (k, v) * numLayer,
+    // kv: [batch, beam, heads, seq_past, kvfeature]
+    // The cache of past sequence. seq-dim-size == |past_seq| + |inputIds|. Will grow.
+    private NDList pastKeyValues;
+
+    BeamBatchTensorList() {}
+
+    BeamBatchTensorList(
+            NDArray nextInputIds,
+            NDArray pastOutputIds,
+            NDList pastKeyValues,
+            NDArray pastAttentionMask,
+            NDArray lastProb) {
+        this.nextInputIds = nextInputIds;
+        this.pastKeyValues = pastKeyValues;
+        this.pastOutputIds = pastOutputIds;
+        this.pastAttentionMask = pastAttentionMask;
+        this.lastProbs = lastProb;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public BatchTensorList fromList(NDList inputList, long[] seqDimOrder) {
+        return new BeamBatchTensorList();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList getList() {
+        return new NDList();
+    }
+
+    /**
+     * Returns the value of the nextInputIds.
+     *
+     * @return the value of nextInputIds
+     */
+    public NDArray getNextInputIds() {
+        return nextInputIds;
+    }
+
+    public void setNextInputIds(NDArray nextInputIds) {
+        this.nextInputIds = nextInputIds;
+    }
+
+    /**
+     * Returns the value of the lastProbs.
+     *
+     * @return the value of lastProbs
+     */
+    public NDArray getLastProbs() {
+        return lastProbs;
+    }
+
+    public void setLastProbs(NDArray lastProbs) {
+        this.lastProbs = lastProbs;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray getPastAttentionMask() {
+        return pastAttentionMask;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void setPastAttentionMask(NDArray pastAttentionMask) {
+        this.pastAttentionMask = pastAttentionMask;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray getPastOutputIds() {
+        return pastOutputIds;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void setPastOutputIds(NDArray pastOutputIds) {
+        this.pastOutputIds = pastOutputIds;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList getPastKeyValues() {
+        return pastKeyValues;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void setPastKeyValues(NDList pastKeyValues) {
+        this.pastKeyValues = pastKeyValues;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/CausalLMOutput.java b/api/src/main/java/ai/djl/modality/nlp/generate/CausalLMOutput.java
new file mode 100644
index 000000000..c6d8d2740
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/CausalLMOutput.java
@@ -0,0 +1,94 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+/** CausalLMOuput is used to contain multiple output of a language model. */
+public class CausalLMOutput {
+
+    // [batch, seq, feature]
+    // The prob. conditional on a sequence that ends at an element in seq-dim. seq-dim-size =
+    // |inputIds|
+    private NDArray logits;
+
+    // [batch, seq, dim] * (layers+1) -> take -1
+    // The vec. rep. of a sequence that ends at an element in seq-dim. seq-dim-size = |inputIds|
+    private NDArray hiddenStates;
+
+    // (k, v) * numLayer,
+    // kv: [batch, heads, seq_past, feature]
+    // The cache of past sequence. seq-dim-size == |seq_past| + |inputIds|
+    private NDList pastKeyValuesList;
+
+    /**
+     * Constructs a new {@code CausalLMOutput} instance.
+     *
+     * @param logits the logits NDArray
+     * @param pastKeyValues the key-value cache
+     */
+    public CausalLMOutput(NDArray logits, NDList pastKeyValues) {
+        this.logits = logits;
+        this.pastKeyValuesList = pastKeyValues;
+    }
+
+    /**
+     * Constructs a new {@code CausalLMOutput} intance.
+     *
+     * @param logits the logits NDArray
+     * @param hiddenState the first layer hiddenStates used as word embedding
+     * @param pastKeyValueList the key-value cache
+     */
+    public CausalLMOutput(NDArray logits, NDArray hiddenState, NDList pastKeyValueList) {
+        this.logits = logits;
+        this.pastKeyValuesList = pastKeyValueList;
+        this.hiddenStates = hiddenState;
+    }
+
+    /**
+     * Returns the value of the logits.
+     *
+     * @return the value of logits
+     */
+    public NDArray getLogits() {
+        return logits;
+    }
+
+    /**
+     * Sets the value of the logits.
+     *
+     * @param logits value of logits NDArray
+     */
+    public void setLogits(NDArray logits) {
+        this.logits = logits;
+    }
+
+    /**
+     * Returns the value of the allHiddenStates.
+     *
+     * @return the value of allHiddenStates
+     */
+    public NDArray getHiddenState() {
+        return hiddenStates;
+    }
+
+    /**
+     * Returns the value of the pastKeyValuesList.
+     *
+     * @return the value of pastKeyValuesList
+     */
+    public NDList getPastKeyValuesList() {
+        return pastKeyValuesList;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/ContrastiveBatchTensorList.java b/api/src/main/java/ai/djl/modality/nlp/generate/ContrastiveBatchTensorList.java
new file mode 100644
index 000000000..1e4b41157
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/ContrastiveBatchTensorList.java
@@ -0,0 +1,89 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+class ContrastiveBatchTensorList extends BatchTensorList {
+    // [batch, seq_past, hiddenDim]
+    // The embed vector of the past seq. seq-dim-size = |past_seq|. Will grow.
+    private NDArray pastHiddenStates;
+
+    // [batch, vacabSize]. Only the last logits, used to recall candidate token.
+    private NDArray logits;
+
+    ContrastiveBatchTensorList(NDList list, long[] seqDimOrder) {
+        super(list.get(0), list.get(1), list.subNDList(4), seqDimOrder);
+        pastHiddenStates = list.get(2);
+        logits = list.get(3);
+    }
+
+    ContrastiveBatchTensorList(
+            NDArray pastOutputIds,
+            NDArray pastAttentionMask,
+            NDArray pastHiddenStates,
+            NDArray logits,
+            NDList pastKeyValues,
+            long[] seqDimOrder) {
+        super(pastOutputIds, pastAttentionMask, pastKeyValues, seqDimOrder);
+        this.pastHiddenStates = pastHiddenStates;
+        this.logits = logits;
+    }
+
+    public ContrastiveBatchTensorList() {}
+
+    /** {@inheritDoc} */
+    @Override
+    public ContrastiveBatchTensorList fromList(NDList inputList, long[] seqDimOrder) {
+        return new ContrastiveBatchTensorList(inputList, seqDimOrder);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList getList() {
+        // The pastOutputIds has to be the first in the output list
+        return new NDList(
+                        getPastOutputIds(),
+                        getPastAttentionMask(),
+                        getPastHiddenStates(),
+                        getLogits())
+                .addAll(getPastKeyValues());
+    }
+
+    /**
+     * Returns the value of the pastHiddenStates.
+     *
+     * @return the value of pastHiddenStates
+     */
+    public NDArray getPastHiddenStates() {
+        return pastHiddenStates;
+    }
+
+    public void setPastHiddenStates(NDArray pastHiddenStates) {
+        this.pastHiddenStates = pastHiddenStates;
+    }
+
+    /**
+     * Returns the value of the logits.
+     *
+     * @return the value of logits
+     */
+    public NDArray getLogits() {
+        return logits;
+    }
+
+    public void setLogits(NDArray logits) {
+        this.logits = logits;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/ContrastiveSeqBatchScheduler.java b/api/src/main/java/ai/djl/modality/nlp/generate/ContrastiveSeqBatchScheduler.java
new file mode 100644
index 000000000..61229d5bf
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/ContrastiveSeqBatchScheduler.java
@@ -0,0 +1,234 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.inference.Predictor;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDScope;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.translate.TranslateException;
+
+import java.util.Arrays;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * {@code ContrastiveSeqBatchScheduler} is a class which implements the contrastive search algorithm
+ * used in SeqBatchScheduler.
+ */
+public class ContrastiveSeqBatchScheduler extends SeqBatchScheduler {
+
+    /**
+     * Constructs a new {@code ContrastiveSeqBatchScheduler} instance.
+     *
+     * @param lmBlock the predictor containing language model
+     * @param config the autoregressive search configuration
+     */
+    public ContrastiveSeqBatchScheduler(
+            Predictor<NDList, CausalLMOutput> lmBlock, SearchConfig config) {
+        super(lmBlock, config);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public SeqBatcher initForward(NDArray inputIds, NDArray batchUids) throws TranslateException {
+        try (NDScope scope = new NDScope()) {
+            scope.suppressNotUsedWarning();
+            manager = inputIds.getManager();
+            NDArray initOffSets = computeOffSets(inputIds, config);
+            NDArray attentionMask = computeAttentionMask(inputIds, config);
+            NDArray positionIds = computePositionIds(inputIds, initOffSets, 0, 1);
+
+            CausalLMOutput output =
+                    predictor.predict(new NDList(inputIds, positionIds, attentionMask));
+            NDArray lastLogits = output.getLogits().get(":, -1, :");
+
+            // Used to mark the sequence dimension's ordinal number for each tensor in the
+            // serialized
+            // batchTensorList
+            long[] seqDimOrder = new long[28];
+            Arrays.fill(seqDimOrder, 0, 3, 1);
+            seqDimOrder[3] = -1; // -1 means no sequence dimension
+            Arrays.fill(seqDimOrder, 4, seqDimOrder.length, 2);
+
+            BatchTensorList batchTensorList =
+                    new ContrastiveBatchTensorList(
+                            inputIds,
+                            attentionMask,
+                            output.getHiddenState(),
+                            lastLogits,
+                            output.getPastKeyValuesList(),
+                            seqDimOrder);
+            SeqBatcher ret = new SeqBatcher(batchTensorList, batchUids, initOffSets, manager);
+
+            // memory management
+            NDScope.unregister(output.getPastKeyValuesList());
+            NDScope.unregister(output.getHiddenState(), attentionMask, lastLogits);
+            NDScope.unregister(ret.offSets, ret.batchUid);
+
+            return ret;
+        }
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray inferenceCall() throws TranslateException {
+        NDArray outputIds;
+        try (NDScope scope = new NDScope()) {
+            scope.suppressNotUsedWarning();
+
+            /* Prepare input for one inference call */
+            NDArray logits = ((ContrastiveBatchTensorList) seqBatcher.getData()).getLogits();
+            NDArray topKIds = logits.topK(config.getK(), -1, true, false).get(1); // [batch, topK]
+            ContrastiveBatchTensorList searchState = (ContrastiveBatchTensorList) seqBatcher.data;
+
+            // Embed the topk dimension into batch dimension for an inference all
+            // [batch, topK] -> [batch * [topK]] -> [[batch * [topK]], seqLength=1]
+            NDArray candidateInputIds = topKIds.flatten().reshape(-1, 1);
+            assert candidateInputIds.getDataType() == DataType.INT64
+                    : "inputIds datatype should be int64";
+            assert candidateInputIds.getShape().getShape().length == 2 : "shape not right";
+
+            // [batch, heads, seq_past, feature] -> [batch * topK, head, seq_past, feature]
+            NDList kCopyPastKeyValues =
+                    new NDList(
+                            searchState.getPastKeyValues().stream()
+                                    .map(ndarray -> ndarray.repeat(0, config.getK()))
+                                    .collect(Collectors.toList()));
+            assert kCopyPastKeyValues.get(0).getDataType() == DataType.FLOAT32
+                    : "inputIds datatype should be Float32";
+
+            // [batch, seq_past] -> [batch * topK, seq_past] -> [batch * topK, seq_past + 1]
+            long numBatch = topKIds.getShape().get(0);
+            NDArray kCopyPastAttentionMask =
+                    searchState.getPastAttentionMask().repeat(0, config.getK());
+            kCopyPastAttentionMask =
+                    kCopyPastAttentionMask.concat(
+                            manager.ones(new Shape(numBatch * config.getK(), 1), DataType.INT64),
+                            1);
+            assert kCopyPastKeyValues.get(0).getShape().get(2) + 1
+                            == kCopyPastAttentionMask.getShape().getLastDimension()
+                    : "attentionMask_seq = past_seq + new_input_seq";
+
+            // Forward with candidates in batch input
+            NDArray candidatePositionIds =
+                    computePositionIds(
+                            candidateInputIds,
+                            seqBatcher.offSets,
+                            searchState.getPastOutputIds().getShape().getLastDimension(),
+                            config.getK());
+            NDList modelInputs =
+                    new NDList(candidateInputIds, candidatePositionIds, kCopyPastAttentionMask);
+            modelInputs.addAll(kCopyPastKeyValues);
+            CausalLMOutput candidateOutput = predictor.predict(modelInputs);
+
+            NDList generatedOutput =
+                    StepGeneration.constrastiveStepGeneration(
+                            topKIds,
+                            logits,
+                            searchState.getPastHiddenStates(),
+                            candidateOutput.getHiddenState(),
+                            seqBatcher.offSets,
+                            config.getAlpha());
+
+            /* Update searchState for next loop */
+            long logitsDim = logits.getShape().get(1);
+            long numHeads = searchState.getPastKeyValues().get(0).getShape().get(1);
+            long kvDim = searchState.getPastKeyValues().get(0).getShape().get(3);
+            long currentSeqLength = searchState.getPastOutputIds().getShape().get(1);
+            long hiddenDim = searchState.getPastHiddenStates().getShape().get(2);
+
+            // [batch, 1]
+            NDArray select = generatedOutput.get(1);
+            NDIndex selectIndex =
+                    new NDIndex(
+                            "{}, {}, ...",
+                            manager.arange(0, numBatch, 1, DataType.INT64),
+                            select.flatten());
+
+            // Take from candidateOutput
+            // [batch, k, inputSeq=1, logitsDim] --select--> [batch, logitDim]
+            NDArray nextLogits =
+                    candidateOutput
+                            .getLogits()
+                            .reshape(numBatch, config.getK(), logitsDim)
+                            .get(selectIndex);
+
+            // Take from candidateOutput
+            // [batch * k, heads, seq_past, feature] --select--> [batch, heads, seq_past, feature]
+            Function<NDArray, NDArray> fn =
+                    ndarray ->
+                            ndarray.reshape(
+                                            numBatch,
+                                            config.getK(),
+                                            numHeads,
+                                            currentSeqLength + 1,
+                                            kvDim)
+                                    .get(selectIndex);
+            NDList nextPastKeyValue =
+                    new NDList(
+                            candidateOutput.getPastKeyValuesList().stream()
+                                    .map(fn)
+                                    .collect(Collectors.toList()));
+
+            // To be concatenated into searchState.pastHiddenStates
+            // [batch * k, inputSeq=1, hiddenDim]
+            NDArray newHiddenState = candidateOutput.getHiddenState();
+            assert newHiddenState.getManager() == manager : "possible leaky memory";
+            NDArray nextPastHiddenStates =
+                    searchState
+                            .getPastHiddenStates()
+                            .concat(
+                                    newHiddenState
+                                            .reshape(numBatch, config.getK(), 1, hiddenDim)
+                                            .get(selectIndex),
+                                    1);
+
+            // To be concatenated into searchState.outputIds
+            // [batch, seq_past]
+            outputIds = generatedOutput.get(0);
+            NDArray nextOutputIds = searchState.getPastOutputIds().concat(outputIds, 1);
+
+            // [batch, seq_past]
+            NDArray nextPastAttentionMask =
+                    searchState
+                            .getPastAttentionMask()
+                            .concat(manager.ones(new Shape(numBatch, 1), DataType.INT64), 1);
+
+            seqBatcher.seqLength++;
+            seqBatcher.data =
+                    new ContrastiveBatchTensorList(
+                            nextOutputIds,
+                            nextPastAttentionMask,
+                            nextPastHiddenStates,
+                            nextLogits,
+                            nextPastKeyValue,
+                            searchState.getSeqDimOrder());
+
+            /* Exit criteria */
+            seqBatcher.exitCriteria(outputIds, config.getMaxSeqLength(), config.getEosTokenId());
+
+            // Memory management
+            NDScope.unregister(nextOutputIds);
+            NDScope.unregister(nextPastAttentionMask);
+            NDScope.unregister(nextPastHiddenStates);
+            NDScope.unregister(nextLogits);
+            NDScope.unregister(nextPastKeyValue);
+            NDScope.unregister(outputIds);
+        }
+        return outputIds;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/GreedyBatchTensorList.java b/api/src/main/java/ai/djl/modality/nlp/generate/GreedyBatchTensorList.java
new file mode 100644
index 000000000..28bceb2df
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/GreedyBatchTensorList.java
@@ -0,0 +1,109 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+class GreedyBatchTensorList extends BatchTensorList {
+    // [batch, 1]
+    private NDArray nextInputIds;
+
+    // [batch, seq_past + new_seq]
+    // The cache of past attentionMask. seq-dim-size == |past_seq| + |inputIds|. Will grow.
+    private NDArray pastAttentionMask;
+
+    /* Variables below are one time step behind the above state variables. Ie, they contain all the past sequence but excludes the time step that corresponds to the above input. */
+
+    // [batch, seq_past]. seq-dim-size == |past_seq| + |inputIds|. Will grow.
+    private NDArray pastOutputIds;
+
+    // (k, v) * numLayer,
+    // kv: [batch, heads, seq_past, kvfeature]
+    // The cache of past sequence. seq-dim-size == |past_seq| + |inputIds|. Will grow.
+    private NDList pastKeyValues;
+
+    GreedyBatchTensorList(
+            NDArray nextInputIds,
+            NDArray pastOutputIds,
+            NDList pastKeyValues,
+            NDArray pastAttentionMask) {
+        this.nextInputIds = nextInputIds;
+        this.pastKeyValues = pastKeyValues;
+        this.pastOutputIds = pastOutputIds;
+        this.pastAttentionMask = pastAttentionMask;
+    }
+
+    public GreedyBatchTensorList() {}
+
+    /** {@inheritDoc} */
+    @Override
+    public BatchTensorList fromList(NDList inputList, long[] seqDimOrder) {
+        return new GreedyBatchTensorList();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList getList() {
+        return new NDList();
+    }
+
+    /**
+     * Returns the value of the nextInputIds.
+     *
+     * @return the value of nextInputIds
+     */
+    public NDArray getNextInputIds() {
+        return nextInputIds;
+    }
+
+    public void setNextInputIds(NDArray nextInputIds) {
+        this.nextInputIds = nextInputIds;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray getPastAttentionMask() {
+        return pastAttentionMask;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void setPastAttentionMask(NDArray pastAttentionMask) {
+        this.pastAttentionMask = pastAttentionMask;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray getPastOutputIds() {
+        return pastOutputIds;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void setPastOutputIds(NDArray pastOutputIds) {
+        this.pastOutputIds = pastOutputIds;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList getPastKeyValues() {
+        return pastKeyValues;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void setPastKeyValues(NDList pastKeyValues) {
+        this.pastKeyValues = pastKeyValues;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/SearchConfig.java b/api/src/main/java/ai/djl/modality/nlp/generate/SearchConfig.java
new file mode 100644
index 000000000..b42fe826b
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/SearchConfig.java
@@ -0,0 +1,155 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+/**
+ * {@code SearchConfig} is a class whose fields are parameters used for autoregressive search / text
+ * generation.
+ */
+public class SearchConfig {
+
+    private int k;
+    private float alpha;
+    private int beam;
+    private int maxSeqLength;
+    private long padTokenId;
+    private long eosTokenId;
+    private boolean suffixPadding;
+
+    /** Constructs a new {@code ContrastiveSearchConfig} instance with default values. */
+    public SearchConfig() {
+        this.k = 4;
+        this.alpha = 0.6f;
+        this.beam = 3;
+        this.maxSeqLength = 30;
+        this.eosTokenId = 50256;
+        this.padTokenId = 50256;
+    }
+
+    /**
+     * Returns the value of the k.
+     *
+     * @return the value of k
+     */
+    public int getK() {
+        return k;
+    }
+
+    /**
+     * Sets the value for the topk choice.
+     *
+     * @param k the value for topk choice
+     */
+    public void setK(int k) {
+        this.k = k;
+    }
+
+    /**
+     * Returns the value of the alpha.
+     *
+     * @return the value of alpha
+     */
+    public float getAlpha() {
+        return alpha;
+    }
+
+    /**
+     * Sets the value of alpha the penalty for repetition.
+     *
+     * @param alpha the value of the penalty for repetition
+     */
+    public void setAlpha(float alpha) {
+        this.alpha = alpha;
+    }
+
+    /**
+     * Returns the value of the beam.
+     *
+     * @return the value of beam
+     */
+    public int getBeam() {
+        return beam;
+    }
+
+    /**
+     * Sets the value of beam size.
+     *
+     * @param beam the value of beam size
+     */
+    public void setBeam(int beam) {
+        this.beam = beam;
+    }
+
+    /**
+     * Returns the value of the maxSeqLength.
+     *
+     * @return the value of maxSeqLength
+     */
+    public int getMaxSeqLength() {
+        return maxSeqLength;
+    }
+
+    /**
+     * Sets the value of max sequence length.
+     *
+     * @param maxSeqLength the value max sequence length
+     */
+    public void setMaxSeqLength(int maxSeqLength) {
+        this.maxSeqLength = maxSeqLength;
+    }
+
+    /**
+     * Returns the value of the padTokenId.
+     *
+     * @return the value of padTokenId
+     */
+    public long getPadTokenId() {
+        return padTokenId;
+    }
+
+    /**
+     * Sets the value of padTokenId.
+     *
+     * @param padTokenId the token id for padding
+     */
+    public void setPadTokenId(long padTokenId) {
+        this.padTokenId = padTokenId;
+    }
+
+    /**
+     * Returns the value of the eosTokenId.
+     *
+     * @return the value of eosTokenId
+     */
+    public long getEosTokenId() {
+        return eosTokenId;
+    }
+
+    /**
+     * Returns the value of the suffixPadding.
+     *
+     * @return the value of suffixPadding
+     */
+    public boolean isSuffixPadding() {
+        return suffixPadding;
+    }
+
+    /**
+     * Sets the value of suffixPadding or rightPadding.
+     *
+     * @param suffixPadding whether the padding is from right
+     */
+    public void setSuffixPadding(boolean suffixPadding) {
+        this.suffixPadding = suffixPadding;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatchScheduler.java b/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatchScheduler.java
new file mode 100644
index 000000000..78f40c5b2
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatchScheduler.java
@@ -0,0 +1,224 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.inference.Predictor;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.translate.TranslateException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * This is a scheduler, serving as an API to the consumer of the system, allowing for three major
+ * actions: initForward, addBatch, fastForward, collectResults. An optimal control sequence should
+ * be solved, after considering the time consumption of each action, the batch size and sequence
+ * length of queueing requests. Such optimal control solver needs additional effort. Primitive
+ * policy is setting several thresholds.
+ */
+public abstract class SeqBatchScheduler {
+    private static final Logger logger = LoggerFactory.getLogger(SeqBatchScheduler.class);
+
+    Predictor<NDList, CausalLMOutput> predictor;
+    SeqBatcher seqBatcher;
+
+    NDManager manager;
+
+    SearchConfig config;
+
+    Map<Long, NDArray> results;
+
+    /**
+     * Constructs a new {@code SeqBatchScheduler} instance.
+     *
+     * @param lmBlock the predictor that cont
+     * @param config the search parameter configuration
+     */
+    public SeqBatchScheduler(Predictor<NDList, CausalLMOutput> lmBlock, SearchConfig config) {
+        this.predictor = lmBlock;
+        this.config = config;
+        results = new ConcurrentHashMap<>();
+    }
+
+    /**
+     * Initializes the iteration and SeqBatcher.
+     *
+     * @param inputIds the input token ids.
+     * @param batchUids the request uid identifying a sequence
+     * @return SeqBatcher Stores the search state and operate on the BatchTensorList
+     * @throws TranslateException if forward fails
+     */
+    public abstract SeqBatcher initForward(NDArray inputIds, NDArray batchUids)
+            throws TranslateException;
+
+    /**
+     * Executes forward for a given number of iterations.
+     *
+     * @param count the time of forward calls
+     * @return boolean Indicate whether the Batch is empty
+     * @throws TranslateException if forward fails
+     */
+    public boolean incrementForward(int count) throws TranslateException {
+        int i = 0;
+        while (i++ < count) {
+            if (seqBatcher == null || seqBatcher.getData() == null) {
+                logger.info(
+                        "seqBatcher not set or is empty. Please call addBatch. Current inference"
+                                + " times is "
+                                + i);
+                return true;
+            }
+
+            inferenceCall();
+            if (seqBatcher.sequenceComplete()) {
+                results.putAll(seqBatcher.collectAndTrim());
+            }
+        }
+        return false;
+    }
+
+    /**
+     * An inference call in an iteration.
+     *
+     * @return the output token ids
+     * @throws TranslateException if forward fails
+     */
+    abstract NDArray inferenceCall() throws TranslateException;
+
+    /**
+     * Adds new batch.
+     *
+     * @param inputIds the input token ids.
+     * @param batchUids the request uid identifying a sequence
+     * @throws TranslateException if forward fails
+     */
+    public void addRequest(NDArray inputIds, NDArray batchUids) throws TranslateException {
+        SeqBatcher seqBatcherNew = initForward(inputIds, batchUids);
+        if (seqBatcher == null) {
+            seqBatcher = seqBatcherNew;
+        } else {
+            seqBatcher.addBatch(seqBatcherNew);
+        }
+    }
+
+    /**
+     * Collects finished results.
+     *
+     * @return the outputs stored as a map from requestUid to output token ids
+     */
+    public Map<Long, NDArray> collectResults() {
+        Map<Long, NDArray> output = results;
+        results = new ConcurrentHashMap<>();
+        return output;
+    }
+
+    /**
+     * Computes the offSets by linear search from the left.
+     *
+     * @param inputIds input token ids
+     * @param config search configuration
+     * @return the offsets NDArray
+     */
+    static NDArray computeOffSets(NDArray inputIds, SearchConfig config) {
+        int numBatch = Math.toIntExact(inputIds.getShape().get(0));
+        int initSeqSize = Math.toIntExact(inputIds.getShape().get(1));
+
+        // Linear search from left to find the first position that's not padTokenId.
+        long[] offSetsArray = new long[numBatch];
+        for (int i = 0; i < numBatch; i++) {
+            long[] aSequence = inputIds.get("{},:", i).toLongArray();
+            int idx = 0;
+            while (idx < initSeqSize) {
+                if (aSequence[idx] != config.getPadTokenId()) {
+                    break;
+                }
+                idx++;
+            }
+            offSetsArray[i] = idx;
+        }
+
+        NDManager manager = inputIds.getManager();
+        return manager.create(offSetsArray).reshape(-1, 1);
+    }
+
+    /**
+     * Computes the attention mask by linear search from the left.
+     *
+     * @param inputIds input token ids
+     * @param config search configuration
+     * @return the attention mask NDArray
+     */
+    static NDArray computeAttentionMask(NDArray inputIds, SearchConfig config) {
+        int numBatch = Math.toIntExact(inputIds.getShape().get(0));
+        int initSeqSize = Math.toIntExact(inputIds.getShape().get(1));
+
+        NDManager manager = inputIds.getManager();
+        NDArray attentionMask =
+                manager.ones(new Shape(1, inputIds.getShape().getLastDimension()), DataType.INT64)
+                        .reshape(1, -1)
+                        .repeat(0, numBatch);
+
+        // Linear search to find the offset and set the mask
+        for (int i = 0; i < numBatch; i++) {
+            long[] aSequence = inputIds.get("{},:", i).toLongArray();
+            int idx = 0;
+            while (idx < initSeqSize) {
+                if (aSequence[idx] != config.getPadTokenId()) {
+                    break;
+                }
+                idx++;
+            }
+            attentionMask.set(new NDIndex("{},{}:{}", i, 0, idx), 0);
+        }
+
+        // [batch, pastSeq]
+        return attentionMask;
+    }
+
+    /**
+     * Computes the position ids by linear search from the left.
+     *
+     * @param inputIds input token ids
+     * @param offSets the offset
+     * @param pastSeqLength past sequence length
+     * @param repeat the number of repeats used in interleave-repeating the position_ids to multiple
+     *     rows
+     * @return the position ids NDArray
+     */
+    static NDArray computePositionIds(
+            NDArray inputIds, NDArray offSets, long pastSeqLength, int repeat) {
+        NDManager manager = inputIds.getManager();
+        NDArray positionIds =
+                manager.arange(
+                                pastSeqLength,
+                                pastSeqLength + inputIds.getShape().getLastDimension(),
+                                1,
+                                DataType.INT64)
+                        .expandDims(0)
+                        .repeat(0, inputIds.getShape().get(0));
+
+        NDArray positionIdsShifted = positionIds.subi(offSets.reshape(-1, 1).repeat(0, repeat));
+        positionIds = positionIdsShifted.maximum(positionIdsShifted.zerosLike());
+
+        // [batch, inputSeq]
+        return positionIds;
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatcher.java b/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatcher.java
new file mode 100644
index 000000000..3a1fc9e4a
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/SeqBatcher.java
@@ -0,0 +1,301 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.NDScope;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.types.Shape;
+
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * {@code SeqBatcher} stores the search state (BatchTensorList), the control variables (e.g.
+ * seqLength, offSets, etc), and batch operations (merge, trim, exitCriteria, etc) on
+ * BatchTensorList.
+ */
+public class SeqBatcher {
+
+    NDManager manager;
+    long batchSize;
+    long seqLength;
+
+    /** [batch] stores the uid and is trimmed or enhanced correspondingly to the batch. */
+    NDArray batchUid;
+
+    /** Minheap with lazy removal, map: batchIdx -> offset. */
+    NDArray offSets;
+
+    /** This is a struct that contains NDArrays with batch dimension. */
+    BatchTensorList data;
+
+    /** batchIndex -> seqEndPosition. */
+    private Map<Long, Long> exitIndexEndPosition;
+
+    SeqBatcher(BatchTensorList data, NDArray batchUid, NDArray offSets, NDManager manager) {
+        this.manager = manager.newSubManager();
+        this.data = data;
+        this.batchUid = batchUid.getShape().dimension() == 2 ? batchUid : batchUid.reshape(-1, 1);
+        this.offSets = offSets.getShape().hashCode() == 2 ? offSets : offSets.reshape(-1, 1);
+        batchSize = data.getPastOutputIds().getShape().get(0);
+        seqLength = data.getPastOutputIds().getShape().get(1);
+        exitIndexEndPosition = new ConcurrentHashMap<>();
+    }
+
+    /**
+     * Returns the batch data which is stored as a {@code BatchTensorList}.
+     *
+     * @return the batch data stored as BatchTensorList
+     */
+    public BatchTensorList getData() {
+        return data;
+    }
+
+    /**
+     * Adds new batch.
+     *
+     * <p>Modify the batch dimension and the left padding.
+     *
+     * @param seqBatcherNew the seqBatcher to add.
+     */
+    public void addBatch(SeqBatcher seqBatcherNew) {
+        merge(this, seqBatcherNew, seqLength - seqBatcherNew.seqLength);
+        // manager and finishedSequences stay the same;
+    }
+
+    /**
+     * Merges two batchers together.
+     *
+     * <p>Modify the batch dimension and the left padding.
+     *
+     * @param seqBatcher1 the first seqBatcher
+     * @param seqBatcher2 the second seqBatcher
+     * @param seqDelta the sequence length difference
+     */
+    private void merge(SeqBatcher seqBatcher1, SeqBatcher seqBatcher2, long seqDelta) {
+        if (seqDelta < 0) {
+            SeqBatcher swapTmp = seqBatcher1;
+            seqBatcher1 = seqBatcher2;
+            seqBatcher2 = swapTmp;
+            seqDelta = -seqDelta;
+        }
+
+        try (NDScope scope = new NDScope()) {
+            scope.suppressNotUsedWarning();
+
+            NDList list1 = seqBatcher1.data.getList();
+            NDList list2 = seqBatcher2.data.getList();
+            NDList merged = new NDList(list1.size());
+            long[] seqDimOrder = seqBatcher1.data.getSeqDimOrder();
+            for (int i = 0; i < list1.size(); i++) {
+                NDArray batch1 = list1.get(i);
+                NDArray batch2 = list2.get(i);
+                if (seqDelta == 0) {
+                    // no need to pad
+                    batch1 = batch1.concat(batch2, 0);
+                    merged.add(batch1);
+                    continue;
+                }
+
+                long[] shape1 = batch1.getShape().getShape();
+                long[] shape2 = batch2.getShape().getShape();
+                long padTokenId = 220;
+
+                // Augment the larger, batch1
+                long[] shapeDelta = batch1.getShape().getShape();
+                shapeDelta[0] = shape2[0];
+                NDArray deltaArray;
+                if (i == 0) {
+                    // The outputTokenIds is padded with padTokenId
+                    deltaArray =
+                            manager.full(new Shape(shapeDelta), padTokenId, batch1.getDataType());
+                } else {
+                    // The rest e.g. attentionMask, kvCache, hiddenStates are padded with 0
+                    deltaArray = manager.zeros(new Shape(shapeDelta), batch1.getDataType());
+                }
+                batch1 = batch1.concat(deltaArray, 0);
+
+                // Get the ndIndex used to set the extended part of batch1 to be batch2.
+                NDIndex ndIndex;
+                // Find the ordinal number of the sequence dimension
+                if (seqDimOrder[i] > 0) {
+                    // Has a valid sequence dimension
+                    ndIndex = new NDIndex("{}:", seqBatcher1.batchSize);
+                    int order = 1;
+                    while (order < seqDimOrder[i]) {
+                        ndIndex = ndIndex.addAllDim();
+                        order++;
+                    }
+                    assert seqDelta + shape2[order] == shape1[order]
+                            : "Wrong shapes. batch1 and batch2 are not mergable";
+                    ndIndex = ndIndex.addSliceDim(seqDelta, shape1[order]).addEllipseDim();
+                } else {
+                    // Only batch dimension, no valid sequence dimension
+                    ndIndex = new NDIndex("{}:, ...", seqBatcher1.batchSize);
+                }
+
+                // Copy batch2 to the extended part in batch1
+                batch1.set(ndIndex, batch2);
+                merged.add(batch1);
+            }
+            data = data.fromList(merged, data.getSeqDimOrder());
+
+            batchSize = seqBatcher1.batchSize + seqBatcher2.batchSize;
+            batchUid = seqBatcher1.batchUid.concat(seqBatcher2.batchUid, 0);
+            offSets = seqBatcher1.offSets.concat(seqBatcher2.offSets.addi(seqDelta), 0);
+            seqLength = seqBatcher1.seqLength;
+
+            // memory
+            NDScope.unregister(batchUid, offSets);
+            NDScope.unregister(merged);
+        }
+    }
+
+    /**
+     * Checks which batch needs to exit, according certain criteria like EOS or maxLength.
+     *
+     * <p>It is an iteration over batch and is thus also considered as batch operation.
+     *
+     * @param outputIds output token ids in an incremental forward call
+     * @param maxLength max total sequence length
+     * @param eosTokenId end of sentence token id
+     */
+    public void exitCriteria(NDArray outputIds, long maxLength, long eosTokenId) {
+        long[] outputIdsArray = outputIds.toLongArray();
+        long[] offSetsArray = offSets.toLongArray();
+        for (int i = 0; i < outputIdsArray.length; i++) {
+            if (seqLength - offSetsArray[i] >= maxLength || outputIdsArray[i] == eosTokenId) {
+                if (!exitIndexEndPosition.containsKey((long) i)) {
+                    exitIndexEndPosition.put((long) i, seqLength);
+                }
+            }
+        }
+    }
+
+    /**
+     * Collects the finished sequences and trim the left padding.
+     *
+     * @return a map that stores request id to output token ids
+     */
+    public Map<Long, NDArray> collectAndTrim() {
+        if (exitIndexEndPosition.isEmpty()) {
+            return new ConcurrentHashMap<>();
+        }
+        Map<Long, NDArray> finishedSequences = new ConcurrentHashMap<>();
+
+        try (NDScope scope = new NDScope()) {
+            scope.suppressNotUsedWarning();
+            // Collect the results into finishedSequences
+            Set<Long> exitIndices = new HashSet<>();
+            for (Map.Entry<Long, Long> entry : exitIndexEndPosition.entrySet()) {
+                // batchIndex -> seqEndPosition
+                long batchIndex = entry.getKey();
+                long seqEndPosition = entry.getValue();
+                long uid = batchUid.getLong(batchIndex);
+                long offSet = offSets.getLong(batchIndex);
+                NDArray output =
+                        data.getPastOutputIds()
+                                .get("{}, {}:{}", batchIndex, offSet, seqEndPosition);
+                finishedSequences.put(uid, output);
+                exitIndices.add(batchIndex);
+
+                NDScope.unregister(output);
+            }
+
+            // Find the batch indices of the non-finished sequences.
+            long[] keepIndices = new long[Math.toIntExact(batchSize) - exitIndices.size()];
+            int j = 0;
+            for (long i = 0; i < batchSize; i++) {
+                if (!exitIndices.contains(i)) {
+                    keepIndices[j++] = i;
+                }
+            }
+
+            if (keepIndices.length == 0) {
+                batchUid = manager.create(new Shape(0, 1), batchUid.getDataType());
+                offSets = manager.create(new Shape(0, 1), offSets.getDataType());
+                data = null;
+                batchSize = 0;
+                seqLength = 0;
+                exitIndexEndPosition = new ConcurrentHashMap<>();
+
+                NDScope.unregister(batchUid, offSets);
+                return finishedSequences;
+            }
+
+            NDIndex ndIndex = new NDIndex("{}", manager.create(keepIndices));
+            batchUid = batchUid.get(ndIndex).reshape(-1, 1);
+            offSets = offSets.get(ndIndex).reshape(-1, 1);
+            long trimSeq = offSets.min(new int[] {0}).toLongArray()[0];
+            offSets = offSets.subi(trimSeq);
+
+            // Trim batch, and sequence dim if needed
+            NDList list = data.getList();
+            NDList newList = new NDList(list.size());
+            long[] seqDimOrder = data.getSeqDimOrder();
+            for (int i = 0; i < list.size(); i++) {
+                NDArray batch = list.get(i);
+                if (trimSeq == 0) {
+                    // no need to trim
+                    ndIndex = new NDIndex("{}, ...", manager.create(keepIndices));
+                    newList.add(batch.get(ndIndex));
+                    continue;
+                }
+
+                // Get the ndIndex used to keep the entries and trim the rest
+                // Find the ordinal number of the sequence dimension
+                if (seqDimOrder[i] > 0) {
+                    // Has a valid sequence dimension
+                    ndIndex = new NDIndex("{}", manager.create(keepIndices));
+                    int order = 1;
+                    while (order < seqDimOrder[i]) {
+                        ndIndex = ndIndex.addAllDim();
+                        order++;
+                    }
+                    ndIndex = ndIndex.addSliceDim(trimSeq, seqLength).addEllipseDim();
+                } else {
+                    // Only batch dimension, no valid sequence dimension
+                    ndIndex = new NDIndex("{}, ...", manager.create(keepIndices));
+                }
+                // Keep the indexed entries and trim the rest
+                newList.add(batch.get(ndIndex));
+            }
+            data = data.fromList(newList, data.getSeqDimOrder());
+
+            batchSize -= exitIndexEndPosition.size();
+            seqLength -= trimSeq;
+
+            exitIndexEndPosition = new ConcurrentHashMap<>();
+
+            // memory
+            NDScope.unregister(newList);
+            NDScope.unregister(batchUid, offSets);
+
+            return finishedSequences;
+        }
+    }
+
+    /**
+     * Computes the position ids by linear search from the left.
+     *
+     * @return the boolean indicating whether all sequences are empty
+     */
+    public boolean sequenceComplete() {
+        return !exitIndexEndPosition.isEmpty();
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/StepGeneration.java b/api/src/main/java/ai/djl/modality/nlp/generate/StepGeneration.java
new file mode 100644
index 000000000..9f56d3c8d
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/StepGeneration.java
@@ -0,0 +1,180 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+
+/**
+ * {@code StepGeneration} is a utility class containing the step generation utility functions used
+ * in autoregressive search.
+ */
+public final class StepGeneration {
+
+    private StepGeneration() {}
+
+    /**
+     * Generate the output token id and selecting indices used in contrastive search.
+     *
+     * @param topKIds the topk candidate token ids
+     * @param logits the logits from the language model
+     * @param contextHiddenStates the embedding of the past generated token ids
+     * @param topkHiddenStates the embedding of the topk candidate token ids
+     * @param offSets the offsets
+     * @param alpha the repetition penalty
+     * @return the output token ids and selecting indices
+     */
+    public static NDList constrastiveStepGeneration(
+            NDArray topKIds,
+            NDArray logits,
+            NDArray contextHiddenStates,
+            NDArray topkHiddenStates,
+            NDArray offSets,
+            float alpha) {
+        /*
+          topKIds: [batch, topK]
+          attentionMask: [batch, past_seq]
+          logits:  [batch, vocabSize]
+          contextHiddenStates: [batch, past_seq, dim]
+          topkHiddenStates: [batch*topK, seq=1, dim]
+          attentionMaskSlice: [batch, 2]: (startPosition, endPosition)
+        */
+
+        long batch = topKIds.getShape().get(0);
+        long topK = topKIds.getShape().get(1);
+        long hiddenDim = topkHiddenStates.getShape().getLastDimension();
+
+        // [batch*topK, seq=1, dim] -> [batch, topK, dim]
+        topkHiddenStates = topkHiddenStates.reshape(batch, topK, hiddenDim);
+
+        //  [batch, topK, dim] * [batch, past_seq, dim] -> [batch, topK, past_seq]
+        topkHiddenStates = topkHiddenStates.normalize(2, 2);
+        contextHiddenStates = contextHiddenStates.normalize(2, 2);
+        NDArray cosSimilarity =
+                topkHiddenStates.batchMatMul(contextHiddenStates.transpose(0, 2, 1));
+
+        // Deactivate entries (batch_idx, :, zero_attention_idx_slice) in max{cosSim} step
+        long[] offSetsArray = offSets.toLongArray();
+        for (int i = 0; i < offSetsArray.length; i++) {
+            cosSimilarity.set(new NDIndex("{}, :, {}:{}", i, 0, offSetsArray[i]), -1);
+        }
+
+        // [batch, topK, past_seq] -> [batch, topK]
+        NDArray topkScorePart1 = cosSimilarity.max(new int[] {2});
+        assert topkScorePart1.getShape().getShape().length == 2 : "Wrong output size";
+        // [batch, logitDim].gather([batch, topK) -> [batch, topK]
+        NDArray topkScorePart2 = logits.softmax(1).gather(topKIds, 1);
+        NDArray topkScore = topkScorePart2.muli(1 - alpha).subi(topkScorePart1.muli(alpha));
+
+        // [batch, topK] => [batch, 1]
+        NDArray select = topkScore.argMax(1);
+        NDIndex selectIndex =
+                new NDIndex(
+                        "{}, {}, ...",
+                        logits.getManager().arange(0, topKIds.getShape().get(0), 1, DataType.INT64),
+                        select);
+        NDArray outputIds = topKIds.get(selectIndex).reshape(-1, 1);
+        return new NDList(outputIds, select);
+    }
+
+    // TODO: add support of Einstein summation:
+    // a = torch.randn(batch, past_seq, dim)
+    // b = torch.randn(batch, topK, dim)
+    // result = torch.einsum('bik,bjk->bij', a, b)
+
+    /**
+     * Generates the output token id for greedy search.
+     *
+     * @param logits the logits from the language model
+     * @return the output token ids
+     */
+    public static NDArray greedyStepGen(NDArray logits) {
+        // logits:  [batch, seq, probDim]
+        assert logits.getShape().getShape().length == 3 : "unexpected input";
+        logits = logits.get(":, -1, :");
+        return logits.argMax(-1).expandDims(1); // [batch, vacDim]
+    }
+
+    /**
+     * Generates the output token id and selecting indices used in beam search.
+     *
+     * @param lastProbs the probabilities of the past prefix sequences
+     * @param logits the logits
+     * @param numBatch number of batch
+     * @param numBeam number of beam
+     * @return the output token ids and selecting indices
+     */
+    public static NDList beamStepGeneration(
+            NDArray lastProbs, NDArray logits, long numBatch, long numBeam) {
+        // [batch * beamSource, seq, probDim] -> [batch, beamSource, probDim]
+        NDArray allProbs = logits.get(":, -1, :").softmax(1).reshape(numBatch, numBeam, -1);
+
+        // Argmax over the probs in the prob dimension.
+        // [batch, beamSource, probDim] -> [batch, beamSource, beamChild]
+        NDList topK = allProbs.topK(Math.toIntExact(numBeam), -1, true, false);
+        NDArray outputIs = topK.get(1);
+        NDArray stepProbs = topK.get(0);
+
+        // Chain the probability
+        // [batch, beamSource] -> [batch, beamSource, 1]
+        lastProbs = lastProbs.reshape(numBatch, numBeam, 1);
+        // [batch, beamSource, beamChild]
+        NDArray newProbs = stepProbs.muli(lastProbs);
+
+        // Argmax over the (beamSource * beamChild) dimension
+        topK =
+                newProbs.reshape(numBatch, numBeam * numBeam)
+                        .topK(Math.toIntExact(numBeam), -1, true, false);
+
+        // The select indices act on (beamSource, beamChild) dimension. Decides how the new
+        // generated tokenIds correspond to the past tokenIds.
+        // [batch, beamNew].
+        NDArray select = topK.get(1);
+        // Act on [batch, beam, ...] dimension and the output will be [batch, beam, ...]
+        NDIndex selectIndex =
+                new NDIndex(
+                        "{}, {}, ...",
+                        logits.getManager()
+                                .arange(0, numBatch, 1, DataType.INT64)
+                                .expandDims(1)
+                                .repeat(1, numBeam),
+                        select);
+
+        // [batch, beamNew]
+        outputIs = outputIs.reshape(numBatch, numBeam * numBeam).get(selectIndex).expandDims(2);
+        // [batch, beamNew]
+        newProbs = newProbs.reshape(numBatch, numBeam * numBeam).get(selectIndex).normalize(1, 1);
+
+        /* During the beam selection process, some source beams are selected several times while
+        some source beams are not selected even once. The pastOutputs should be reselected to
+        have the right correspondence to the newInputIds.
+        */
+        // [batch, beamNew]
+        assert select.getDataType() == DataType.INT64 : "Wrong output! Expect integer division";
+        assert select.getShape().getShape().length == 2 : "Wrong size. Expect [batch, beamNew]";
+        // For each batch, convert the index1 in beamSource*beamChild dimension to its index2 in
+        // beamSource dimension: index2 = index1 / numBeam.
+        long[] index = select.toLongArray();
+        for (int i = 0; i < index.length; i++) {
+            index[i] = Math.floorDiv(index[i], numBeam);
+        }
+        NDArray sourceBeamSelected =
+                logits.getManager().create(index, new Shape(numBatch, numBeam));
+
+        return new NDList(outputIs, newProbs, sourceBeamSelected);
+    }
+    // TODO: implement pytorch floor_divide.
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java b/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java
new file mode 100644
index 000000000..3d79168ae
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/TextGenerator.java
@@ -0,0 +1,558 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.modality.nlp.generate;
+
+import ai.djl.inference.Predictor;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.NDScope;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.translate.TranslateException;
+
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+/**
+ * {@code TextGenerator} is an LMSearch (language model search) which contains multiple
+ * autoregressive search methods.
+ *
+ * <p>It has a Predictor from NDList to CausalLMOutput, which is called inside an autoregressive
+ * inference loop.
+ */
+public class TextGenerator {
+
+    private String searchName;
+    private SearchConfig config;
+    private Predictor<NDList, CausalLMOutput> predictor;
+
+    private NDArray positionOffset;
+
+    /**
+     * Constructs a new {@code TextGenerator} instance.
+     *
+     * @param predictor the language model
+     * @param searchName the autoregressive search name
+     * @param searchConfig the autoregressive search configuration
+     */
+    public TextGenerator(
+            Predictor<NDList, CausalLMOutput> predictor,
+            String searchName,
+            SearchConfig searchConfig) {
+        this.predictor = predictor;
+        this.searchName = searchName;
+        this.config = searchConfig;
+    }
+
+    /**
+     * Executes greedy search.
+     *
+     * @param inputIds the input token ids.
+     * @return the output token ids stored as NDArray
+     * @throws TranslateException if forward fails
+     */
+    @SuppressWarnings("try")
+    public NDArray greedySearch(NDArray inputIds) throws TranslateException {
+        NDArray attentionMask = prepareAttentionMaskOffset(inputIds, config);
+        NDManager manager = inputIds.getManager();
+        GreedyBatchTensorList searchState =
+                new GreedyBatchTensorList(inputIds, null, null, attentionMask);
+        while (true) {
+            try (NDScope ignore = new NDScope()) {
+                NDArray pastOutputIds = searchState.getPastOutputIds();
+                NDArray nextInputIds = searchState.getNextInputIds();
+                NDArray pastAttentionMask = searchState.getPastAttentionMask();
+                NDList pastKeyValues = searchState.getPastKeyValues();
+                long pastSeqLength =
+                        pastOutputIds == null ? 0 : pastOutputIds.getShape().getLastDimension();
+                NDList modelInput = prepareInput(nextInputIds, pastAttentionMask, pastSeqLength, 1);
+                if (pastKeyValues != null) {
+                    modelInput.addAll(pastKeyValues);
+                }
+                CausalLMOutput modelOutput = predictor.predict(modelInput);
+
+                NDArray outputIds = StepGeneration.greedyStepGen(modelOutput.getLogits());
+
+                // Update searchState
+                if (pastOutputIds == null) {
+                    pastOutputIds = nextInputIds;
+                    searchState.setPastOutputIds(pastOutputIds);
+                } else {
+                    pastOutputIds = pastOutputIds.concat(nextInputIds, 1);
+                    searchState.setPastOutputIds(pastOutputIds);
+                }
+                nextInputIds = outputIds;
+                searchState.setNextInputIds(nextInputIds);
+                pastKeyValues = modelOutput.getPastKeyValuesList();
+                searchState.setPastKeyValues(pastKeyValues);
+                pastAttentionMask =
+                        pastAttentionMask.concat(
+                                manager.ones(
+                                        new Shape(inputIds.getShape().get(0), 1), DataType.INT64),
+                                1);
+                searchState.setPastAttentionMask(pastAttentionMask);
+
+                // memory management
+                NDScope.unregister(nextInputIds, pastAttentionMask, pastOutputIds);
+                NDScope.unregister(pastKeyValues);
+            }
+
+            // Termination Criteria
+            // TODO: <EOS>, delete the sentence and add it to result.
+            if (searchState.getPastOutputIds().getShape().get(1) + 1 >= config.getMaxSeqLength()) {
+                break;
+            }
+        }
+        return searchState.getPastOutputIds().concat(searchState.getNextInputIds(), 1);
+    }
+
+    /**
+     * Generates text using beam search.
+     *
+     * @param inputIds input tokens ids
+     * @see <a href="https://huggingface.co/blog/how-to-generate">Beam Search</a>
+     * @return output tensor
+     * @throws TranslateException if failed run forward
+     */
+    @SuppressWarnings("try")
+    public NDArray beamSearch(NDArray inputIds) throws TranslateException {
+        NDArray attentionMask = prepareAttentionMaskOffset(inputIds, config);
+        NDManager manager = inputIds.getManager();
+        long numBeam = config.getBeam();
+        long numBatch = inputIds.getShape().get(0);
+        BeamBatchTensorList searchState = new BeamBatchTensorList();
+
+        long numHeads = 0;
+        long kvDim = 0;
+        while (true) {
+            if (searchState.getPastAttentionMask() == null) {
+                // Initial beams
+                NDList modelInput = prepareInput(inputIds, attentionMask, 0, 1);
+                CausalLMOutput modelOutput = predictor.predict(modelInput);
+
+                // [batch, probDim]
+                NDArray allProbs = modelOutput.getLogits().get(":, -1, :").softmax(1);
+
+                // [batch, beam]
+                NDList topK = allProbs.topK(Math.toIntExact(numBeam), -1, true, false);
+                NDArray outputIds = topK.get(1).expandDims(2);
+                NDArray lastProbs = topK.get(0).normalize(1, 1);
+                assert outputIds.getShape().getShape().length == 3 : "Wrong shape";
+                assert lastProbs.getShape().getShape().length == 2 : "Wrong Shape";
+
+                // [batch, beam, seq + 1]
+                attentionMask =
+                        attentionMask
+                                .concat(manager.ones(new Shape(numBatch, 1), DataType.INT64), -1)
+                                .expandDims(1)
+                                .repeat(1, numBeam);
+
+                // [batch, beam, heads, seq_past, kvFeature]
+                Function<NDArray, NDArray> fn = ndarray -> ndarray.expandDims(1).repeat(1, numBeam);
+                NDList pastKeyValues =
+                        new NDList(
+                                modelOutput.getPastKeyValuesList().stream()
+                                        .map(fn)
+                                        .collect(Collectors.toList()));
+                // [batch, beam, seq_past]
+                NDArray pastOutputIds = inputIds.expandDims(1).repeat(1, numBeam);
+
+                searchState =
+                        new BeamBatchTensorList(
+                                outputIds, pastOutputIds, pastKeyValues, attentionMask, lastProbs);
+
+                numHeads = pastKeyValues.get(0).getShape().get(2);
+                kvDim = pastKeyValues.get(0).getShape().getLastDimension();
+            }
+
+            try (NDScope ignore = new NDScope()) {
+                long pastSeqLength = searchState.getPastOutputIds().getShape().getLastDimension();
+                NDList modelInput =
+                        prepareInput(
+                                searchState.getNextInputIds().reshape(numBatch * numBeam, 1),
+                                searchState.getPastAttentionMask().reshape(numBatch * numBeam, -1),
+                                pastSeqLength,
+                                config.getBeam());
+
+                final long finalNumHeads = numHeads;
+                final long finalKvDim = kvDim;
+                Function<NDArray, NDArray> fn =
+                        ndarray ->
+                                ndarray.reshape(
+                                        numBatch * numBeam,
+                                        finalNumHeads,
+                                        pastSeqLength,
+                                        finalKvDim);
+                NDList pastKeyValues =
+                        new NDList(
+                                searchState.getPastKeyValues().stream()
+                                        .map(fn)
+                                        .collect(Collectors.toList()));
+                modelInput.addAll(pastKeyValues);
+                CausalLMOutput modelOutput = predictor.predict(modelInput);
+
+                NDList generatedOutput =
+                        StepGeneration.beamStepGeneration(
+                                searchState.getLastProbs(),
+                                modelOutput.getLogits(),
+                                numBatch,
+                                numBeam);
+
+                // Update searchState
+                searchState = updateSearchState(searchState, modelOutput, generatedOutput, manager);
+
+                // Memory management
+                NDScope.unregister(
+                        searchState.getNextInputIds(),
+                        searchState.getPastOutputIds(),
+                        searchState.getPastAttentionMask(),
+                        searchState.getLastProbs());
+                NDScope.unregister(searchState.getPastKeyValues());
+            }
+
+            // Termination Criteria
+            // TODO: <EOS>, delete the sentence and add it to result.
+            if (searchState.getPastOutputIds().getShape().getLastDimension() + 1
+                    >= config.getMaxSeqLength()) {
+                break;
+            }
+        }
+
+        return searchState
+                .getPastOutputIds()
+                .concat(searchState.getNextInputIds(), -1)
+                .reshape(numBatch * numBeam, -1);
+    }
+
+    /**
+     * Generates text using contrastive search.
+     *
+     * @param inputIds input token ids
+     * @see <a href="https://huggingface.co/blog/introducing-csearch">Contrastive Search</a>
+     * @return the generated {@code NDArray}
+     * @throws TranslateException if forward failed
+     */
+    @SuppressWarnings("try")
+    public NDArray contrastiveSearch(NDArray inputIds) throws TranslateException {
+        // inputIds: [batchSize, seqLength: t_init]
+        // attentionMask: [batchSize, pastSeq]. seq-dim-size = |past_seq| + |inputIds|.
+
+        NDManager manager = inputIds.getManager();
+        NDArray attentionMask = prepareAttentionMaskOffset(inputIds, config);
+        ContrastiveBatchTensorList searchState = new ContrastiveBatchTensorList();
+        while (true) {
+            if (searchState.getPastKeyValues() == null) {
+                NDList modelInput = prepareInput(inputIds, attentionMask, 0, 1);
+                CausalLMOutput output = predictor.predict(modelInput);
+                NDArray lastLogits = output.getLogits().get(":, -1, :");
+                searchState =
+                        new ContrastiveBatchTensorList(
+                                inputIds,
+                                attentionMask,
+                                output.getHiddenState(),
+                                lastLogits,
+                                output.getPastKeyValuesList(),
+                                new long[] {});
+            }
+
+            /* Contrastive search loop main part */
+            // (1) candidate tokens recall;
+            // (2) candidate re-rank by degeneration penalty
+
+            try (NDScope ignore = new NDScope()) {
+                NDArray topKIds =
+                        searchState
+                                .getLogits()
+                                .topK(config.getK(), -1, true, false)
+                                .get(1); // [batch, topK]
+
+                // Generate model inputs and put candidates together into batch
+                // [batch, topK] -> [batch * [topK]] -> [[batch * [topK]], seqLength=1]
+                NDArray candidateInputIds = topKIds.flatten().reshape(-1, 1);
+                assert candidateInputIds.getDataType() == DataType.INT64
+                        : "inputIds datatype should be int64";
+                assert candidateInputIds.getShape().getShape().length == 2 : "shape not right";
+
+                // [batch, heads, seq_past, feature] -> [batch * topK, head, seq_past, feature]
+                NDList kCopyPastKeyValues =
+                        new NDList(
+                                searchState.getPastKeyValues().stream()
+                                        .map(ndarray -> ndarray.repeat(0, config.getK()))
+                                        .collect(Collectors.toList()));
+                assert kCopyPastKeyValues.get(0).getDataType() == DataType.FLOAT32
+                        : "inputIds datatype should be Float32";
+
+                // [batch, seq_past] -> [batch * topK, seq_past] -> [batch * topK, seq_past + 1]
+                long numBatch = topKIds.getShape().get(0);
+                NDArray kCopyPastAttentionMask =
+                        searchState.getPastAttentionMask().repeat(0, config.getK());
+                kCopyPastAttentionMask =
+                        kCopyPastAttentionMask.concat(
+                                manager.ones(
+                                        new Shape(numBatch * config.getK(), 1), DataType.INT64),
+                                1);
+                assert kCopyPastKeyValues.get(0).getShape().get(2) + 1
+                                == kCopyPastAttentionMask.getShape().getLastDimension()
+                        : "attentionMask_seq = past_seq + new_input_seq";
+
+                // Forward with candidates in batch input
+                NDList candidateModelInput =
+                        prepareInput(
+                                candidateInputIds,
+                                kCopyPastAttentionMask,
+                                searchState.getPastOutputIds().getShape().getLastDimension(),
+                                config.getK());
+                candidateModelInput.addAll(kCopyPastKeyValues);
+                CausalLMOutput candidateOutput = predictor.predict(candidateModelInput);
+
+                NDList generatedOutput =
+                        StepGeneration.constrastiveStepGeneration(
+                                topKIds,
+                                searchState.getLogits(),
+                                searchState.getPastHiddenStates(),
+                                candidateOutput.getHiddenState(),
+                                positionOffset,
+                                config.getAlpha());
+
+                // Update searchState for next loop
+                searchState =
+                        updateSearchState(searchState, candidateOutput, generatedOutput, manager);
+
+                // Memory
+                NDScope.unregister(
+                        searchState.getPastOutputIds(),
+                        searchState.getPastAttentionMask(),
+                        searchState.getLogits(),
+                        searchState.getPastHiddenStates());
+                NDScope.unregister(searchState.getPastKeyValues());
+            }
+
+            // TODO: <EOS>, delete the sentence and add it to result.
+            if (searchState.getPastOutputIds().getShape().get(1) >= config.getMaxSeqLength()) {
+                break;
+            }
+        }
+
+        return searchState.getPastOutputIds();
+    }
+
+    private static BeamBatchTensorList updateSearchState(
+            BeamBatchTensorList searchState,
+            CausalLMOutput modelOutput,
+            NDList generatedOutput,
+            NDManager manager) {
+
+        NDList pastKeyValues = searchState.getPastKeyValues();
+        long numHeads = pastKeyValues.get(0).getShape().get(2);
+        long kvDim = pastKeyValues.get(0).getShape().getLastDimension();
+        long numBatch = searchState.getPastOutputIds().getShape().get(0);
+        long numBeam = searchState.getPastOutputIds().getShape().get(1);
+        long pastSeqLength = searchState.getPastOutputIds().getShape().getLastDimension();
+
+        NDArray nextInputIds = generatedOutput.get(0);
+        assert nextInputIds.getShape().getShape().length == 3 : "Wrong Shape";
+        NDArray newProbs = generatedOutput.get(1);
+        // [batch, beamNew]
+        NDArray sourceBeamSelected = generatedOutput.get(2);
+        // Act on [batch, beam, ...] dimension and the output will be [batch, beam, ...]
+        NDIndex sourceBeamIndex =
+                new NDIndex(
+                        "{}, {}, ...",
+                        manager.arange(0, numBatch, 1, DataType.INT64)
+                                .expandDims(1)
+                                .repeat(1, numBeam),
+                        sourceBeamSelected);
+
+        // A simple concatenation is not enough. During the beam selection process, some source
+        // beams are selected several times while some source beams are not selected even once.
+        // The pastOutput should be reselected to have the right correspondence to the
+        // newInputIds.
+        NDArray pastOutputIds =
+                searchState
+                        .getPastOutputIds()
+                        .concat(searchState.getNextInputIds(), -1)
+                        .get(sourceBeamIndex);
+        Function<NDArray, NDArray> fn =
+                ndarray ->
+                        ndarray.reshape(numBatch, numBeam, numHeads, pastSeqLength + 1, kvDim)
+                                .get(sourceBeamIndex);
+        pastKeyValues =
+                new NDList(
+                        modelOutput.getPastKeyValuesList().stream()
+                                .map(fn)
+                                .collect(Collectors.toList()));
+
+        NDArray pastAttentionMask =
+                searchState
+                        .getPastAttentionMask()
+                        .concat(manager.ones(new Shape(numBatch, numBeam, 1), DataType.INT64), -1)
+                        .get(sourceBeamIndex);
+
+        return new BeamBatchTensorList(
+                nextInputIds, pastOutputIds, pastKeyValues, pastAttentionMask, newProbs);
+    }
+
+    private static ContrastiveBatchTensorList updateSearchState(
+            ContrastiveBatchTensorList searchState,
+            CausalLMOutput candidateOutput,
+            NDList generatedOutput,
+            NDManager manager) {
+        // Update searchState for next iteration
+        assert candidateOutput.getLogits().getShape().get(1) == 1
+                : "dimension check: here, outputLogits corresponds to inputSeq == 1";
+        long numBatch = searchState.getLogits().getShape().get(0);
+        long logitsDim = searchState.getLogits().getShape().get(1);
+        long pastSeqLengthPriorUpdate = searchState.getPastOutputIds().getShape().get(1);
+        long numHeads = searchState.getPastKeyValues().get(0).getShape().get(1);
+        long kvDim = searchState.getPastKeyValues().get(0).getShape().get(3);
+        long hiddenDim = searchState.getPastHiddenStates().getShape().get(2);
+        long k = candidateOutput.getLogits().getShape().get(0) / numBatch;
+
+        // [batch, 1]
+        NDArray select = generatedOutput.get(1);
+        NDIndex selectIndex =
+                new NDIndex(
+                        "{}, {}, ...",
+                        manager.arange(0, numBatch, 1, DataType.INT64),
+                        select.flatten());
+
+        // Take from candidateOutput
+        // [batch, k, inputSeq=1, logitsDim] --select--> [batch, logitDim]
+        NDArray nextLogits =
+                candidateOutput.getLogits().reshape(numBatch, k, logitsDim).get(selectIndex);
+
+        // Take from candidateOutput
+        // [batch * k, heads, seq_past, feature] --select--> [batch, heads, seq_past, feature]
+        Function<NDArray, NDArray> fn =
+                ndarray ->
+                        ndarray.reshape(numBatch, k, numHeads, pastSeqLengthPriorUpdate + 1, kvDim)
+                                .get(selectIndex);
+        NDList nextPastKeyValue =
+                new NDList(
+                        candidateOutput.getPastKeyValuesList().stream()
+                                .map(fn)
+                                .collect(Collectors.toList()));
+
+        // To be concatenated into searchState.pastHiddenStates
+        // [batch * k, inputSeq=1, hiddenDim]
+        NDArray newHiddenState = candidateOutput.getHiddenState();
+        assert newHiddenState.getManager() == manager : "possible leaky memory";
+        NDArray nextPastHiddenStates =
+                searchState
+                        .getPastHiddenStates()
+                        .concat(
+                                newHiddenState.reshape(numBatch, k, 1, hiddenDim).get(selectIndex),
+                                1);
+
+        // To be concatenated into searchState.outputIds
+        // [batch, seq_past]
+        NDArray outputIds = generatedOutput.get(0);
+        NDArray nextOutputIds = searchState.getPastOutputIds().concat(outputIds, 1);
+
+        // [batch, seq_past]
+        NDArray nextPastAttentionMask =
+                searchState
+                        .getPastAttentionMask()
+                        .concat(manager.ones(new Shape(numBatch, 1), DataType.INT64), 1);
+
+        return new ContrastiveBatchTensorList(
+                nextOutputIds,
+                nextPastAttentionMask,
+                nextPastHiddenStates,
+                nextLogits,
+                nextPastKeyValue,
+                new long[] {});
+    }
+
+    private NDArray prepareAttentionMaskOffset(NDArray inputIds, SearchConfig config) {
+        // prepare attentionMask and positionOffset
+        // Used to initialize the search
+        boolean suffixPadding = config.isSuffixPadding();
+        NDManager manager = inputIds.getManager();
+        int numBatch = Math.toIntExact(inputIds.getShape().get(0));
+        int initSeqSize = Math.toIntExact(inputIds.getShape().get(1));
+        NDArray attentionMask =
+                manager.ones(new Shape(1, inputIds.getShape().getLastDimension()), DataType.INT64)
+                        .reshape(1, -1)
+                        .repeat(0, numBatch);
+
+        // Linear search from left to find the first position that's not padTokenId.
+        long[][] offset = new long[numBatch][1];
+        for (int i = 0; i < numBatch; i++) {
+            long[] aSequence = inputIds.get("{},:", i).toLongArray();
+            int idx = 0;
+            while (idx < initSeqSize) {
+                if (suffixPadding && aSequence[idx] == config.getPadTokenId()
+                        || !suffixPadding && aSequence[idx] != config.getPadTokenId()) {
+                    break;
+                }
+                idx++;
+            }
+            attentionMask.set(
+                    new NDIndex(
+                            "{},{}:{}",
+                            i,
+                            suffixPadding ? idx : 0,
+                            suffixPadding ? initSeqSize : idx),
+                    0);
+            if (!suffixPadding) {
+                offset[i][0] = idx;
+            }
+        }
+        positionOffset = manager.create(offset);
+        return attentionMask;
+    }
+
+    private NDList prepareInput(
+            NDArray inputIds, NDArray attentionMask, long pastSeqLength, int repeat) {
+        // Pack the model input
+        NDArray positionIds =
+                inputIds.getManager()
+                        .arange(
+                                pastSeqLength,
+                                pastSeqLength + inputIds.getShape().getLastDimension(),
+                                1,
+                                DataType.INT64)
+                        .expandDims(0)
+                        .repeat(0, inputIds.getShape().get(0));
+
+        NDArray positionIdsShifted = positionIds.subi(positionOffset.repeat(0, repeat));
+        positionIds = positionIdsShifted.maximum(positionIdsShifted.zerosLike());
+
+        return new NDList(inputIds, positionIds, attentionMask);
+    }
+
+    /**
+     * Forward function call to generate text.
+     *
+     * @param inputIds the input token ids
+     * @return generated token ids
+     * @throws TranslateException if prediction fails
+     */
+    public NDArray forward(NDArray inputIds) throws TranslateException {
+        switch (searchName) {
+            case "greedy":
+                return greedySearch(inputIds);
+            case "beam":
+                return beamSearch(inputIds);
+            case "contrastive":
+                return contrastiveSearch(inputIds);
+            default:
+                throw new IllegalArgumentException(
+                        "searchName not correctly specified. Please choose among: {greedy, beam,"
+                                + " contrastive}");
+        }
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/nlp/generate/package-info.java b/api/src/main/java/ai/djl/modality/nlp/generate/package-info.java
new file mode 100644
index 000000000..aa997f6a2
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/nlp/generate/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+
+/** Contains utility classes for image manipulation. */
+package ai.djl.modality.nlp.generate;
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslator.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslator.java
new file mode 100644
index 000000000..9c82fa6ff
--- /dev/null
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslator.java
@@ -0,0 +1,106 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.pytorch.zoo.nlp.textgeneration;
+
+import ai.djl.modality.nlp.generate.CausalLMOutput;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.translate.NoBatchifyTranslator;
+import ai.djl.translate.TranslatorContext;
+
+import java.util.stream.Collectors;
+
+/** The {@link ai.djl.translate.Translator} for PyTorch GPT2 model. */
+public class PtGptTranslator implements NoBatchifyTranslator<NDList, CausalLMOutput> {
+
+    private long kvDim;
+    private int numAttentionHeads;
+    private int numLayers;
+    private String tupleName;
+
+    /**
+     * Constructs a new instance of {@code PtGptTranslator}.
+     *
+     * @param kvDim the kv dimension
+     * @param numAttentionHeads the number of attention heads
+     * @param numLayers the number of layers
+     */
+    public PtGptTranslator(long kvDim, int numAttentionHeads, int numLayers) {
+        this.kvDim = kvDim;
+        this.numAttentionHeads = numAttentionHeads;
+        this.numLayers = numLayers;
+        tupleName = "past_key_values(" + numLayers + ',' + 2 + ')';
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDList processInput(TranslatorContext ctx, NDList input) throws Exception {
+        NDManager manager = ctx.getNDManager();
+        if (input.size() == 3) {
+            ctx.setAttachment("initialCall", Boolean.TRUE);
+            long batchSize = input.get(0).getShape().get(0);
+            NDArray attentionMask = input.get(2);
+            attentionMask =
+                    manager.zeros(new Shape(batchSize, 1), DataType.INT64)
+                            .concat(attentionMask, -1);
+            input.set(2, attentionMask);
+            addInitialDummyPastKeyValues(input, input.get(0), manager);
+        }
+
+        return input;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public CausalLMOutput processOutput(TranslatorContext ctx, NDList output) throws Exception {
+        NDArray logitsOutput = output.get(0);
+        NDManager manager = output.getManager();
+        NDList pastKeyValuesOutput = output.subNDList(1, numLayers * 2 + 1);
+        NDArray hiddenStatesOutput;
+        if (output.size() > numLayers * 2 + 1) {
+            hiddenStatesOutput = output.get(numLayers * 2 + 1);
+        } else {
+            // If the traced_GPT2 model outputs hiddenStates, then this is not executed. If the
+            // provided traced model doesn't output hiddenStates, we can throw a warning here.
+            hiddenStatesOutput = manager.zeros(new Shape(1));
+        }
+
+        if (ctx.getAttachment("initialCall") != null) {
+            NDIndex index2 = new NDIndex(":, :, 1:, ...");
+            pastKeyValuesOutput =
+                    new NDList(
+                            pastKeyValuesOutput.stream()
+                                    .map(object -> object.get(index2))
+                                    .collect(Collectors.toList()));
+        }
+
+        for (NDArray array : pastKeyValuesOutput) {
+            array.setName(tupleName);
+        }
+
+        return new CausalLMOutput(logitsOutput, hiddenStatesOutput, pastKeyValuesOutput);
+    }
+
+    private void addInitialDummyPastKeyValues(NDList list, NDArray inputIds, NDManager manager) {
+        long numBatch = inputIds.getShape().get(0);
+        for (int i = 0; i < numLayers * 2; ++i) {
+            NDArray array = manager.zeros(new Shape(numBatch, numAttentionHeads, 1, kvDim));
+            array.setName(tupleName);
+            list.add(array);
+        }
+    }
+}
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslatorFactory.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslatorFactory.java
new file mode 100644
index 000000000..ca3f693ad
--- /dev/null
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/PtGptTranslatorFactory.java
@@ -0,0 +1,57 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.pytorch.zoo.nlp.textgeneration;
+
+import ai.djl.Model;
+import ai.djl.modality.nlp.generate.CausalLMOutput;
+import ai.djl.ndarray.NDList;
+import ai.djl.translate.ArgumentsUtil;
+import ai.djl.translate.Translator;
+import ai.djl.translate.TranslatorFactory;
+import ai.djl.util.Pair;
+
+import java.lang.reflect.Type;
+import java.util.HashSet;
+import java.util.Map;
+import java.util.Set;
+
+/** An {@link TranslatorFactory} that creates a {@link PtGptTranslator} instance. */
+public class PtGptTranslatorFactory implements TranslatorFactory {
+
+    private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
+
+    static {
+        SUPPORTED_TYPES.add(new Pair<>(NDList.class, CausalLMOutput.class));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Set<Pair<Type, Type>> getSupportedTypes() {
+        return SUPPORTED_TYPES;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    @SuppressWarnings("unchecked")
+    public <I, O> Translator<I, O> newInstance(
+            Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
+        if (!isSupported(input, output)) {
+            throw new IllegalArgumentException("Unsupported input/output types.");
+        }
+        long kvDim = ArgumentsUtil.longValue(arguments, "kvDim", 64);
+        int numAttentionHeads = ArgumentsUtil.intValue(arguments, "numAttentionHeads", 12);
+        int numLayers = ArgumentsUtil.intValue(arguments, "numLayers", 12);
+
+        return (Translator<I, O>) (new PtGptTranslator(kvDim, numAttentionHeads, numLayers));
+    }
+}
diff --git a/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java
new file mode 100644
index 000000000..fdfee9d4e
--- /dev/null
+++ b/engines/pytorch/pytorch-model-zoo/src/main/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+
+/** Contains classes for the {@link ai.djl.Application.NLP#TEXT_GENERATION} models. */
+package ai.djl.pytorch.zoo.nlp.textgeneration;
diff --git a/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java b/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java
new file mode 100644
index 000000000..4bab3a9cf
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/nlp/TextGeneration.java
@@ -0,0 +1,83 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.examples.inference.nlp;
+
+import ai.djl.MalformedModelException;
+import ai.djl.examples.inference.ImageClassification;
+import ai.djl.huggingface.tokenizers.Encoding;
+import ai.djl.huggingface.tokenizers.HuggingFaceTokenizer;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.nlp.generate.CausalLMOutput;
+import ai.djl.modality.nlp.generate.SearchConfig;
+import ai.djl.modality.nlp.generate.TextGenerator;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ModelNotFoundException;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.translate.DeferredTranslatorFactory;
+import ai.djl.translate.TranslateException;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.io.IOException;
+
+public final class TextGeneration {
+
+    private static final Logger logger = LoggerFactory.getLogger(ImageClassification.class);
+
+    private TextGeneration() {}
+
+    public static void main(String[] args)
+            throws ModelNotFoundException, MalformedModelException, IOException,
+                    TranslateException {
+        String ret = generateTextWithPyTorch();
+        logger.info("{}", ret);
+    }
+
+    public static String generateTextWithPyTorch()
+            throws ModelNotFoundException, MalformedModelException, IOException,
+                    TranslateException {
+        SearchConfig config = new SearchConfig();
+        config.setMaxSeqLength(60);
+
+        String url = "https://djl-misc.s3.amazonaws.com/test/models/gpt2/gpt2_pt.zip";
+
+        Criteria<NDList, CausalLMOutput> criteria =
+                Criteria.builder()
+                        .setTypes(NDList.class, CausalLMOutput.class)
+                        .optModelUrls(url)
+                        .optEngine("PyTorch")
+                        .optTranslatorFactory(new DeferredTranslatorFactory())
+                        .build();
+        String input = "DeepMind Company is";
+
+        try (ZooModel<NDList, CausalLMOutput> model = criteria.loadModel();
+                Predictor<NDList, CausalLMOutput> predictor = model.newPredictor();
+                NDManager manager = model.getNDManager().newSubManager();
+                HuggingFaceTokenizer tokenizer = HuggingFaceTokenizer.newInstance("gpt2")) {
+
+            TextGenerator generator = new TextGenerator(predictor, "greedy", config);
+
+            Encoding encoding = tokenizer.encode(input);
+            long[] inputIds = encoding.getIds();
+            NDArray inputIdArray = manager.create(inputIds).expandDims(0);
+
+            NDArray output = generator.greedySearch(inputIdArray);
+            long[] outputIds = output.toLongArray();
+            return tokenizer.decode(outputIds);
+        }
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/inference/nlp/package-info.java b/examples/src/main/java/ai/djl/examples/inference/nlp/package-info.java
new file mode 100644
index 000000000..96e4766ac
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/inference/nlp/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+
+/** Contains examples of NLP. */
+package ai.djl.examples.inference.nlp;


Test patch:
diff --git a/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/TextGenerationTest.java b/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/TextGenerationTest.java
new file mode 100644
index 000000000..800843e95
--- /dev/null
+++ b/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/TextGenerationTest.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.pytorch.zoo.nlp.textgeneration;
+
+import ai.djl.ModelException;
+import ai.djl.inference.Predictor;
+import ai.djl.modality.nlp.generate.CausalLMOutput;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.nn.LambdaBlock;
+import ai.djl.repository.zoo.Criteria;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.translate.TranslateException;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+
+public class TextGenerationTest {
+
+    @Test
+    public void testGpt2() throws TranslateException, ModelException, IOException {
+        Block block =
+                new LambdaBlock(
+                        a -> {
+                            NDList list = new NDList(25);
+                            NDManager manager = a.getManager();
+                            long[][] logits = new long[4][50257];
+                            logits[3][257] = 1;
+                            NDArray arr = manager.create(logits).expandDims(0);
+                            list.add(arr);
+
+                            for (int i = 0; i < 12 * 2; ++i) {
+                                NDArray array = manager.zeros(new Shape(1, 12, 1, 64));
+                                list.add(array);
+                            }
+                            return list;
+                        },
+                        "model");
+
+        Path modelDir = Paths.get("build/text_generation");
+        Files.createDirectories(modelDir);
+
+        Criteria<NDList, CausalLMOutput> criteria =
+                Criteria.builder()
+                        .setTypes(NDList.class, CausalLMOutput.class)
+                        .optModelPath(modelDir)
+                        .optBlock(block)
+                        .optOption("hasParameter", "false")
+                        .optTranslatorFactory(new PtGptTranslatorFactory())
+                        .build();
+
+        try (ZooModel<NDList, CausalLMOutput> model = criteria.loadModel();
+                Predictor<NDList, CausalLMOutput> predictor = model.newPredictor();
+                NDManager manager = NDManager.newBaseManager()) {
+            long[][] inputIds = {{29744, 28478, 5834, 318}};
+            int len = inputIds[0].length;
+            NDArray input = manager.create(inputIds);
+            NDArray attentionMask = manager.ones(new Shape(1, len), DataType.INT64);
+            NDArray positionIds = manager.arange(0, len, 1, DataType.INT64).expandDims(0);
+            CausalLMOutput res = predictor.predict(new NDList(input, attentionMask, positionIds));
+            NDArray logits = res.getLogits();
+            long nextTokenId = logits.get(":, -1, :").argMax().getLong();
+            Assert.assertEquals(nextTokenId, 257);
+            NDList list = res.getPastKeyValuesList();
+            Assert.assertEquals(list.size(), 24);
+            Assert.assertEquals(res.getHiddenState().getShape().get(0), 1);
+        }
+    }
+}
diff --git a/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java b/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java
new file mode 100644
index 000000000..e1c0aaa17
--- /dev/null
+++ b/engines/pytorch/pytorch-model-zoo/src/test/java/ai/djl/pytorch/zoo/nlp/textgeneration/package-info.java
@@ -0,0 +1,16 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+/**
+ * Contains unit test classes for text generation.
+ */
+package ai.djl.pytorch.zoo.nlp.textgeneration;
diff --git a/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java b/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java
new file mode 100644
index 000000000..a79573fcd
--- /dev/null
+++ b/examples/src/test/java/ai/djl/examples/inference/nlp/TextGenerationTest.java
@@ -0,0 +1,41 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.examples.inference.nlp;
+
+import ai.djl.ModelException;
+import ai.djl.testing.TestRequirements;
+import ai.djl.translate.TranslateException;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+
+public class TextGenerationTest {
+
+    @Test
+    public void testTextGeneration() throws TranslateException, ModelException, IOException {
+        TestRequirements.weekly();
+        TestRequirements.engine("PyTorch");
+
+        String expected =
+                "DeepMind Company is a global leader in the field of artificial"
+                        + " intelligence and artificial intelligence. We are a leading provider"
+                        + " of advanced AI solutions for the automotive industry, including the"
+                        + " latest in advanced AI solutions for the automotive industry. We are"
+                        + " also a leading provider of advanced AI solutions for the automotive"
+                        + " industry, including the";
+
+        Assert.assertEquals(TextGeneration.generateTextWithPyTorch(), expected);
+    }
+}
diff --git a/examples/src/test/java/ai/djl/examples/inference/nlp/package-info.java b/examples/src/test/java/ai/djl/examples/inference/nlp/package-info.java
new file mode 100644
index 000000000..fa1d19a80
--- /dev/null
+++ b/examples/src/test/java/ai/djl/examples/inference/nlp/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+
+/** Contains tests for the inference examples. */
+package ai.djl.examples.inference.nlp;