Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add support for exporting custom models #556

Merged
merged 4 commits into from
Feb 19, 2024
Merged

Add support for exporting custom models #556

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
95600ad
Add support for exporting custom models
xenova Feb 1, 2024
cff27c9
Require task to be specified
xenova Feb 1, 2024
3d1e0cc
Pass `trust_remote_code` to all `from_pretrained` calls.
xenova Feb 2, 2024
2fa6218
Add support for `nomic_bert`
xenova Feb 15, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 98 additions & 55 deletions scripts/convert.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -280,13 +296,30 @@ 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)
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)
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.
Expand All @@ -302,13 +335,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
Expand Down Expand Up @@ -368,63 +408,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, **from_pretrained_kwargs)
vision_model = CLIPVisionModelWithProjection.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)),
},
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, **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`.')

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:
Expand Down Expand Up @@ -457,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)

Expand Down
7 changes: 7 additions & 0 deletions src/models.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 { }
Expand Down Expand Up @@ -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]],
Expand Down
Loading