gpt2bs.py

####################################################################################
#
# !!! This script is replaced by the latest onnxruntime contrib op solution, which is
# https://github.com/microsoft/onnxruntime/blob/ad9d2e2e891714e0911ccc3fa8b70f42025b4d56/onnxruntime/python/tools/transformers/convert_beam_search.py
#
###################################################################################

importos
importonnx
importnumpy
importargparse
importonnxruntimeas_ort
importonnxruntime_extensionsas_ortex
fromtransformersimportAutoConfig
fromdistutils.versionimportStrictVersion

ifStrictVersion(_ort.__version__) <StrictVersion('1.8.1'):
raiseRuntimeError('Full GPT-2 model is only available on onxruntime 1.8.1 and higher version.')

model_name_or_path="gpt2"
device="cpu"
default_beam_width=4
default_batch_size=1
onnx_model_path="gpt2_one_step_search.onnx"
gpt2_full_model_path="gpt2_full.onnx"

# Create a cache directory to store pretrained model.
cache_dir=os.path.expanduser('~/.cache/huggingface/')
ifnotos.path.exists(cache_dir):
cache_dir=os.path.join(".", "cache_models")
ifnotos.path.exists(cache_dir):
os.makedirs(cache_dir)


def_extract_endict(tokenizer_endict):
_1, _2= [tokenizer_endict.get(ky_) forky_in ('input_ids', 'attention_mask')]
return_1.astype(numpy.int64), _2.astype(numpy.float32)


defget_tokenizer(model_name_or_path, enable_tokenizer, cache_dir):
fromtransformersimportGPT2Tokenizer# noqa
fromonnxruntime_extensions.onnxprocessimportbuild_customop_model, pyfunc_from_model

tokenizer=GPT2Tokenizer.from_pretrained(model_name_or_path, cache_dir=cache_dir)

tokenizer.padding_side="left"
tokenizer.pad_token=tokenizer.eos_token
ifenable_tokenizer:
gpt2_encoder_model_path='./gpt2_tok.onnx'
build_customop_model('GPT2Tokenizer', gpt2_encoder_model_path, model=tokenizer)
returntokenizer, pyfunc_from_model(gpt2_encoder_model_path)
else:
returntokenizer, None


defconvert_gpt2():
fromonnxruntime.transformers.gpt2_beamsearch_helperimportGpt2BeamSearchHelper, GPT2LMHeadModel_BeamSearchStep
config=AutoConfig.from_pretrained(model_name_or_path, cache_dir=cache_dir)
model=GPT2LMHeadModel_BeamSearchStep.from_pretrained(model_name_or_path,
config=config,
batch_size=default_batch_size,
beam_size=default_beam_width,
cache_dir=cache_dir)
model.eval().to(device)
Gpt2BeamSearchHelper.export_onnx(model, device, onnx_model_path)


definference_and_dump_full_model(tokenizer, func_tokenizer, input_text, num_tokens_to_produce=30):
fromonnxruntime_extensions.onnxprocessimporttrace_for_onnx, pyfunc_from_model

# a hot fix for the dynamic axes of the converted model
gpt2_core=onnx.load_model(onnx_model_path)
for_viingpt2_core.graph.output:
if_vi.name=='last_state':
_vi.type.tensor_type.shape.dim[1].dim_param='seq_len'

func_one_step=pyfunc_from_model(gpt2_core)
config=AutoConfig.from_pretrained(model_name_or_path, cache_dir=cache_dir)
num_attention_heads=config.n_head
hidden_size=config.n_embd
num_layer=config.n_layer
iffunc_tokenizerisNone:
input_ids, attention_mask=_extract_endict(tokenizer(input_text, padding=True, return_tensors='np'))
withtrace_for_onnx(input_ids, attention_mask,
num_tokens_to_produce,
names=["input_ids", "attention_mask", "out_token_num"], target_opset=12) astc_sess:
input_ids, attention_mask, num_tokens=tc_sess.get_inputs()
input_ids.symbolic_shape= ['batch_size', 'seq_len']
attention_mask.symbolic_shape= ['batch_size', 'seq_len']

full_model=_beam_search(tokenizer, func_one_step,
num_attention_heads, hidden_size, num_layer,
tc_sess, num_tokens, input_ids, attention_mask)
else:
withtrace_for_onnx(input_text, num_tokens_to_produce,
names=func_tokenizer.input_names, target_opset=12) astc_sess:
inputs, num_tokens=tc_sess.get_inputs()
input_ids, attention_mask=func_tokenizer(inputs, padding=True)
full_model=_beam_search(tokenizer, func_one_step,
num_attention_heads, hidden_size, num_layer,
tc_sess, num_tokens, input_ids, attention_mask)

_ortex.optimize_model(full_model, gpt2_full_model_path)


def_beam_search(tokenizer, func_one_step,
num_attention_heads, hidden_size, num_layer, tc_sess, num_tokens, input_ids, attention_mask):
fromonnxruntime_extensions.onnxprocessimporttorch_wrapperastorch

ifattention_mask.dtypeisnottorch.float32:
attention_mask=attention_mask.type(torch.float)
position_ids= (attention_mask.long().cumsum(-1) -1)
batch_size=default_batch_size
past_shape= [2, batch_size, num_attention_heads, 0, hidden_size//num_attention_heads]
empty_past= []
for_inrange(num_layer):
empty_past.append(torch.empty(*past_shape).type(torch.float32).to(device))

beam_select_idx=torch.zeros([1, batch_size]).long()
input_log_probs=torch.zeros([batch_size, 1])
input_unfinished_sents=torch.ones([batch_size, 1], dtype=torch.bool)
prev_step_scores=torch.zeros([batch_size, 1])
beam_size=default_beam_width
prev_step_results=input_ids.clone().detach().to(device)

cfg=torch.control_flow()
forstatesincfg.loop(num_tokens, torch.tensor(True), input_ids, position_ids,
attention_mask, beam_select_idx, input_log_probs,
input_unfinished_sents, prev_step_results, prev_step_scores, *empty_past):
step=states[0]
states[1].symbolic_shape= ['batch_size', 'seq_len']
states[2].symbolic_shape= ['batch_size', 'seq_len']
states[3].symbolic_shape= ['batch_size', 'all_seq_len']
states[4].symbolic_shape= [1, 'batch_size']

# prev_step_results
states[7].symbolic_shape= ['batch_size', 'total_seq_len']

forst_instates[-num_layer:]:
st_.symbolic_shape= [2, 'batch_size', num_attention_heads,
'past_seq_len', hidden_size//num_attention_heads]

prev_attention_mask=states[3]
outputs=func_one_step(*states[1:])
last_state=outputs[0].clone().detach().cpu()
input_ids=last_state.reshape([batch_size*beam_size, -1]).to(device)

input_unfinished_sents_id=-3
prev_step_results=outputs[-2].clone().detach().to(device)
# position_ids = (torch.tensor([context_length + step - 1
# ]).unsqueeze(0).repeat(batch_size * beam_size, 1).to(device))
position_ids=torch.zeros([batch_size*beam_size, 1], dtype=torch.int64) +attention_mask.size()[-1]
factor= (~step.type(torch.bool)).type(torch.int64)
prev_attention_mask=prev_attention_mask.repeat(factor* (batch_size*beam_size-1) +1, 1).to(device)
attention_mask=torch.cat(
 [
prev_attention_mask,
torch.ones([batch_size*beam_size, 1], dtype=torch.float),
 ],
1,
 ).to(device)

beam_select_idx=outputs[input_unfinished_sents_id-2].clone().detach().to(device)
input_log_probs=outputs[input_unfinished_sents_id-1].clone().detach().to(device)
input_unfinished_sents=outputs[input_unfinished_sents_id].clone().detach().to(device)
prev_step_scores=outputs[-1].clone().detach().to(device)

past= []
foriinrange(num_layer):
past_i=outputs[i+1].clone().detach()
past.append(past_i.to(device))

any_unfinished=input_unfinished_sents.any()
input_ids.symbolic_shape= ['total_batch_size', 'seq_len']
position_ids.symbolic_shape= ['total_batch_size', 'seq_len']
attention_mask.symbolic_shape= ['total_batch_size', 'all_seq_len']
prev_step_results.symbolic_shape= ['total_batch_size', 'step_seq_len']
forst_inpast:
st_.symbolic_shape= [2, 'total_batch_size', num_attention_heads,
'all_seq_len', hidden_size//num_attention_heads]
cfg.flow_output(any_unfinished, input_ids,
position_ids, attention_mask, beam_select_idx,
input_log_probs, input_unfinished_sents, prev_step_results, prev_step_scores, *past)

result_id=6
all_token_ids=cfg.finalize()[result_id]
mdl=tc_sess.save_as_onnx(None, all_token_ids)
print(tokenizer.decode(all_token_ids.t[0], skip_special_tokens=True))
returnmdl


defverify_bsfull_model(input_text, tokenizer, enable_tokenizer):
importtime
fromonnxruntime_extensionsimportPyOrtFunction
gpt2_all=PyOrtFunction.from_model(gpt2_full_model_path)
gpt2_all._ensure_ort_session()
ifenable_tokenizer:
start_time=time.perf_counter()
outputs=gpt2_all(input_text, 30)
else:
input_ids, attention_mask=_extract_endict(tokenizer(input_text, padding=True, return_tensors='np'))
start_time=time.perf_counter()
outputs=gpt2_all(input_ids, attention_mask, 30)
print("total time: {}".format(time.perf_counter() -start_time))
print(tokenizer.decode(outputs[0], skip_special_tokens=True))


defmain(enable_tokenizer):
tokenizer, func_tokenizer=get_tokenizer(model_name_or_path, enable_tokenizer, cache_dir)
input_text= ['best hotel in bay area.']
ifnotos.path.exists(onnx_model_path):
convert_gpt2()
inference_and_dump_full_model(tokenizer, func_tokenizer, input_text)
verify_bsfull_model(input_text, tokenizer, enable_tokenizer)


if__name__=="__main__":
parser=argparse.ArgumentParser()
parser.add_argument("--disable-tokenizer", '-d', help="No tokenizer operator for the full model",
action="store_true")
parser.add_argument("--output", '-o', help="The output file name")
args=parser.parse_args()
ifargs.outputisnotNone:
gpt2_full_model_path=args.output
main(notargs.disable_tokenizer)