""" This script is used to convert HuggingFace models to ONNX format and optionally quantize the model using dynamic quantization. This is courtesy of Joshua aka @Xenova """ import json import os import shutil from dataclasses import dataclass, field from typing import Optional, Set import onnx from onnxruntime.quantization import QuantType, quantize_dynamic from optimum.exporters.onnx import export_models, main_export from optimum.exporters.tasks import TasksManager from tqdm import tqdm from transformers import AutoConfig, AutoTokenizer, HfArgumentParser DEFAULT_QUANTIZE_PARAMS = { "per_channel": True, "reduce_range": True, } MODEL_SPECIFIC_QUANTIZE_PARAMS = { # Decoder-only models "codegen": { "per_channel": False, "reduce_range": False, }, "gpt2": { "per_channel": False, "reduce_range": False, }, "gpt_bigcode": { "per_channel": False, "reduce_range": False, }, "gptj": { "per_channel": False, "reduce_range": False, }, "gpt-neo": { "per_channel": False, "reduce_range": False, }, "gpt-neox": { "per_channel": False, "reduce_range": False, }, "mpt": { "per_channel": False, "reduce_range": False, }, "bloom": { "per_channel": False, "reduce_range": False, }, "llama": { "per_channel": False, "reduce_range": False, }, "opt": { "per_channel": False, "reduce_range": False, }, "mistral": { "per_channel": False, "reduce_range": False, }, "falcon": { "per_channel": False, "reduce_range": False, }, "phi": { "per_channel": False, "reduce_range": False, }, "qwen2": { "per_channel": False, "reduce_range": False, }, # Encoder-decoder models "whisper": { "per_channel": False, "reduce_range": False, }, "vision-encoder-decoder": { "per_channel": False, "reduce_range": False, }, # Encoder-only models "owlv2": { "per_channel": False, "reduce_range": False, }, } MODELS_WITHOUT_TOKENIZERS = [ "wav2vec2", "wav2vec2-bert", "wavlm", "hubert", ] @dataclass class ConversionArguments: """ Arguments used for converting HuggingFace models to onnx. """ model_id: str = field(metadata={"help": "Model identifier"}) tokenizer_id: str = field(default=None, metadata={"help": "Tokenizer identifier (if different to `model_id`)"}) quantize: bool = field(default=False, metadata={"help": "Whether to quantize the model."}) output_parent_dir: str = field( default="./models/", metadata={"help": "Path where the converted model will be saved to."} ) task: Optional[str] = field( default="auto", metadata={ "help": ( "The task to export the model for. If not specified, the task will be auto-inferred based on the model. Available tasks depend on the model, but are among:" f" {str(TasksManager.get_all_tasks())}. For decoder models, use `xxx-with-past` to export the model using past key values in the decoder." ) }, ) opset: int = field( default=None, metadata={ "help": ( "If specified, ONNX opset version to export the model with. Otherwise, the default opset will be used." ) }, ) device: str = field(default="cpu", metadata={"help": "The device to use to do the export."}) skip_validation: bool = field(default=False, metadata={"help": "Whether to skip validation of the converted model"}) per_channel: bool = field(default=None, metadata={"help": "Whether to quantize weights per channel"}) reduce_range: bool = field( default=None, metadata={ "help": "Whether to quantize weights with 7-bits. It may improve the accuracy for some models running on non-VNNI machine, especially for per-channel mode" }, ) output_attentions: bool = field( default=False, metadata={ "help": "Whether to output attentions from the model. NOTE: This is only supported for whisper models right now." }, ) split_modalities: bool = field( default=False, metadata={ "help": "Whether to split multimodal models. NOTE: This is only supported for CLIP models right now." }, ) trust_remote_code: bool = field( default=False, metadata={ "help": "Allows to use custom code for the modeling hosted in the model repository. This option should only be set for repositories" "you trust and in which you have read the code, as it will execute on your local machine arbitrary code present in the model repository." }, ) custom_onnx_configs: str = field( default=None, metadata={ "help": "Experimental usage: override the default ONNX config used for the given model. This argument may be useful for advanced users " "that desire a finer-grained control on the export." }, ) def get_operators(model: onnx.ModelProto) -> Set[str]: operators = set() def traverse_graph(graph): for node in graph.node: operators.add(node.op_type) for attr in node.attribute: if attr.type == onnx.AttributeProto.GRAPH: subgraph = attr.g traverse_graph(subgraph) traverse_graph(model.graph) return operators def quantize(model_names_or_paths, **quantize_kwargs): """ Quantize the weights of the model from float32 to int8 to allow very efficient inference on modern CPU Uses unsigned ints for activation values, signed ints for weights, per https://onnxruntime.ai/docs/performance/quantization.html#data-type-selection it is faster on most CPU architectures Args: onnx_model_path: Path to location the exported ONNX model is stored Returns: The Path generated for the quantized """ quantize_config = dict(**quantize_kwargs, per_model_config={}) for model in tqdm(model_names_or_paths, desc="Quantizing"): directory_path = os.path.dirname(model) file_name_without_extension = os.path.splitext(os.path.basename(model))[0] # NOTE: # As of 2023/04/20, the current latest version of onnxruntime-web is 1.14.0, and does not support INT8 weights for Conv layers. # For this reason, we choose model weight types to ensure compatibility with onnxruntime-web. # # As per docs, signed weight type (QInt8) is faster on most CPUs, so, we use that unless the model contains a Conv layer. # For more information, see: # - https://github.com/microsoft/onnxruntime/issues/3130#issuecomment-1105200621 # - https://github.com/microsoft/onnxruntime/issues/2339 loaded_model = onnx.load_model(model) op_types = get_operators(loaded_model) weight_type = QuantType.QUInt8 if "Conv" in op_types else QuantType.QInt8 quantize_dynamic( model_input=model, model_output=os.path.join(directory_path, f"{file_name_without_extension}_quantized.onnx"), weight_type=weight_type, # TODO allow user to specify these # op_types_to_quantize=['MatMul', 'Add', 'Conv'], extra_options=dict(EnableSubgraph=True), **quantize_kwargs, ) quantize_config["per_model_config"][file_name_without_extension] = dict( op_types=list(op_types), weight_type=str(weight_type), ) # Save quantization config with open(os.path.join(directory_path, "quantize_config.json"), "w") as fp: json.dump(quantize_config, fp, indent=4) def main(): parser = HfArgumentParser((ConversionArguments,)) (conv_args,) = parser.parse_args_into_dataclasses() model_id = conv_args.model_id tokenizer_id = conv_args.tokenizer_id or model_id output_model_folder = os.path.join(conv_args.output_parent_dir, model_id) # Create output folder os.makedirs(output_model_folder, exist_ok=True) from_pretrained_kwargs = dict( trust_remote_code=conv_args.trust_remote_code, ) # Saving the model config config = AutoConfig.from_pretrained(model_id, **from_pretrained_kwargs) custom_kwargs = {} if conv_args.custom_onnx_configs is not None: if conv_args.task == "auto": raise Exception("`--task` must be set when exporting with `--custom_onnx_configs`") custom_onnx_configs = json.loads(conv_args.custom_onnx_configs) for key in custom_onnx_configs: onnx_configs = TasksManager._SUPPORTED_MODEL_TYPE[custom_onnx_configs[key]]["onnx"] mapping = onnx_configs[conv_args.task] custom_onnx_configs[key] = mapping.func(config, **mapping.keywords) custom_kwargs["custom_onnx_configs"] = custom_onnx_configs tokenizer = None try: # Load tokenizer tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, **from_pretrained_kwargs) # To avoid inserting all chat templates into tokenizers.js, we save the chat template # to the tokenizer_config.json file, and load it when the tokenizer is loaded. if getattr(tokenizer, "chat_template", None) is None and getattr(tokenizer, "use_default_system_prompt", False): # No chat template specified, and we use the default setattr(tokenizer, "chat_template", tokenizer.default_chat_template) except KeyError: pass # No Tokenizer except Exception as e: if config.model_type not in MODELS_WITHOUT_TOKENIZERS: raise e core_export_kwargs = dict( opset=conv_args.opset, device=conv_args.device, trust_remote_code=conv_args.trust_remote_code, **custom_kwargs, ) export_kwargs = dict( model_name_or_path=model_id, output=output_model_folder, task=conv_args.task, do_validation=not conv_args.skip_validation, library_name="transformers", **core_export_kwargs, ) # Handle special cases if config.model_type == "marian": from .extra.marian import generate_tokenizer_json tokenizer_json = generate_tokenizer_json(model_id, tokenizer) with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp: json.dump(tokenizer_json, fp, indent=4) elif config.model_type == "esm": from .extra.esm import generate_fast_tokenizer fast_tokenizer = generate_fast_tokenizer(tokenizer) fast_tokenizer.save(os.path.join(output_model_folder, "tokenizer.json")) elif config.model_type == "whisper": if conv_args.output_attentions: from .extra.whisper import get_main_export_kwargs export_kwargs.update(**get_main_export_kwargs(config, "automatic-speech-recognition")) elif config.model_type in ("wav2vec2", "wav2vec2-bert", "hubert"): if tokenizer is not None: from .extra.wav2vec2 import generate_tokenizer_json tokenizer_json = generate_tokenizer_json(tokenizer) with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp: json.dump(tokenizer_json, fp, indent=4) elif config.model_type == "vits": if tokenizer is not None: from .extra.vits import generate_tokenizer_json tokenizer_json = generate_tokenizer_json(tokenizer) with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp: json.dump(tokenizer_json, fp, indent=4) elif config.model_type == "speecht5": # TODO allow user to specify vocoder path export_kwargs["model_kwargs"] = {"vocoder": "microsoft/speecht5_hifigan"} if tokenizer is not None: from .extra.speecht5 import generate_tokenizer_json tokenizer_json = generate_tokenizer_json(tokenizer) with open(os.path.join(output_model_folder, "tokenizer.json"), "w", encoding="utf-8") as fp: json.dump(tokenizer_json, fp, indent=4) elif config.model_type in ("owlvit", "owlv2"): # Override default batch size to 1, needed because non-maximum suppression is performed for exporting. # For more information, see https://github.com/huggingface/optimum/blob/e3b7efb1257c011db907ef40ab340e795cc5684c/optimum/exporters/onnx/model_configs.py#L1028-L1032 export_kwargs["batch_size"] = 1 else: pass # TODO # Step 1. convert huggingface model to onnx if not conv_args.split_modalities: main_export(**export_kwargs) else: custom_export_kwargs = dict( output_dir=output_model_folder, **core_export_kwargs, ) if config.model_type == "clip": # Handle special case for exporting text and vision models separately from transformers.models.clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection from .extra.clip import CLIPTextModelWithProjectionOnnxConfig, CLIPVisionModelWithProjectionOnnxConfig text_model = CLIPTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs) vision_model = CLIPVisionModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs) export_models( models_and_onnx_configs={ "text_model": (text_model, CLIPTextModelWithProjectionOnnxConfig(text_model.config)), "vision_model": (vision_model, CLIPVisionModelWithProjectionOnnxConfig(vision_model.config)), }, **custom_export_kwargs, ) elif config.model_type == "siglip": # Handle special case for exporting text and vision models separately from transformers.models.siglip import SiglipTextModel, SiglipVisionModel from .extra.siglip import SiglipTextModelOnnxConfig, SiglipVisionModelOnnxConfig text_model = SiglipTextModel.from_pretrained(model_id, **from_pretrained_kwargs) vision_model = SiglipVisionModel.from_pretrained(model_id, **from_pretrained_kwargs) export_models( models_and_onnx_configs={ "text_model": (text_model, SiglipTextModelOnnxConfig(text_model.config)), "vision_model": (vision_model, SiglipVisionModelOnnxConfig(vision_model.config)), }, **custom_export_kwargs, ) # TODO: Enable once https://github.com/huggingface/optimum/pull/1552 is merged # elif config.model_type == 'clap': # # Handle special case for exporting text and audio models separately # from .extra.clap import ClapTextModelWithProjectionOnnxConfig, ClapAudioModelWithProjectionOnnxConfig # from transformers.models.clap import ClapTextModelWithProjection, ClapAudioModelWithProjection # text_model = ClapTextModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs) # audio_model = ClapAudioModelWithProjection.from_pretrained(model_id, **from_pretrained_kwargs) # export_models( # models_and_onnx_configs={ # "text_model": (text_model, ClapTextModelWithProjectionOnnxConfig(text_model.config)), # "audio_model": (audio_model, ClapAudioModelWithProjectionOnnxConfig(audio_model.config)), # }, # **custom_export_kwargs, # ) else: raise Exception(f"Unable to export {config.model_type} model with `--split_modalities`.") # Step 2. (optional, recommended) quantize the converted model for fast inference and to reduce model size. if conv_args.quantize: # Update quantize config with model specific defaults quantize_config = MODEL_SPECIFIC_QUANTIZE_PARAMS.get(config.model_type, DEFAULT_QUANTIZE_PARAMS) # Update if user specified values if conv_args.per_channel is not None: quantize_config["per_channel"] = conv_args.per_channel if conv_args.reduce_range is not None: quantize_config["reduce_range"] = conv_args.reduce_range quantize( [ os.path.join(output_model_folder, x) for x in os.listdir(output_model_folder) if x.endswith(".onnx") and not x.endswith("_quantized.onnx") ], **quantize_config, ) # Step 3. Move .onnx files to the 'onnx' subfolder os.makedirs(os.path.join(output_model_folder, "onnx"), exist_ok=True) for file in os.listdir(output_model_folder): if file.endswith((".onnx", ".onnx_data")): shutil.move(os.path.join(output_model_folder, file), os.path.join(output_model_folder, "onnx", file)) # Step 4. Update the generation config if necessary if config.model_type == "whisper": from transformers import GenerationConfig from .extra.whisper import get_alignment_heads generation_config = GenerationConfig.from_pretrained(model_id, **from_pretrained_kwargs) generation_config.alignment_heads = get_alignment_heads(config) generation_config.save_pretrained(output_model_folder) if __name__ == "__main__": main()