diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md
index 19caa8a68a..9c5c60114d 100644
--- a/docs/ContribOperators.md
+++ b/docs/ContribOperators.md
@@ -1569,11 +1569,13 @@ This version of the operator has been available since version 1 of the 'com.micr
The id of the end-of-sequence token
model_type : int
model type: 0 for decoder only like GPT-2; 1 for encoder decoder like Bart
+no_repeat_ngram_size : int
+no repeat ngrams size
pad_token_id : int (required)
The id of the padding token
-#### Inputs (2 - 4)
+#### Inputs (2 - 6)
- input_ids : I
@@ -1584,6 +1586,10 @@ This version of the operator has been available since version 1 of the 'com.micr
- The minimum length below which the score of eos_token_id is set to -Inf. Shape is (1)
- repetition_penalty (optional) : T
- The parameter for repetition penalty. Default value 1.0 means no penalty. Accepts value > 0.0. Shape is (1)
+- vocab_mask (optional) : I
+- Mask of vocabulary. Words that masked with 0 are not allowed to be generated, and 1 is allowed. Shape is (vacab_size)
+- prefix_vocab_mask (optional) : I
+- 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)
#### Outputs
diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md
index 4f5f7341ee..e5123f172c 100644
--- a/docs/OperatorKernels.md
+++ b/docs/OperatorKernels.md
@@ -412,7 +412,7 @@ Do not modify directly.*
|FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|GatherND|*in* data:**T**
*in* indices:**Tind**
*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)
**Tind** = tensor(int32), tensor(int64)|
|Gelu|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
-|GreedySearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*out* sequences:**I**|1+|**T** = tensor(float)|
+|GreedySearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*in* vocab_mask:**I**
*in* prefix_vocab_mask:**I**
*out* sequences:**I**|1+|**T** = tensor(float)|
|GridSample|*in* X:**T1**
*in* Grid:**T1**
*out* Y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(float)|
|Inverse|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|MatMulInteger16|*in* A:**T1**
*in* B:**T2**
*out* Y:**T3**|1+|**T1** = tensor(int16)
**T2** = tensor(int16)
**T3** = tensor(int32)|
@@ -770,7 +770,7 @@ Do not modify directly.*
|FusedConv|*in* X:**T**
*in* W:**T**
*in* B:**T**
*in* Z:**T**
*out* Y:**T**|1+|**T** = tensor(float)|
|FusedMatMul|*in* A:**T**
*in* B:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(double), tensor(float), tensor(float16)|
|Gelu|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
-|GreedySearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*out* sequences:**I**|1+|**T** = tensor(float), tensor(float16)|
+|GreedySearch|*in* input_ids:**I**
*in* max_length:**I**
*in* min_length:**I**
*in* repetition_penalty:**T**
*in* vocab_mask:**I**
*in* prefix_vocab_mask:**I**
*out* sequences:**I**|1+|**T** = tensor(float), tensor(float16)|
|GridSample|*in* X:**T1**
*in* Grid:**T1**
*out* Y:**T2**|1+|**T1** = tensor(float)
**T2** = tensor(float)|
|Inverse|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
|Irfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)|
diff --git a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h
index 6b7a1b9c87..6653df838e 100644
--- a/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h
+++ b/onnxruntime/contrib_ops/cpu/transformers/beam_search_impl_base.h
@@ -177,68 +177,11 @@ Status BeamSearchBase::CheckInputs(const OpKernelContextInternal& context) {
// Input shapes:
// input_ids : (batch_size, sequence_length)
// vocab_mask : (vocab_size) or nullptr
-
- const Tensor* input_ids = context.Input(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(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(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();
- }
-
- const Tensor* prefix_vocab_mask = context.Input(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(vocab_mask_dims[0]) != static_cast(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(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();
- }
-
- const Tensor* attention_mask = context.Input(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(0), // input_ids
+ context.Input(7), // vocab_mask
+ context.Input(8), // prefix_vocab_mask
+ context.Input(10))); // attention_mask
return Status::OK();
}
diff --git a/onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h
index b7f88e5ea8..7889d5bd4b 100644
--- a/onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h
+++ b/onnxruntime/contrib_ops/cpu/transformers/generate_impl_base.h
@@ -79,6 +79,77 @@ class GenerateBase {
return Status::OK();
}
+ template
+ 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(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();
+ }
+
+ 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(vocab_mask_dims[0]) != static_cast(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(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();
+ }
+
+ 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; }
diff --git a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h
index de68c84708..ec408dc86f 100644
--- a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h
+++ b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_impl_base.h
@@ -122,13 +122,12 @@ template
Status GreedySearchBase::CheckInputs(const OpKernelContextInternal& context) {
// Input shapes:
// input_ids : (batch_size, sequence_length)
-
- const Tensor* input_ids = context.Input(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(0), // input_ids
+ context.Input(4), // vocab_mask
+ context.Input(5), // prefix_vocab_mask
+ nullptr)); // attention_mask
return Status::OK();
}
@@ -145,9 +144,6 @@ Status GreedySearchBase::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.
diff --git a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_parameters.cc b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_parameters.cc
index b851bff1b5..5496511de5 100644
--- a/onnxruntime/contrib_ops/cpu/transformers/greedy_search_parameters.cc
+++ b/onnxruntime/contrib_ops/cpu/transformers/greedy_search_parameters.cc
@@ -13,7 +13,7 @@ void GreedySearchParameters::ParseFromAttributes(const OpKernelInfo& info) {
eos_token_id = static_cast(info.GetAttrOrDefault("eos_token_id", -1));
pad_token_id = static_cast(info.GetAttrOrDefault("pad_token_id", -1));
decoder_start_token_id = static_cast(info.GetAttrOrDefault("decoder_start_token_id", -1));
- no_repeat_ngram_size = static_cast(0);
+ no_repeat_ngram_size = static_cast(info.GetAttrOrDefault("no_repeat_ngram_size", 0));
}
void GreedySearchParameters::ParseFromInputs(OpKernelContext* context) {
diff --git a/onnxruntime/contrib_ops/cpu/transformers/logits_processor.cc b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.cc
index d3487dfe6f..df2a2215c2 100644
--- a/onnxruntime/contrib_ops/cpu/transformers/logits_processor.cc
+++ b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.cc
@@ -186,57 +186,11 @@ void PrefixVocabMaskLogitsProcessor::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>(
- 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>(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>(parameters.vocab_mask);
- processor_list_.push_back(vocab_mask_processor_.get());
- }
-
- if (!parameters.prefix_vocab_mask.empty()) {
- prefix_vocab_mask_processor_ = std::make_unique>(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>(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(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>(
- parameters.repetition_penalty);
- processor_list_.push_back(repetition_penalty_processor_.get());
- }
-
- if (parameters.min_length > 0) {
- min_length_processor_ = std::make_unique>(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(parameters);
}
void LogitsProcessorList::Process(const ISequences* sequences,
diff --git a/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h
index 00ea85a278..1a2fba19bf 100644
--- a/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h
+++ b/onnxruntime/contrib_ops/cpu/transformers/logits_processor.h
@@ -104,6 +104,46 @@ class LogitsProcessorList : public ILogitsProcessorList {
void Process(const ISequences* sequences, gsl::span& next_token_scores, int step);
private:
+ template
+ 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>(
+ 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
+ >(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>(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
+ >(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>(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*> processor_list_;
diff --git a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc b/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
index f01342e969..4a79792559 100644
--- a/onnxruntime/core/graph/contrib_ops/contrib_defs.cc
+++ b/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(-1))
+ .Attr("no_repeat_ngram_size", "no repeat ngrams size", AttributeProto::INT, static_cast(0))
.Attr("model_type", "model type: 0 for decoder only like GPT-2; 1 for encoder decoder like Bart", AttributeProto::INT, static_cast(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) {
diff --git a/onnxruntime/python/tools/transformers/convert_generation.py b/onnxruntime/python/tools/transformers/convert_generation.py
index 8a2787315a..6bea5ec3bd 100644
--- a/onnxruntime/python/tools/transformers/convert_generation.py
+++ b/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):