From 95600adc486267699d0cccc7c76a83d4f34da271 Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Thu, 1 Feb 2024 22:36:25 +0000 Subject: [PATCH 1/4] Add support for exporting custom models --- scripts/convert.py | 142 ++++++++++++++++++++++++++++----------------- 1 file changed, 89 insertions(+), 53 deletions(-) diff --git a/scripts/convert.py b/scripts/convert.py index 78232aab7..9b5311046 100644 --- a/scripts/convert.py +++ b/scripts/convert.py @@ -188,6 +188,22 @@ class ConversionArguments: } ) + 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() @@ -281,7 +297,17 @@ def main(): os.makedirs(output_model_folder, exist_ok=True) # Saving the model config - config = AutoConfig.from_pretrained(model_id) + config = AutoConfig.from_pretrained(model_id, trust_remote_code=conv_args.trust_remote_code) + + custom_kwargs={} + if conv_args.custom_onnx_configs is not None: + custom_onnx_configs = json.loads(conv_args.custom_onnx_configs) + + for key in custom_onnx_configs: + mapping = TasksManager._SUPPORTED_MODEL_TYPE[custom_onnx_configs[key]]['onnx'][conv_args.task] + custom_onnx_configs[key] = mapping.func(config, **mapping.keywords) + + custom_kwargs['custom_onnx_configs'] = custom_onnx_configs tokenizer = None try: @@ -302,13 +328,20 @@ def main(): 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, - opset=conv_args.opset, - device=conv_args.device, do_validation=not conv_args.skip_validation, + library_name='transformers', + **core_export_kwargs, ) # Handle special cases @@ -368,63 +401,66 @@ def main(): pass # TODO # Step 1. convert huggingface model to onnx - if config.model_type == 'clip' and conv_args.split_modalities: - # Handle special case for exporting text and vision models separately - from .extra.clip import CLIPTextModelWithProjectionOnnxConfig, CLIPVisionModelWithProjectionOnnxConfig - from transformers.models.clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection - - text_model = CLIPTextModelWithProjection.from_pretrained(model_id) - vision_model = CLIPVisionModelWithProjection.from_pretrained(model_id) - - export_models( - models_and_onnx_configs={ - "text_model": (text_model, CLIPTextModelWithProjectionOnnxConfig(text_model.config)), - "vision_model": (vision_model, CLIPVisionModelWithProjectionOnnxConfig(vision_model.config)), - }, + if not conv_args.split_modalities: + main_export(**export_kwargs) + else: + custom_export_kwargs = dict( output_dir=output_model_folder, - opset=conv_args.opset, - device=conv_args.device, + **core_export_kwargs, ) - elif config.model_type == 'siglip' and conv_args.split_modalities: - # Handle special case for exporting text and vision models separately - from .extra.siglip import SiglipTextModelOnnxConfig, SiglipVisionModelOnnxConfig - from transformers.models.siglip import SiglipTextModel, SiglipVisionModel + if config.model_type == 'clip': + # Handle special case for exporting text and vision models separately + from .extra.clip import CLIPTextModelWithProjectionOnnxConfig, CLIPVisionModelWithProjectionOnnxConfig + from transformers.models.clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection - text_model = SiglipTextModel.from_pretrained(model_id) - vision_model = SiglipVisionModel.from_pretrained(model_id) + text_model = CLIPTextModelWithProjection.from_pretrained(model_id) + vision_model = CLIPVisionModelWithProjection.from_pretrained(model_id) - export_models( - models_and_onnx_configs={ - "text_model": (text_model, SiglipTextModelOnnxConfig(text_model.config)), - "vision_model": (vision_model, SiglipVisionModelOnnxConfig(vision_model.config)), - }, - output_dir=output_model_folder, - opset=conv_args.opset, - device=conv_args.device, - ) + 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, + ) - # TODO: Enable once https://github.com/huggingface/optimum/pull/1552 is merged - # elif config.model_type == 'clap' and conv_args.split_modalities: - # # 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) - # audio_model = ClapAudioModelWithProjection.from_pretrained(model_id) - - # export_models( - # models_and_onnx_configs={ - # "text_model": (text_model, ClapTextModelWithProjectionOnnxConfig(text_model.config)), - # "audio_model": (audio_model, ClapAudioModelWithProjectionOnnxConfig(audio_model.config)), - # }, - # output_dir=output_model_folder, - # opset=conv_args.opset, - # device=conv_args.device, - # ) + elif config.model_type == 'siglip': + # Handle special case for exporting text and vision models separately + from .extra.siglip import SiglipTextModelOnnxConfig, SiglipVisionModelOnnxConfig + from transformers.models.siglip import SiglipTextModel, SiglipVisionModel + + text_model = SiglipTextModel.from_pretrained(model_id) + vision_model = SiglipVisionModel.from_pretrained(model_id) + + 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) + # audio_model = ClapAudioModelWithProjection.from_pretrained(model_id) + + # 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`.') - else: - main_export(**export_kwargs) # Step 2. (optional, recommended) quantize the converted model for fast inference and to reduce model size. if conv_args.quantize: From cff27c94a14c290d0223717b7f3862b3cbc297a3 Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Thu, 1 Feb 2024 22:46:55 +0000 Subject: [PATCH 2/4] Require task to be specified --- scripts/convert.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/scripts/convert.py b/scripts/convert.py index 9b5311046..f4c0fa5e4 100644 --- a/scripts/convert.py +++ b/scripts/convert.py @@ -301,10 +301,13 @@ def main(): 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: - mapping = TasksManager._SUPPORTED_MODEL_TYPE[custom_onnx_configs[key]]['onnx'][conv_args.task] + 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 From 3d1e0ccf4107bd0853eb6c65ccf9bba2c1acc72e Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Fri, 2 Feb 2024 14:06:32 +0000 Subject: [PATCH 3/4] Pass `trust_remote_code` to all `from_pretrained` calls. --- scripts/convert.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/scripts/convert.py b/scripts/convert.py index f4c0fa5e4..35b92ae62 100644 --- a/scripts/convert.py +++ b/scripts/convert.py @@ -296,8 +296,12 @@ def main(): # 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, trust_remote_code=conv_args.trust_remote_code) + config = AutoConfig.from_pretrained(model_id, **from_pretrained_kwargs) custom_kwargs={} if conv_args.custom_onnx_configs is not None: @@ -315,7 +319,7 @@ def main(): tokenizer = None try: # Load tokenizer - tokenizer = AutoTokenizer.from_pretrained(tokenizer_id) + 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. @@ -417,8 +421,8 @@ def main(): from .extra.clip import CLIPTextModelWithProjectionOnnxConfig, CLIPVisionModelWithProjectionOnnxConfig from transformers.models.clip import CLIPTextModelWithProjection, CLIPVisionModelWithProjection - text_model = CLIPTextModelWithProjection.from_pretrained(model_id) - vision_model = CLIPVisionModelWithProjection.from_pretrained(model_id) + 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={ @@ -433,8 +437,8 @@ def main(): from .extra.siglip import SiglipTextModelOnnxConfig, SiglipVisionModelOnnxConfig from transformers.models.siglip import SiglipTextModel, SiglipVisionModel - text_model = SiglipTextModel.from_pretrained(model_id) - vision_model = SiglipVisionModel.from_pretrained(model_id) + 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={ @@ -450,8 +454,8 @@ def main(): # from .extra.clap import ClapTextModelWithProjectionOnnxConfig, ClapAudioModelWithProjectionOnnxConfig # from transformers.models.clap import ClapTextModelWithProjection, ClapAudioModelWithProjection - # text_model = ClapTextModelWithProjection.from_pretrained(model_id) - # audio_model = ClapAudioModelWithProjection.from_pretrained(model_id) + # 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={ @@ -496,7 +500,7 @@ def main(): from transformers import GenerationConfig from .extra.whisper import get_alignment_heads - generation_config = GenerationConfig.from_pretrained(model_id) + 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) From 2fa6218459a2d1af2e6b7d29f37001c00d751e37 Mon Sep 17 00:00:00 2001 From: Joshua Lochner Date: Thu, 15 Feb 2024 19:53:42 +0200 Subject: [PATCH 4/4] Add support for `nomic_bert` Prevents unsupported model warnings --- src/models.js | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/models.js b/src/models.js index 789f853fe..ec1e73f55 100644 --- a/src/models.js +++ b/src/models.js @@ -1471,6 +1471,12 @@ export class BertForQuestionAnswering extends BertPreTrainedModel { } ////////////////////////////////////////////////// +////////////////////////////////////////////////// +// NomicBert models +export class NomicBertPreTrainedModel extends PreTrainedModel { } +export class NomicBertModel extends NomicBertPreTrainedModel { } +////////////////////////////////////////////////// + ////////////////////////////////////////////////// // RoFormer models export class RoFormerPreTrainedModel extends PreTrainedModel { } @@ -5177,6 +5183,7 @@ export class PretrainedMixin { const MODEL_MAPPING_NAMES_ENCODER_ONLY = new Map([ ['bert', ['BertModel', BertModel]], + ['nomic_bert', ['NomicBertModel', NomicBertModel]], ['roformer', ['RoFormerModel', RoFormerModel]], ['electra', ['ElectraModel', ElectraModel]], ['esm', ['EsmModel', EsmModel]],