Support vocab_mask/prefix_vocab_mask/no_repeat_number in greedysearch op (#12327)

* support more inputs for greedy search

* fix docs

* refactor test

* lint

* review comments

docs/ContribOperators.md

@@ -1569,11 +1569,13 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dd>The id of the end-of-sequence token</dd>
 <dt><tt>model_type</tt> : int</dt>
 <dd>model type: 0 for decoder only like GPT-2; 1 for encoder decoder like Bart</dd>
+<dt><tt>no_repeat_ngram_size</tt> : int</dt>
+<dd>no repeat ngrams size</dd>
 <dt><tt>pad_token_id</tt> : int (required)</dt>
 <dd>The id of the padding token</dd>
 </dl>
 
-#### Inputs (2 - 4)
+#### Inputs (2 - 6)
 
 <dl>
 <dt><tt>input_ids</tt> : I</dt>
@@ -1584,6 +1586,10 @@ This version of the operator has been available since version 1 of the 'com.micr
 <dd>The minimum length below which the score of eos_token_id is set to -Inf. Shape is (1)</dd>
 <dt><tt>repetition_penalty</tt> (optional) : T</dt>
 <dd>The parameter for repetition penalty. Default value 1.0 means no penalty. Accepts value > 0.0. Shape is (1)</dd>
+<dt><tt>vocab_mask</tt> (optional) : I</dt>
+<dd>Mask of vocabulary. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (vacab_size)</dd>
+<dt><tt>prefix_vocab_mask</tt> (optional) : I</dt>
+<dd>Mask of vocabulary for first step. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (batch_size, vocab_size)</dd>
 </dl>
 
 #### Outputs

docs/OperatorKernels.md

@@ -412,7 +412,7 @@ Do not modify directly.*
 |FusedMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
 |GatherND|*in* data:**T**<br> *in* indices:**Tind**<br> *out* output:**T**|1+|**T** = tensor(bfloat16), tensor(bool), tensor(double), tensor(float), tensor(float16), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(string), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)<br/> **Tind** = tensor(int32), tensor(int64)|
 |Gelu|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
-|GreedySearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* repetition_penalty:**T**<br> *out* sequences:**I**|1+|**T** = tensor(float)|
+|GreedySearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**I**<br> *in* prefix_vocab_mask:**I**<br> *out* sequences:**I**|1+|**T** = tensor(float)|
 |GridSample|*in* X:**T1**<br> *in* Grid:**T1**<br> *out* Y:**T2**|1+|**T1** = tensor(float)<br/> **T2** = tensor(float)|
 |Inverse|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
 |MatMulInteger16|*in* A:**T1**<br> *in* B:**T2**<br> *out* Y:**T3**|1+|**T1** = tensor(int16)<br/> **T2** = tensor(int16)<br/> **T3** = tensor(int32)|
@@ -770,7 +770,7 @@ Do not modify directly.*
 |FusedConv|*in* X:**T**<br> *in* W:**T**<br> *in* B:**T**<br> *in* Z:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
 |FusedMatMul|*in* A:**T**<br> *in* B:**T**<br> *out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
 |Gelu|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
-|GreedySearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* repetition_penalty:**T**<br> *out* sequences:**I**|1+|**T** = tensor(float), tensor(float16)|
+|GreedySearch|*in* input_ids:**I**<br> *in* max_length:**I**<br> *in* min_length:**I**<br> *in* repetition_penalty:**T**<br> *in* vocab_mask:**I**<br> *in* prefix_vocab_mask:**I**<br> *out* sequences:**I**|1+|**T** = tensor(float), tensor(float16)|
 |GridSample|*in* X:**T1**<br> *in* Grid:**T1**<br> *out* Y:**T2**|1+|**T1** = tensor(float)<br/> **T2** = tensor(float)|
 |Inverse|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
 |Irfft|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|

onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h

@@ -177,68 +177,11 @@ Status BeamSearchBase<T>::CheckInputs(const OpKernelContextInternal& context) {
   // Input shapes:
   //   input_ids  : (batch_size, sequence_length)
   //   vocab_mask : (vocab_size) or nullptr
-
-  const Tensor* input_ids = context.Input<Tensor>(0);
-  const auto& dims = input_ids->Shape().GetDims();
-  if (dims.size() != 2) {
-    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                           "Input 'input_ids' is expected to have 2 dimensions, got ", dims.size());
-  }
-
-  const Tensor* vocab_mask = context.Input<Tensor>(7);
-  if (vocab_mask != nullptr) {  // vocab_mask is optional
-    const auto& vocab_mask_dims = vocab_mask->Shape().GetDims();
-    if (vocab_mask_dims.size() != 1) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                             "Input 'vocab_mask' is expected to have 1 dimension, got ", vocab_mask_dims.size());
-    }
-
-    // There is dependency on vocab_size parameter, which shall be set before calling this function.
-    if (static_cast<int>(vocab_mask_dims[0]) != parameters_->vocab_size) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                             "Input 'vocab_mask' shape does not match with vocab_size, got ", vocab_mask_dims[0]);
-    }
-
-    // store vocab mask in parameters.
-    parameters_->vocab_mask = vocab_mask->DataAsSpan<int32_t>();
-  }
-
-  const Tensor* prefix_vocab_mask = context.Input<Tensor>(8);
-  if (prefix_vocab_mask != nullptr) {  // prefix_vocab_mask is optional
-    const auto& vocab_mask_dims = prefix_vocab_mask->Shape().GetDims();
-    if (vocab_mask_dims.size() != 2) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                             "Input 'prefix_vocab_mask' is expected to be 2 dimensions, got ", vocab_mask_dims.size());
-    }
-
-    // prefix_vocab_mask first dimension should be same as the first dimension of input_ids
-    if (static_cast<int>(vocab_mask_dims[0]) != static_cast<int>(dims[0])) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                             "input_ids and prefix_vocab_mask must have the same batch_size");
-    }
-
-    // There is dependency on vocab_size parameter, which shall be set before calling this function.
-    if (static_cast<int>(vocab_mask_dims[1]) != parameters_->vocab_size) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                             "Input 'prefix_vocab_mask' shape[1] shall be vocab_size, got ", vocab_mask_dims[1]);
-    }
-
-    // store prefix vocab mask in parameters.
-    parameters_->prefix_vocab_mask = prefix_vocab_mask->DataAsSpan<int32_t>();
-  }
-
-  const Tensor* attention_mask = context.Input<Tensor>(9);
-  if (attention_mask != nullptr) {
-    const auto& dims_attn = attention_mask->Shape().GetDims();
-    if (dims_attn.size() != 2) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                             "Input 'attention_mask' is expected to have 2 dimensions, got ", dims_attn.size());
-    }
-    if (dims_attn != dims) {
-      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
-                             "Input 'attention_mask' is expected to have same shape as input_ids");
-    }
-  }
+  ORT_RETURN_IF_ERROR(this->CheckInputsImpl(parameters_,
+                                            context.Input<Tensor>(0),     // input_ids
+                                            context.Input<Tensor>(7),     // vocab_mask
+                                            context.Input<Tensor>(8),     // prefix_vocab_mask
+                                            context.Input<Tensor>(10)));  // attention_mask
 
   return Status::OK();
 }

onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h

@@ -79,6 +79,77 @@ class GenerateBase {
       return Status::OK();
   }
 
+  template <typename ParametersT>
+  Status CheckInputsImpl(const ParametersT& parameters,
+                         const Tensor* input_ids,
+                         const Tensor* vocab_mask,
+                         const Tensor* prefix_vocab_mask,
+                         const Tensor* attention_mask) const {
+    const auto& dims = input_ids->Shape().GetDims();
+    if (dims.size() != 2) {
+      return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                             "Input 'input_ids' is expected to have 2 dimensions, got ", dims.size());
+    }
+
+    if (vocab_mask != nullptr) {  // vocab_mask is optional
+      const auto& vocab_mask_dims = vocab_mask->Shape().GetDims();
+      if (vocab_mask_dims.size() != 1) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Input 'vocab_mask' is expected to have 1 dimension, got ", vocab_mask_dims.size());
+      }
+
+      // There is dependency on vocab_size parameter, which shall be set before calling this function.
+      if (static_cast<int>(vocab_mask_dims[0]) != parameters->vocab_size) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Input 'vocab_mask'  dimension 0 does not match with vocab_size's, got ",
+                               vocab_mask_dims[0]);
+      }
+
+      // store vocab mask in parameters.
+      parameters->vocab_mask = vocab_mask->DataAsSpan<int32_t>();
+    }
+
+    if (prefix_vocab_mask != nullptr) {  // prefix_vocab_mask is optional
+      const auto& vocab_mask_dims = prefix_vocab_mask->Shape().GetDims();
+      if (vocab_mask_dims.size() != 2) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME,
+                               INVALID_ARGUMENT,
+                               "Input 'prefix_vocab_mask' is expected to be 2 dimensions, got ",
+                               vocab_mask_dims.size());
+      }
+
+      // prefix_vocab_mask first dimension should be same as the first dimension of input_ids
+      if (static_cast<int>(vocab_mask_dims[0]) != static_cast<int>(dims[0])) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "input_ids and prefix_vocab_mask must have the same batch_size");
+      }
+
+      // There is dependency on vocab_size parameter, which shall be set before calling this function.
+      if (static_cast<int>(vocab_mask_dims[1]) != parameters->vocab_size) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Input 'prefix_vocab_mask' shape[1] shall be vocab_size, got ", vocab_mask_dims[1]);
+      }
+
+      // store prefix vocab mask in parameters.
+      parameters->prefix_vocab_mask = prefix_vocab_mask->DataAsSpan<int32_t>();
+    }
+
+    if (attention_mask != nullptr) {
+      const auto& dims_attn = attention_mask->Shape().GetDims();
+      if (dims_attn.size() != 2) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Input 'attention_mask' is expected to have 2 dimensions, got ", dims_attn.size());
+      }
+      if (dims_attn != dims) {
+        return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT,
+                               "Input 'attention_mask' is expected to have same shape as input_ids");
+      }
+    }
+
+    return Status::OK();
+  }
+
+
  protected:
   bool IsCuda() const { return cuda_stream_ != nullptr; }
 

onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h

@@ -122,13 +122,12 @@ template <typename T>
 Status GreedySearchBase<T>::CheckInputs(const OpKernelContextInternal& context) {
   // Input shapes:
   //   input_ids  : (batch_size, sequence_length)
-
-  const Tensor* input_ids = context.Input<Tensor>(0);
-  const auto& dims = input_ids->Shape().GetDims();
-  if (dims.size() != 2) {
-    return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Input 'input_ids' is expected to have 2 dimensions, got ",
-                           dims.size());
-  }
+  //   vocab_mask : (vocab_size) or nullptr
+  ORT_RETURN_IF_ERROR(this->CheckInputsImpl(parameters_,
+                                            context.Input<Tensor>(0),   // input_ids
+                                            context.Input<Tensor>(4),   // vocab_mask
+                                            context.Input<Tensor>(5),   // prefix_vocab_mask
+                                            nullptr));                  // attention_mask
 
   return Status::OK();
 }
@@ -145,9 +144,6 @@ Status GreedySearchBase<T>::Initialize() {
   // This flag will be updated later when the scores output exists.
   parameters_->output_scores = false;
 
-  // no_repeat_ngram_size is currently not supported in Greedy search
-  parameters_->no_repeat_ngram_size = 0;
-
   if (!this->IsCuda()) {
     // Logits processor is used in CPU only. In CUDA, cuda kernels are used instead.
     // Initialize processors after CheckInputs so that parameters_->vocab_mask is ready.

onnxruntime/contrib_ops/cpu/transformers/greedy_search_parameters.cc

@@ -13,7 +13,7 @@ void GreedySearchParameters::ParseFromAttributes(const OpKernelInfo& info) {
   eos_token_id = static_cast<int>(info.GetAttrOrDefault<int64_t>("eos_token_id", -1));
   pad_token_id = static_cast<int>(info.GetAttrOrDefault<int64_t>("pad_token_id", -1));
   decoder_start_token_id = static_cast<int>(info.GetAttrOrDefault<int64_t>("decoder_start_token_id", -1));
-  no_repeat_ngram_size = static_cast<int>(0);
+  no_repeat_ngram_size = static_cast<int>(info.GetAttrOrDefault<int64_t>("no_repeat_ngram_size", 0));
 }
 
 void GreedySearchParameters::ParseFromInputs(OpKernelContext* context) {

onnxruntime/contrib_ops/cpu/transformers/logits_processor.cc

@@ -186,57 +186,11 @@ void PrefixVocabMaskLogitsProcessor<T>::Process(const ISequences* /*sequences*/,
 }
 
 void LogitsProcessorList::Init(const BeamSearchParameters& parameters) {
-  processor_list_.clear();
-
-  if (parameters.repetition_penalty != 1.0f) {  // 1.0 means no penalty
-    repetition_penalty_processor_ = std::make_unique<RepetitionPenaltyLogitsProcessor<float>>(
-        parameters.repetition_penalty);
-    processor_list_.push_back(repetition_penalty_processor_.get());
-  }
-
-  if (parameters.no_repeat_ngram_size > 0) {
-    no_repeat_ngram_processor_ = std::make_unique<NoRepeatNGramLogitsProcessor<float>>(parameters.no_repeat_ngram_size);
-    processor_list_.push_back(no_repeat_ngram_processor_.get());
-  }
-
-  if (!parameters.vocab_mask.empty()) {
-    vocab_mask_processor_ = std::make_unique<VocabMaskLogitsProcessor<float>>(parameters.vocab_mask);
-    processor_list_.push_back(vocab_mask_processor_.get());
-  }
-
-  if (!parameters.prefix_vocab_mask.empty()) {
-    prefix_vocab_mask_processor_ = std::make_unique<PrefixVocabMaskLogitsProcessor<float>>(parameters.prefix_vocab_mask,
-                                                                                           parameters.batch_size);
-    processor_list_.push_back(prefix_vocab_mask_processor_.get());
-  }
-
-  if (parameters.min_length > 0) {
-    min_length_processor_ = std::make_unique<MinLengthLogitsProcessor<float>>(parameters.min_length,
-                                                                              parameters.eos_token_id);
-    processor_list_.push_back(min_length_processor_.get());
-  }
-
-  batch_beam_size_ = parameters.BatchBeamSize();
-  vocab_size_ = parameters.vocab_size;
+  LogitsProcessorInitImpl<BeamSearchParameters>(parameters);
 }
 
 void LogitsProcessorList::Init(const GreedySearchParameters& parameters) {
-  processor_list_.clear();
-
-  if (parameters.repetition_penalty != 1.0f) {  // 1.0 means no penalty
-    repetition_penalty_processor_ = std::make_unique<RepetitionPenaltyLogitsProcessor<float>>(
-                                      parameters.repetition_penalty);
-    processor_list_.push_back(repetition_penalty_processor_.get());
-  }
-
-  if (parameters.min_length > 0) {
-    min_length_processor_ = std::make_unique<MinLengthLogitsProcessor<float>>(parameters.min_length,
-                                                                              parameters.eos_token_id);
-    processor_list_.push_back(min_length_processor_.get());
-  }
-
-  batch_beam_size_ = parameters.BatchBeamSize();
-  vocab_size_ = parameters.vocab_size;
+  LogitsProcessorInitImpl<GreedySearchParameters>(parameters);
 }
 
 void LogitsProcessorList::Process(const ISequences* sequences,

onnxruntime/contrib_ops/cpu/transformers/logits_processor.h

@@ -104,6 +104,46 @@ class LogitsProcessorList : public ILogitsProcessorList {
   void Process(const ISequences* sequences, gsl::span<float>& next_token_scores, int step);
 
  private:
+  template<typename GenerationParametersT>
+  void LogitsProcessorInitImpl(const GenerationParametersT& parameters) {
+    processor_list_.clear();
+
+    if (parameters.repetition_penalty != 1.0f) {  // 1.0 means no penalty
+      repetition_penalty_processor_ = std::make_unique<RepetitionPenaltyLogitsProcessor<float>>(
+          parameters.repetition_penalty);
+      processor_list_.push_back(repetition_penalty_processor_.get());
+    }
+
+    if (parameters.no_repeat_ngram_size > 0) {
+      no_repeat_ngram_processor_ = std::make_unique<
+                                     NoRepeatNGramLogitsProcessor<float>
+                                   >(parameters.no_repeat_ngram_size);
+      processor_list_.push_back(no_repeat_ngram_processor_.get());
+    }
+
+    if (!parameters.vocab_mask.empty()) {
+      vocab_mask_processor_ = std::make_unique<VocabMaskLogitsProcessor<float>>(parameters.vocab_mask);
+      processor_list_.push_back(vocab_mask_processor_.get());
+    }
+
+    if (!parameters.prefix_vocab_mask.empty()) {
+      prefix_vocab_mask_processor_ = std::make_unique<
+                                       PrefixVocabMaskLogitsProcessor<float>
+                                     >(parameters.prefix_vocab_mask,
+                                       parameters.batch_size);
+      processor_list_.push_back(prefix_vocab_mask_processor_.get());
+    }
+
+    if (parameters.min_length > 0) {
+      min_length_processor_ = std::make_unique<MinLengthLogitsProcessor<float>>(parameters.min_length,
+                                                                                parameters.eos_token_id);
+      processor_list_.push_back(min_length_processor_.get());
+    }
+
+    batch_beam_size_ = parameters.BatchBeamSize();
+    vocab_size_ = parameters.vocab_size;
+  }
+
   int batch_beam_size_;
   int vocab_size_;
   InlinedVector<ILogitsProcessor<float>*> processor_list_;

onnxruntime/core/graph/contrib_ops/contrib_defs.cc

@@ -1033,6 +1033,7 @@ ONNX_MS_OPERATOR_SET_SCHEMA(GreedySearch, 1,
                                 .Attr("eos_token_id", "The id of the end-of-sequence token", AttributeProto::INT)
                                 .Attr("pad_token_id", "The id of the padding token", AttributeProto::INT)
                                 .Attr("decoder_start_token_id", "The id of the token that indicates decoding starts.", AttributeProto::INT, static_cast<int64_t>(-1))
+                                .Attr("no_repeat_ngram_size", "no repeat ngrams size", AttributeProto::INT, static_cast<int64_t>(0))
                                 .Attr("model_type", "model type: 0 for decoder only like GPT-2; 1 for encoder decoder like Bart", AttributeProto::INT, static_cast<int64_t>(0))
                                 .Attr("encoder", "The subgraph for initialization of encoder and decoder. It will be called once before decoder subgraph.", AttributeProto::GRAPH, OPTIONAL_VALUE)
                                 .Attr("decoder", "Decoder subgraph to execute in a loop.", AttributeProto::GRAPH)
@@ -1040,8 +1041,10 @@ ONNX_MS_OPERATOR_SET_SCHEMA(GreedySearch, 1,
                                 .Input(1, "max_length", "The maximum length of the sequence to be generated. Shape is (1)", "I")
                                 .Input(2, "min_length", "The minimum length below which the score of eos_token_id is set to -Inf. Shape is (1)", "I", OpSchema::Optional)
                                 .Input(3, "repetition_penalty", "The parameter for repetition penalty. Default value 1.0 means no penalty. Accepts value > 0.0. Shape is (1)", "T", OpSchema::Optional)
+                                .Input(4, "vocab_mask", "Mask of vocabulary. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (vacab_size)", "I", OpSchema::Optional)
+                                .Input(5, "prefix_vocab_mask", "Mask of vocabulary for first step. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (batch_size, vocab_size)", "I", OpSchema::Optional)
                                 .Output(0, "sequences", "Word IDs of generated sequences. Shape is (batch_size, max_sequence_length)", "I")
-                                // TODO(wy): support no_repeat_ngram_size, vocab_mask, prefix_vocab_mask, scores,
+                                // TODO(wy): support scores if needed.
                                 .TypeConstraint("T", {"tensor(float)"}, "Constrain input and output types to float tensors.")
                                 .TypeConstraint("I", {"tensor(int32)"}, "Constrain to integer types")
                                 .TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {

onnxruntime/python/tools/transformers/convert_generation.py

@@ -790,10 +790,6 @@ def convert_generation_model(args: argparse.Namespace, generation_type: Generati
     if is_greedysearch:
         if not is_gpt2:
             raise NotImplementedError("Currently only gpt2 with greedy search is supported")
-        if args.vocab_mask:
-            raise NotImplementedError("vocab_mask currently is not supported in greedy search")
-        if args.prefix_vocab_mask:
-            raise NotImplementedError("prefix_vocab_mask currently is not supported in greedy search")
         if args.output_sequences_scores:
             raise NotImplementedError("output_sequences_scores currently is not supported in greedy search")
         if args.output_token_scores:
@@ -869,12 +865,12 @@ def convert_generation_model(args: argparse.Namespace, generation_type: Generati
 
     if args.vocab_mask:
         inputs.append("vocab_mask")
-    elif not is_greedysearch:
+    else:
         inputs.append("")
 
     if args.prefix_vocab_mask:
         inputs.append("prefix_vocab_mask")
-    elif not is_greedysearch:
+    else:
         inputs.append("")
 
     if args.custom_attention_mask:
@@ -921,6 +917,7 @@ def convert_generation_model(args: argparse.Namespace, generation_type: Generati
             onnx.helper.make_attribute("eos_token_id", eos_token_id),
             onnx.helper.make_attribute("pad_token_id", pad_token_id),
             onnx.helper.make_attribute("model_type", 0 if args.model_type == "gpt2" else 1),
+            onnx.helper.make_attribute("no_repeat_ngram_size", args.no_repeat_ngram_size),
         ]
     )
     node.attribute.extend(attr_to_extend)
@@ -1235,43 +1232,6 @@ def test_gpt_model(args: argparse.Namespace, sentences: Optional[List[str]] = No
             "min_length": np.array([args.min_length], dtype=np.int32),
             "repetition_penalty": np.array([args.repetition_penalty], dtype=np.float32),
         }
-
-        logger.debug("ORT inputs", inputs)
-        result = ort_session.run(None, inputs)
-
-        if args.save_test_data:
-            test_data_dir = Path(args.output).parent.as_posix()
-            logger.debug("test_data_dir", test_data_dir)
-            from bert_test_data import output_test_data
-
-            all_inputs = [inputs]
-            for i, inputs in enumerate(all_inputs):
-                dir = os.path.join(test_data_dir, "test_data_set_" + str(i))
-                output_test_data(dir, inputs)
-
-        # Test performance
-        latency = []
-        for _ in range(args.total_runs):
-            start = time.time()
-            _ = ort_session.run(None, inputs)
-            latency.append(time.time() - start)
-
-        from benchmark_helper import get_latency_result
-
-        batch_size = input_ids.shape[0]
-        output = get_latency_result(latency, batch_size)
-
-        print("ORT outputs:")
-        sequences = result[0]
-        print("sequences", sequences)
-
-        (batch_size, max_length) = sequences.shape
-        ort_decoded_sequences = []
-        for i in range(batch_size):
-            decoded_sequence = tokenizer.decode(sequences[i], skip_special_tokens=True)
-            ort_decoded_sequences.append(decoded_sequence)
-            print(f"batch {i} sequence: {decoded_sequence}")
-
     else:
         inputs = {
             "input_ids": input_ids.cpu().numpy().astype(np.int32),
@@ -1283,51 +1243,60 @@ def test_gpt_model(args: argparse.Namespace, sentences: Optional[List[str]] = No
             "repetition_penalty": np.array([args.repetition_penalty], dtype=np.float32),
         }
 
+    if args.vocab_mask:
+        vocab_mask = np.ones((vocab_size), dtype=np.int32)
         if args.vocab_mask:
-            vocab_mask = np.ones((vocab_size), dtype=np.int32)
-            if args.vocab_mask:
-                for bad_word_id in bad_words_ids:
-                    vocab_mask[bad_word_id] = 0
-            inputs["vocab_mask"] = vocab_mask
+            for bad_word_id in bad_words_ids:
+                vocab_mask[bad_word_id] = 0
+        inputs["vocab_mask"] = vocab_mask
 
-        batch_size = input_ids.shape[0]
-        if args.prefix_vocab_mask:
-            logger.info("Use prefix vocab mask with all ones in ORT, but no corresponding setting for Torch model.")
-            prefix_vocab_mask = np.ones((batch_size, vocab_size), dtype=np.int32)
-            inputs["prefix_vocab_mask"] = prefix_vocab_mask
+    batch_size = input_ids.shape[0]
+    if args.prefix_vocab_mask:
+        logger.info("Use prefix vocab mask with all ones in ORT, but no corresponding setting for Torch model.")
+        prefix_vocab_mask = np.ones((batch_size, vocab_size), dtype=np.int32)
+        inputs["prefix_vocab_mask"] = prefix_vocab_mask
 
-        logger.debug("ORT inputs", inputs)
-        result = ort_session.run(None, inputs)
+    logger.debug("ORT inputs", inputs)
+    result = ort_session.run(None, inputs)
 
-        if args.save_test_data:
-            test_data_dir = Path(args.output).parent.as_posix()
-            logger.debug("test_data_dir", test_data_dir)
-            from bert_test_data import output_test_data
+    if args.save_test_data:
+        test_data_dir = Path(args.output).parent.as_posix()
+        logger.debug("test_data_dir", test_data_dir)
+        from bert_test_data import output_test_data
 
-            all_inputs = [inputs]
-            for i, inputs in enumerate(all_inputs):
-                dir = os.path.join(test_data_dir, "test_data_set_" + str(i))
-                output_test_data(dir, inputs)
+        all_inputs = [inputs]
+        for i, inputs in enumerate(all_inputs):
+            dir = os.path.join(test_data_dir, "test_data_set_" + str(i))
+            output_test_data(dir, inputs)
 
-        # Test performance
-        latency = []
-        for _ in range(args.total_runs):
-            start = time.time()
-            _ = ort_session.run(None, inputs)
-            latency.append(time.time() - start)
+    # Test performance
+    latency = []
+    for _ in range(args.total_runs):
+        start = time.time()
+        _ = ort_session.run(None, inputs)
+        latency.append(time.time() - start)
 
-        from benchmark_helper import get_latency_result
+    from benchmark_helper import get_latency_result
 
-        output = get_latency_result(latency, batch_size)
+    batch_size = input_ids.shape[0]
+    output = get_latency_result(latency, batch_size)
 
-        print("ORT outputs:")
-        sequences = result[0]
-        print("sequences", sequences)
-        if args.output_sequences_scores:
-            print("sequences_scores", result[1])
-        if args.output_token_scores:
-            print("scores", result[2])
+    print("ORT outputs:")
+    sequences = result[0]
+    print("sequences", sequences)
+    if args.output_sequences_scores:
+        print("sequences_scores", result[1])
+    if args.output_token_scores:
+        print("scores", result[2])
 
+    if is_greedy:
+        (batch_size, max_length) = sequences.shape
+        ort_decoded_sequences = []
+        for i in range(batch_size):
+            decoded_sequence = tokenizer.decode(sequences[i], skip_special_tokens=True)
+            ort_decoded_sequences.append(decoded_sequence)
+            print(f"batch {i} sequence: {decoded_sequence}")
+    else:
         (batch_size, num_sequences, max_length) = sequences.shape
         ort_decoded_sequences = []
         for i in range(batch_size):