Spaces:
Runtime error
Runtime error
| import gradio as gr | |
| import json | |
| from pathlib import Path | |
| from huggingface_hub import hf_hub_download, HfApi | |
| from coremltools import ComputeUnit | |
| from transformers.onnx.utils import get_preprocessor | |
| from exporters.coreml import export | |
| from exporters.coreml.features import FeaturesManager | |
| from exporters.coreml.validate import validate_model_outputs | |
| compute_units_mapping = { | |
| "All": ComputeUnit.ALL, | |
| "CPU": ComputeUnit.CPU_ONLY, | |
| "CPU + GPU": ComputeUnit.CPU_AND_GPU, | |
| "CPU + NE": ComputeUnit.CPU_AND_NE, | |
| } | |
| compute_units_labels = list(compute_units_mapping.keys()) | |
| framework_mapping = { | |
| "PyTorch": "pt", | |
| "TensorFlow": "tf", | |
| } | |
| framework_labels = list(framework_mapping.keys()) | |
| precision_mapping = { | |
| "Float32": "float32", | |
| "Float16 quantization": "float16", | |
| } | |
| precision_labels = list(precision_mapping.keys()) | |
| tolerance_mapping = { | |
| "Model default": None, | |
| "1e-2": 1e-2, | |
| "1e-3": 1e-3, | |
| "1e-4": 1e-4, | |
| } | |
| tolerance_labels = list(tolerance_mapping.keys()) | |
| def error_str(error, title="Error"): | |
| return f"""#### {title} | |
| {error}""" if error else "" | |
| def url_to_model_id(model_id_str): | |
| if not model_id_str.startswith("https://huggingface.co/"): return model_id_str | |
| return model_id_str.split("/")[-2] + "/" + model_id_str.split("/")[-1] | |
| def supported_frameworks(model_id): | |
| """ | |
| Return a list of supported frameworks (`PyTorch` or `TensorFlow`) for a given model_id. | |
| Only PyTorch and Tensorflow are supported. | |
| """ | |
| api = HfApi() | |
| model_info = api.model_info(model_id) | |
| tags = model_info.tags | |
| frameworks = [tag for tag in tags if tag in ["pytorch", "tf"]] | |
| return sorted(["PyTorch" if f == "pytorch" else "TensorFlow" for f in frameworks]) | |
| def on_model_change(model): | |
| model = url_to_model_id(model) | |
| tasks = None | |
| error = None | |
| try: | |
| config_file = hf_hub_download(model, filename="config.json") | |
| if config_file is None: | |
| raise Exception(f"Model {model} not found") | |
| with open(config_file, "r") as f: | |
| config_json = f.read() | |
| config = json.loads(config_json) | |
| model_type = config["model_type"] | |
| features = FeaturesManager.get_supported_features_for_model_type(model_type) | |
| tasks = list(features.keys()) | |
| frameworks = supported_frameworks(model) | |
| selected_framework = frameworks[0] if len(frameworks) > 0 else None | |
| return ( | |
| gr.update(visible=bool(model_type)), # Settings column | |
| gr.update(choices=tasks, value=tasks[0] if tasks else None), # Tasks | |
| gr.update(visible=len(frameworks)>1, choices=frameworks, value=selected_framework), # Frameworks | |
| gr.update(value=error_str(error)), # Error | |
| ) | |
| except Exception as e: | |
| error = e | |
| model_type = None | |
| def convert_model(preprocessor, model, model_coreml_config, compute_units, precision, tolerance, output, use_past=False, seq2seq=None): | |
| coreml_config = model_coreml_config(model.config, use_past=use_past, seq2seq=seq2seq) | |
| mlmodel = export( | |
| preprocessor, | |
| model, | |
| coreml_config, | |
| quantize=precision, | |
| compute_units=compute_units, | |
| ) | |
| filename = output | |
| if seq2seq == "encoder": | |
| filename = filename.parent / ("encoder_" + filename.name) | |
| elif seq2seq == "decoder": | |
| filename = filename.parent / ("decoder_" + filename.name) | |
| filename = filename.as_posix() | |
| mlmodel.save(filename) | |
| if tolerance is None: | |
| tolerance = coreml_config.atol_for_validation | |
| validate_model_outputs(coreml_config, preprocessor, model, mlmodel, tolerance) | |
| def convert(model, task, compute_units, precision, tolerance, framework): | |
| model = url_to_model_id(model) | |
| compute_units = compute_units_mapping[compute_units] | |
| precision = precision_mapping[precision] | |
| tolerance = tolerance_mapping[tolerance] | |
| framework = framework_mapping[framework] | |
| # TODO: support legacy format | |
| output = Path("exported")/model/"coreml"/task | |
| output.mkdir(parents=True, exist_ok=True) | |
| output = output/f"{precision}_model.mlpackage" | |
| try: | |
| preprocessor = get_preprocessor(model) | |
| model = FeaturesManager.get_model_from_feature(task, model, framework=framework) | |
| _, model_coreml_config = FeaturesManager.check_supported_model_or_raise(model, feature=task) | |
| if task in ["seq2seq-lm", "speech-seq2seq"]: | |
| # Convert encoder / decoder | |
| convert_model( | |
| preprocessor, | |
| model, | |
| model_coreml_config, | |
| compute_units, | |
| precision, | |
| tolerance, | |
| output, | |
| seq2seq="encoder" | |
| ) | |
| convert_model( | |
| preprocessor, | |
| model, | |
| model_coreml_config, | |
| compute_units, | |
| precision, | |
| tolerance, | |
| output, | |
| seq2seq="decoder" | |
| ) | |
| else: | |
| convert_model( | |
| preprocessor, | |
| model, | |
| model_coreml_config, | |
| compute_units, | |
| precision, | |
| tolerance, | |
| output, | |
| ) | |
| # TODO: push to hub, whatever | |
| return "Done" | |
| except Exception as e: | |
| return error_str(e) | |
| DESCRIPTION = """ | |
| ## Convert a transformers model to Core ML | |
| With this Space you can try to convert a transformers model to Core ML. It uses the 🤗 Hugging Face [Exporters repo](https://huggingface.co/exporters) under the hood. | |
| Note that not all models are supported. If you get an error on a model you'd like to convert, please open an issue on the [repo](https://github.com/huggingface/exporters). | |
| After conversion, you can choose to submit a PR to the original repo, or create your own repo with just the converted Core ML weights. | |
| """ | |
| with gr.Blocks() as demo: | |
| gr.Markdown(DESCRIPTION) | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| gr.Markdown("## 1. Load model info") | |
| input_model = gr.Textbox( | |
| max_lines=1, | |
| label="Model name or URL, such as apple/mobilevit-small", | |
| placeholder="distilbert-base-uncased", | |
| value="distilbert-base-uncased", | |
| ) | |
| btn_get_tasks = gr.Button("Load") | |
| with gr.Column(scale=3): | |
| with gr.Column(visible=False) as group_settings: | |
| gr.Markdown("## 2. Select Task") | |
| radio_tasks = gr.Radio(label="Choose the task for the converted model.") | |
| gr.Markdown("The `default` task is suitable for feature extraction.") | |
| radio_framework = gr.Radio( | |
| visible=False, | |
| label="Framework", | |
| choices=framework_labels, | |
| value=framework_labels[0], | |
| ) | |
| radio_compute = gr.Radio( | |
| label="Compute Units", | |
| choices=compute_units_labels, | |
| value=compute_units_labels[0], | |
| ) | |
| radio_precision = gr.Radio( | |
| label="Precision", | |
| choices=precision_labels, | |
| value=precision_labels[0], | |
| ) | |
| radio_tolerance = gr.Radio( | |
| label="Absolute Tolerance for Validation", | |
| choices=tolerance_labels, | |
| value=tolerance_labels[0], | |
| ) | |
| btn_convert = gr.Button("Convert") | |
| gr.Markdown("Conversion will take a few minutes.") | |
| error_output = gr.Markdown(label="Output") | |
| btn_get_tasks.click( | |
| fn=on_model_change, | |
| inputs=input_model, | |
| outputs=[group_settings, radio_tasks, radio_framework, error_output], | |
| queue=False, | |
| scroll_to_output=True | |
| ) | |
| btn_convert.click( | |
| fn=convert, | |
| inputs=[input_model, radio_tasks, radio_compute, radio_precision, radio_tolerance, radio_framework], | |
| outputs=error_output, | |
| scroll_to_output=True | |
| ) | |
| # gr.HTML(""" | |
| # <div style="border-top: 1px solid #303030;"> | |
| # <br> | |
| # <p>Footer</p><br> | |
| # <p><img src="https://visitor-badge.glitch.me/badge?page_id=pcuenq.transformers-to-coreml" alt="visitors"></p> | |
| # </div> | |
| # """) | |
| demo.queue(concurrency_count=1, max_size=10) | |
| demo.launch(debug=True, share=False) | |