Repository: deepjavalibrary/djl
Pull request #2678 (2023-06-25 22:06:04)

[api] Improves DJL URL error message
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here


codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2678?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Patch coverage: **`71.51`**% and project coverage change: **`+1.48`** :tada:
> 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 [(`a1e9804`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2678?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 73.57%.

: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    #2678      +/-   ##
============================================
+ Coverage     72.08%   73.57%   +1.48%     
- Complexity     5126     6975    +1849     
============================================
  Files           473      685     +212     
  Lines         21970    30459    +8489     
  Branches       2351     3157     +806     
============================================
+ Hits          15838    22409    +6571     
- Misses         4925     6520    +1595     
- Partials       1207     1530     +323     
```


| [Impacted Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2678?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/2678?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/2678?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/2678?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/2678?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/2678?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/2678?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/2678?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/2678?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/2678?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/2678?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/2678?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

... and [355 files with indirect coverage changes](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2678/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/2678?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/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java b/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
index 7b6c3b7c9..71f394d6d 100644
--- a/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
+++ b/api/src/main/java/ai/djl/repository/RepositoryFactoryImpl.java
@@ -253,14 +253,12 @@ class RepositoryFactoryImpl implements RepositoryFactory {
 
             ModelZoo zoo = ModelZoo.getModelZoo(groupId);
             if (zoo == null) {
-                logger.warn("ModelZoo not found: {}", groupId);
-                return repo;
+                throw new IllegalArgumentException("ModelZoo not found in classpath: " + groupId);
             }
 
             ModelLoader loader = zoo.getModelLoader(artifactId);
             if (loader == null) {
-                logger.warn("Artifact not found: {}/{}", groupId, artifactId);
-                return repo;
+                throw new IllegalArgumentException("Invalid djl URL: " + uri);
             }
 
             MRL mrl =
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/api/src/test/java/ai/djl/repository/DjlRepositoryTest.java b/api/src/test/java/ai/djl/repository/DjlRepositoryTest.java
index 24a0b1c6c..5a53da4f4 100644
--- a/api/src/test/java/ai/djl/repository/DjlRepositoryTest.java
+++ b/api/src/test/java/ai/djl/repository/DjlRepositoryTest.java
@@ -37,11 +37,13 @@ public class DjlRepositoryTest {
         repo = Repository.newInstance("DJL", "djl://ai.djl.pytorch/resnet/0.0.1/traced_resnet18");
         Assert.assertEquals(repo.getResources().size(), 1);
 
-        repo = Repository.newInstance("DJL", "djl://ai.djl.pytorch/fake/0.0.1");
-        Assert.assertEquals(repo.getResources().size(), 0);
+        Assert.assertThrows(
+                IllegalArgumentException.class,
+                () -> Repository.newInstance("DJL", "djl://ai.djl.pytorch/fake/0.0.1"));
 
-        repo = Repository.newInstance("DJL", "djl://ai.djl.fake/mlp/0.0.1");
-        Assert.assertEquals(repo.getResources().size(), 0);
+        Assert.assertThrows(
+                IllegalArgumentException.class,
+                () -> Repository.newInstance("DJL", "djl://ai.djl.fake/mlp/0.0.1"));
 
         Assert.expectThrows(
                 IllegalArgumentException.class, () -> Repository.newInstance("DJL", "djl://"));
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;