microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
onnxruntime_extensions/pnp/_nlp.py
148lines · modecode
| 1 | import json |
| 2 | from collections import OrderedDict |
| 3 | |
| 4 | from ._base import ProcessingTracedModule, tensor_data_type as _dt |
| 5 | from ._torchext import create_op_function |
| 6 | from ._onnx_ops import schema |
| 7 | from .._ocos import default_opset_domain |
| 8 | |
| 9 | |
| 10 | def make_custom_op(ctx, op_type, input_names, output_names, container, operator_name=None, **kwargs): |
| 11 | op_name = container.get_unique_operator_name(op_type) if operator_name is None else operator_name |
| 12 | container.add_node(op_type, input_names, output_names, |
| 13 | op_version=1, name=op_name, op_domain=default_opset_domain(), **kwargs) |
| 14 | |
| 15 | |
| 16 | def create_bert_tokenizer(ctx, name, input_names, output_names, container, operator_name=None, **kwargs): |
| 17 | if 'hf_tok' in kwargs: |
| 18 | hf_bert_tokenizer = kwargs['hf_tok'] |
| 19 | ordered_vocab = OrderedDict(sorted(hf_bert_tokenizer.vocab.items(), key=lambda item: int(item[1]))) |
| 20 | vocab = '\n'.join(ordered_vocab.keys()) |
| 21 | attrs = dict(vocab_file=vocab) |
| 22 | # Unfortunately, there's no specific accessor function on |
| 23 | # transformers.BertTokenizer to query for strip_accents. |
| 24 | attrs['strip_accents'] = 1 if 'strip_accents' in hf_bert_tokenizer.init_kwargs and hf_bert_tokenizer.init_kwargs.get('strip_accents') else 0 |
| 25 | attrs['do_lower_case'] = 1 if hasattr(hf_bert_tokenizer, 'do_lower_case') and hf_bert_tokenizer.do_lower_case else 0 |
| 26 | elif 'vocab_file' in kwargs: |
| 27 | vocab = None |
| 28 | vocab_file = kwargs['vocab_file'] |
| 29 | with open(vocab_file, "r", encoding='utf-8') as vf: |
| 30 | lines = vf.readlines() |
| 31 | vocab = '\n'.join(lines) |
| 32 | if vocab is None: |
| 33 | raise RuntimeError("Cannot load vocabulary file {}!".format(vocab_file)) |
| 34 | attrs = dict(vocab_file=vocab) |
| 35 | if 'strip_accents' in kwargs: |
| 36 | attrs['strip_accents'] = kwargs['strip_accents'] |
| 37 | if 'do_lower_case' in kwargs: |
| 38 | attrs['do_lower_case'] = kwargs['do_lower_case'] |
| 39 | else: |
| 40 | raise RuntimeError("Need hf_tok/vocab_file parameter to build the tokenizer") |
| 41 | |
| 42 | return make_custom_op(ctx, name, input_names, |
| 43 | output_names, container, operator_name=operator_name, **attrs) |
| 44 | |
| 45 | |
| 46 | @schema(inputs=((_dt.STRING, []),), |
| 47 | outputs=((_dt.INT64, []), (_dt.INT64, []), (_dt.INT64, []))) |
| 48 | def bert_tokenizer(ctx, input_names, output_names, container, operator_name=None, **kwargs): |
| 49 | return create_bert_tokenizer(ctx, 'BertTokenizer', input_names, output_names, |
| 50 | container, operator_name=operator_name, **kwargs) |
| 51 | |
| 52 | |
| 53 | @schema(inputs=((_dt.STRING, []),), |
| 54 | outputs=((_dt.INT64, []), (_dt.INT64, []), (_dt.INT64, []))) |
| 55 | def hf_bert_tokenizer(ctx, input_names, output_names, container, operator_name=None, **kwargs): |
| 56 | return create_bert_tokenizer(ctx, 'HfBertTokenizer', input_names, output_names, |
| 57 | container, operator_name=operator_name, **kwargs) |
| 58 | |
| 59 | |
| 60 | @schema(inputs=((_dt.STRING, []),), |
| 61 | outputs=((_dt.INT64, []), (_dt.INT64, []))) |
| 62 | def gpt2_tokenize(ctx, input_names, output_names, container, operator_name=None, **kwargs): |
| 63 | if 'hf_tok' in kwargs: |
| 64 | hf_gpt2_tokenizer = kwargs['hf_tok'] |
| 65 | attrs = {'vocab': json.dumps(hf_gpt2_tokenizer.encoder, separators=(',', ':'))} |
| 66 | sorted_merges = {v_: k_ for k_, v_ in hf_gpt2_tokenizer.bpe_ranks.items()} |
| 67 | attrs['merges'] = '\n'.join("{} {}".format(*sorted_merges[n_]) for n_ in range(len(sorted_merges))) |
| 68 | elif 'vocab' in kwargs: |
| 69 | attrs = dict( |
| 70 | vocab=kwargs['vocab'], |
| 71 | merges=kwargs['merges']) |
| 72 | else: |
| 73 | raise RuntimeError("Need hf_tok/vocab parameter to build the tokenizer") |
| 74 | padding_len = -1 |
| 75 | if 'padding_length' in kwargs: |
| 76 | padding_len = kwargs['padding_length'] |
| 77 | attrs['padding_length'] = padding_len |
| 78 | |
| 79 | return make_custom_op(ctx, 'GPT2Tokenizer', input_names, |
| 80 | output_names, container, operator_name=operator_name, **attrs) |
| 81 | |
| 82 | |
| 83 | def _get_file_content(path): |
| 84 | with open(path, "rb") as file: |
| 85 | return file.read() |
| 86 | |
| 87 | |
| 88 | def _get_bound_object(func): |
| 89 | return func.__self__ |
| 90 | |
| 91 | # v1. Order of outputs - input_ids, token_type_ids, attention_mask |
| 92 | # (this is NOT consistent with the HuggingFace implementation of the tokenizer) |
| 93 | class PreHuggingFaceBert(ProcessingTracedModule): |
| 94 | def __init__(self, hf_tok=None, vocab_file=None, do_lower_case=0, strip_accents=1): |
| 95 | super(PreHuggingFaceBert, self).__init__() |
| 96 | if hf_tok is None: |
| 97 | self.onnx_bert_tokenizer = create_op_function('BertTokenizer', bert_tokenizer, |
| 98 | vocab_file=vocab_file, |
| 99 | do_lower_case=do_lower_case, |
| 100 | strip_accents=strip_accents) |
| 101 | else: |
| 102 | self.onnx_bert_tokenizer = create_op_function('BertTokenizer', bert_tokenizer, |
| 103 | hf_tok=hf_tok) |
| 104 | |
| 105 | def forward(self, text): |
| 106 | return self.onnx_bert_tokenizer(text) |
| 107 | |
| 108 | def export(self, *args, **kwargs): |
| 109 | return _get_bound_object(self.onnx_bert_tokenizer).build_model(kwargs.get('opset_version', 0), *args) |
| 110 | |
| 111 | |
| 112 | # v2. Order of outputs - input_ids, attention_mask, token_type_ids |
| 113 | # (this is consistent with the HuggingFace implementation of the tokenizer) |
| 114 | class HfBertTokenizer(ProcessingTracedModule): |
| 115 | def __init__(self, hf_tok=None, vocab_file=None, do_lower_case=0, strip_accents=1): |
| 116 | super(HfBertTokenizer, self).__init__() |
| 117 | if hf_tok is None: |
| 118 | self.onnx_bert_tokenizer = create_op_function('HfBertTokenizer', hf_bert_tokenizer, |
| 119 | vocab_file=vocab_file, |
| 120 | do_lower_case=do_lower_case, |
| 121 | strip_accents=strip_accents) |
| 122 | else: |
| 123 | self.onnx_bert_tokenizer = create_op_function('HfBertTokenizer', hf_bert_tokenizer, |
| 124 | hf_tok=hf_tok) |
| 125 | |
| 126 | def forward(self, text): |
| 127 | return self.onnx_bert_tokenizer(text) |
| 128 | |
| 129 | def export(self, *args, **kwargs): |
| 130 | return _get_bound_object(self.onnx_bert_tokenizer).build_model(kwargs.get('opset_version', 0), *args) |
| 131 | |
| 132 | |
| 133 | class PreHuggingFaceGPT2(ProcessingTracedModule): |
| 134 | def __init__(self, hf_tok=None, vocab_file=None, merges_file=None, padding_length=-1): |
| 135 | super(PreHuggingFaceGPT2, self).__init__() |
| 136 | if hf_tok is None: |
| 137 | self.onnx_gpt2_tokenize = create_op_function('GPT2Tokenizer', gpt2_tokenize, |
| 138 | vocab=_get_file_content(vocab_file), |
| 139 | merges=_get_file_content(merges_file), |
| 140 | padding_length=padding_length) |
| 141 | else: |
| 142 | self.onnx_gpt2_tokenize = create_op_function('GPT2Tokenizer', gpt2_tokenize, hf_tok=hf_tok) |
| 143 | |
| 144 | def forward(self, text): |
| 145 | return self.onnx_gpt2_tokenize(text) |
| 146 | |
| 147 | def export(self, *args, **kwargs): |
| 148 | return _get_bound_object(self.onnx_gpt2_tokenize).build_model(kwargs.get('opset_version', 0), *args) |
| 149 | |