onnxruntime_extensions 0.14.0__cp313-cp313-macosx_11_0_universal2.whl

onnxruntime_extensions/cvt.py

@@ -0,0 +1,306 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+###############################################################################
+
+"""
+cvt.py: Processing Graph Converter and Generator
+"""
+
+from typing import Union
+
+from ._hf_cvt import HFTokenizerConverter, HFTokenizerOnnxGraph  # noqa
+from ._ortapi2 import make_onnx_model, SingleOpGraph
+
+import os
+import numpy as np
+import tempfile
+import shutil
+
+# edit environment variables to avoid protobuf version mismatch
+os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
+
+from transformers.convert_slow_tokenizer import SpmConverter  # noqa: E402
+from transformers import AutoTokenizer  # noqa: E402
+from tokenizers import decoders, normalizers, pre_tokenizers, Regex  # noqa: E402
+
+
+OrtxTokenizer = None
+try:
+    from onnxruntime_extensions.pp_api import Tokenizer as OrtxTokenizer
+except ImportError:
+    pass
+
+_is_torch_available = False
+try:
+    import torch    # noqa
+    _is_torch_available = True
+    from ._torch_cvt import WhisperDataProcGraph
+except ImportError:
+    WhisperDataProcGraph = None
+
+
+_PRE_POST_PAIR = {'TrieTokenizer': "TrieDetokenizer"}
+
+def _get_prepend_scheme(add_prefix_space: bool, original_tokenizer) -> str:
+    if add_prefix_space:
+        prepend_scheme = "always"
+        if not getattr(original_tokenizer, "legacy", True):
+            prepend_scheme = "first"
+    else:
+        prepend_scheme = "never"
+    return prepend_scheme
+
+
+class Baichuan2Converter(SpmConverter):
+    handle_byte_fallback = True
+
+    def __init__(self, original_tokenizer):
+        super().__init__(original_tokenizer)
+        original_tokenizer.add_prefix_space = False
+
+    def vocab(self, proto):
+        vocab = [
+            (self.original_tokenizer.convert_ids_to_tokens(0), 0.0),
+            (self.original_tokenizer.convert_ids_to_tokens(1), 0.0),
+            (self.original_tokenizer.convert_ids_to_tokens(2), 0.0),
+        ]
+        vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
+        return vocab
+
+    def unk_id(self, proto):
+        unk_id = 0
+        return unk_id
+
+    def decoder(self, replacement, add_prefix_space):
+        sequence = [
+            decoders.Replace("▁", " "),
+            decoders.ByteFallback(),
+            decoders.Fuse(),
+        ]
+        if add_prefix_space:
+            sequence += [decoders.Strip(content=" ", left=1)]
+        return decoders.Sequence(sequence)
+
+    def normalizer(self, proto):
+        if getattr(self.original_tokenizer, "legacy", True):
+            sequence = []
+            if getattr(self.original_tokenizer, "add_prefix_space", True):
+                sequence += [normalizers.Prepend(prepend="▁")]
+            sequence += [normalizers.Replace(pattern=" ", content="▁")]
+            return normalizers.Sequence(sequence)
+        return None  # non-legacy, no normalizer
+
+    def pre_tokenizer(self, replacement, add_prefix_space):
+        if not getattr(self.original_tokenizer, "legacy", True):  # non-legacy, we need a replace
+            prepend_scheme = _get_prepend_scheme(add_prefix_space, self.original_tokenizer)
+            return pre_tokenizers.Metaspace(replacement=replacement, prepend_scheme=prepend_scheme, split=False)
+        else:
+            return super().pre_tokenizer(replacement, add_prefix_space)
+
+
+class ChatGlmConverter(SpmConverter):
+    def normalizer(self, proto):
+        precompiled_charsmap = proto.normalizer_spec.precompiled_charsmap
+        _normalizers = [
+            normalizers.Strip(left=False, right=True),  # stripping is important
+            normalizers.Replace(Regex(" {2,}"), "▁"),
+        ]
+        return normalizers.Sequence([normalizers.Precompiled(precompiled_charsmap)] + _normalizers)
+
+    def pre_tokenizer(self, replacement, add_prefix_space):
+        prepend_scheme = "always"
+        if hasattr(self.original_tokenizer, "legacy") and not self.original_tokenizer.legacy:
+            prepend_scheme = "first"
+        return pre_tokenizers.Metaspace(
+            replacement=replacement, add_prefix_space=add_prefix_space, prepend_scheme=prepend_scheme
+        )
+
+
+JSON_TOKEN_CONVERTERS = {
+    "BaichuanTokenizer": Baichuan2Converter,
+    "ChatGLMTokenizer": ChatGlmConverter,
+}
+
+# Save tokenizer JSON files using HuggingFace AutoTokenizer
+def convert_tokenizer(model_path, output_dir):
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+    if output_dir is None:
+        if os.path.isdir(model_path):
+            output_dir = model_path
+        else:
+            # create a temporary directory
+            output_dir = tempfile.mkdtemp()
+            tokenizer.save_pretrained(output_dir)
+        json_path = os.path.join(output_dir, "tokenizer.json")
+
+    if type(tokenizer).__name__ in JSON_TOKEN_CONVERTERS:
+        GenericSpmConverter = JSON_TOKEN_CONVERTERS[type(tokenizer).__name__]
+
+    converted = GenericSpmConverter(tokenizer).converted()
+    converted.save(json_path)
+    print(f"**Tokenizer saved to {json_path}")
+    return output_dir
+
+# Validate tokenizer files downloaded from memory
+def validate_tokenizer(model_path, output_dir):
+    test_sentence = "I like walking my cute dog\n and\x17 then, 生活的真谛是   \t\t\t\t \n\n61"
+    if OrtxTokenizer is None:
+        print("onnxruntime_extensions package was built with C API enabled, skipping tokenization test")
+    ortx_tokenizer = OrtxTokenizer(output_dir)
+    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True, use_fast=False)
+    expected_ids = tokenizer(test_sentence, return_tensors="np")["input_ids"]
+    ortx_ids = np.asarray(ortx_tokenizer.tokenize(test_sentence))
+    assert np.array_equal(expected_ids[0], ortx_ids), f"Tokenization mismatch: {expected_ids[0]} != {ortx_ids}"
+    print("Tokenization test passed")
+
+# Download tokenizer JSON files from memory
+def download_tokenizer(tokenizer_dir, output_dir):
+    try:
+        from transformers.utils import cached_file
+
+        resolved_full_file = cached_file(tokenizer_dir, "tokenizer.json")
+        resolved_config_file = cached_file(tokenizer_dir, "tokenizer_config.json")
+    except ImportError:
+        raise ValueError(f"Directory '{tokenizer_dir}' not found and transformers is not available")
+    if not os.path.exists(resolved_full_file):
+        raise FileNotFoundError(f"Downloaded HF file '{resolved_full_file}' cannot be found")
+    if os.path.dirname(resolved_full_file) != os.path.dirname(resolved_config_file):
+        raise FileNotFoundError(
+            f"Downloaded HF files '{resolved_full_file}' " f"and '{resolved_config_file}' are not in the same directory"
+        )
+
+    if output_dir is None or len(output_dir) == 0:
+        output_dir = os.path.dirname(resolved_full_file)
+        print(f"Using {output_dir} as output directory")
+        return output_dir
+    else:
+        # copy the files to the output directory
+        shutil.copy(resolved_full_file, output_dir)
+        shutil.copy(resolved_config_file, output_dir)
+        return output_dir
+
+
+def gen_processing_models(processor: Union[str, object],
+                          pre_kwargs: dict = None,
+                          post_kwargs: dict = None,
+                          opset: int = None,
+                          schema_v2: bool = False,
+                          **kwargs):
+    """
+    Generate the pre- and post-processing ONNX model, basing on the name or HF class.
+
+    Parameters
+    ----------
+    processor:
+        the HF processor/tokenizer instance, or the name (str) of a Data Processor
+        the instance is preferred, otherwise when name was given, the corresponding configuration for the processor
+        has to be provided in the kwargs
+    pre_kwargs: dict
+        Keyword arguments for generating the pre-processing model
+        WITH_DEFAULT_INPUTS: bool, add default inputs to the graph, default is True
+        CAST_TOKEN_ID: bool, add a cast op to output token IDs to be int64 if needed, default is False
+    post_kwargs: dict
+        Keyword arguments for generating the post-processing model
+    opset: int
+        the target opset version of the model
+    schema_v2: bool
+        the flag for using embedded tokenizer files; this option leverages the blob-loading functionality
+        which loads HF tokenizers from memory rather than using the tokenizer files in HF JSON format.
+    kwargs:
+        The additional arguments for generating models
+
+    Returns
+    -------
+    ONNX-Models
+        The pre- and post-processing ONNX models
+    """
+    if pre_kwargs is None and post_kwargs is None:
+        raise ValueError(
+            "Either pre_kwargs or post_kwargs should be provided. None means no processing graph output.")
+
+    # If true, we get the tokenizer JSON files by either downloading from cache or using HuggingFace AutoTokenizer
+    # to convert them, and then create an ONNX model with the JSON files as strings in the model attributes (attrs).
+    if schema_v2:
+        model_name = processor if isinstance(processor, str) else type(processor).__name__
+
+        converted_tokenizer = {"Baichuan2", "chatglm"}
+        need_convert = False
+        for token in converted_tokenizer:
+            if model_name.find(token) != -1:
+                need_convert = True
+                break
+
+        if need_convert:
+            model_dir = convert_tokenizer(model_name)
+            validate_tokenizer(model_name, None)
+        else:
+            model_dir = download_tokenizer(model_name, None)
+
+        # Load the content of tokenizer.json into a string
+        with open(f"{model_dir}/tokenizer.json", "r", encoding="utf-8") as f:
+            tokenizer_vocab = f.read()
+
+        # Load the content of tokenizer_config.json into a string
+        with open(f"{model_dir}/tokenizer_config.json", "r", encoding="utf-8") as f:
+            tokenizer_config = f.read()
+
+        # Create an ONNX model with these JSON file strings in attrs
+        g_pre, g_post = (None, None)
+        if pre_kwargs is not None:
+            # Add tokenizer_vocab and tokenizer_config to the kwargs
+            # so they are added to attrs in build_graph
+            pre_kwargs['tokenizer_vocab'] = tokenizer_vocab
+            pre_kwargs['tokenizer_config'] = tokenizer_config
+            g_pre = SingleOpGraph.build_graph("HfJsonTokenizer", **pre_kwargs)
+        if post_kwargs is not None:
+            if pre_kwargs is None:
+                cls_name = processor
+            else:
+                if processor not in _PRE_POST_PAIR:
+                    raise RuntimeError(
+                        f"Cannot locate the post processing operator name from {processor}")
+                cls_name = _PRE_POST_PAIR[processor]
+            # Add tokenizer_vocab and tokenizer_config to the kwargs
+            # so they are added to attrs in build_graph
+            post_kwargs['tokenizer_vocab'] = tokenizer_vocab
+            post_kwargs['tokenizer_config'] = tokenizer_config
+            g_post = SingleOpGraph.build_graph(cls_name, **post_kwargs)
+        return make_onnx_model(g_pre) if g_pre else None, make_onnx_model(g_post) if g_post else None
+    else:
+        if isinstance(processor, str):
+            g_pre, g_post = (None, None)
+            if pre_kwargs:
+                g_pre = SingleOpGraph.build_graph(processor, **pre_kwargs)
+            if post_kwargs:
+                if pre_kwargs is None:
+                    cls_name = processor
+                else:
+                    if processor not in _PRE_POST_PAIR:
+                        raise RuntimeError(
+                            f"Cannot locate the post processing operator name from {processor}")
+                    cls_name = _PRE_POST_PAIR[processor]
+                g_post = SingleOpGraph.build_graph(cls_name, **post_kwargs)
+            return make_onnx_model(g_pre) if g_pre else None, make_onnx_model(g_post) if g_post else None
+
+        cls_name = type(processor).__name__
+        if cls_name == "WhisperProcessor":
+            if WhisperDataProcGraph is None:
+                raise ValueError(
+                    "The Whisper processor needs torch.onnx support, please install pytorch 2.0 and above")
+            _converter = WhisperDataProcGraph(processor, opset=opset, **kwargs)
+            pre_m = _converter.pre_processing(
+                **pre_kwargs) if pre_kwargs is not None else None
+            post_m = _converter.post_processing(
+                **post_kwargs) if post_kwargs is not None else None
+            return pre_m, post_m
+        elif HFTokenizerOnnxGraph.is_supported(processor):
+            _converter = HFTokenizerOnnxGraph(processor)
+            pre_g = _converter.pre_processing(
+                **pre_kwargs) if pre_kwargs is not None else None
+            post_g = _converter.post_processing(
+                **post_kwargs) if post_kwargs is not None else None
+            return make_onnx_model(pre_g) if pre_g else None, \
+                make_onnx_model(post_g) if post_g else None
+        else:
+            raise ValueError(f"Unsupported processor/tokenizer: {cls_name}")

onnxruntime_extensions/onnxprocess/__init__.py

@@ -0,0 +1,12 @@
+"""
+override the torch importing, to dump all torch operators during the processing code.
+!!!This package depends on onnxruntime_extensions root package, but not vice versa.!!!
+, since this package fully relies on pytorch, while the onnxruntime_extensions doesn't
+"""
+
+from . _tensor import op_from_customop as pyfunc_from_custom_op
+from . _tensor import op_from_model as pyfunc_from_model
+from ._builder import build_customop_model
+from ._session import ONNXTraceSession
+
+trace_for_onnx = ONNXTraceSession.trace_for_onnx

onnxruntime_extensions/onnxprocess/_builder.py

@@ -0,0 +1,53 @@
+import json
+import pathlib
+from ._onnx_ops import make_model_ex
+from .._cuops import SingleOpGraph, GPT2Tokenizer, VectorToString
+from .._ortapi2 import default_opset_domain
+
+
+def is_path(name_or_buffer):
+    return isinstance(name_or_buffer, str) or isinstance(name_or_buffer, pathlib.Path)
+
+
+class _GPT2Tokenizer(GPT2Tokenizer):
+    @classmethod
+    def serialize_attr(cls, kwargs):
+        assert 'model' in kwargs, "Need model parameter to build the tokenizer"
+        hf_gpt2_tokenizer = kwargs['model']
+        attrs = {'vocab': json.dumps(hf_gpt2_tokenizer.encoder, separators=(',', ':'))}
+        sorted_merges = {v_: k_ for k_, v_ in hf_gpt2_tokenizer.bpe_ranks.items()}
+        attrs['merges'] = '\n'.join("{} {}".format(*sorted_merges[n_]) for n_ in range(len(sorted_merges)))
+        return attrs
+
+
+class _VectorToString(VectorToString):
+    @classmethod
+    def serialize_attr(cls, kwargs):
+        assert 'decoder' in kwargs, "Need decoder parameter to build the tokenizer"
+        decoder = kwargs['decoder']
+        remapped = {v: [k] for k, v in decoder.items()}
+        attrs = dict(map=remapped, unk='<unknown>')
+        return super().serialize_attr(attrs)
+
+
+customop_mbuilder = {
+    c_.op_type(): c_ for c_ in (
+        _GPT2Tokenizer,
+        _VectorToString
+    )
+}
+
+
+def build_customop_model(op_type, f, opset_version=11, **attrs):
+    op_class = SingleOpGraph.get_op_class(op_type)
+    if op_type in customop_mbuilder:
+        op_class = customop_mbuilder[op_type]
+
+    graph = SingleOpGraph.build_my_graph(op_class, **attrs)
+    m = make_model_ex(graph, [(default_opset_domain(), 1)], opset_version)
+    if is_path(f):
+        with open(f, 'wb') as f_:
+            f_.write(m.SerializeToString())
+    else:
+        f.write(m.SerializeToString())
+        f.flush()