Spaces:
Running
Running
File size: 2,347 Bytes
9bf1d31 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import os
from transformers import AutoTokenizer, WhisperProcessor, AutoFeatureExtractor
from optimum.onnxruntime import (
ORTModelForQuestionAnswering,
ORTModelForCausalLM,
ORTModelForSequenceClassification,
ORTModelForTokenClassification,
ORTModelForSpeechSeq2Seq,
ORTOptimizer,
ORTModelForMaskedLM,
ORTModelForSeq2SeqLM,
ORTModelForMultipleChoice,
ORTModelForImageClassification,
)
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
TASK_MAPPING = {
# NLP models
"ner": (ORTModelForTokenClassification, AutoTokenizer),
"text_classification": (ORTModelForSequenceClassification, AutoTokenizer),
"question_answering": (ORTModelForQuestionAnswering, AutoTokenizer),
"causal_lm": (ORTModelForCausalLM, AutoTokenizer),
"mask_lm": (ORTModelForMaskedLM, AutoTokenizer),
"seq2seq_lm": (ORTModelForSeq2SeqLM, AutoTokenizer),
"multiple_choice": (ORTModelForMultipleChoice, AutoTokenizer),
# Audio models
"whisper_finetuning": (ORTModelForSpeechSeq2Seq, WhisperProcessor),
# Vision models
"image_classification": (ORTModelForImageClassification, AutoFeatureExtractor),
}
def convert_to_onnx(model_name, task, output_dir):
"""
Convert model to ONNX format for the specified task.
"""
logger.info(f"Converting model: {model_name} for task: {task}")
os.makedirs(output_dir, exist_ok=True)
if task not in TASK_MAPPING:
logger.error(f"Task {task} is not supported for ONNX conversion in this script.")
return None
ORTModelClass, ProcessorClass = TASK_MAPPING[task]
try:
if task == "embedding":
ort_optimizer = ORTOptimizer.from_pretrained(model_name)
ort_optimizer.export(output_dir=output_dir, task="feature-extraction")
else:
ort_model = ORTModelClass.from_pretrained(model_name, export=True)
ort_model.save_pretrained(output_dir)
processor = ProcessorClass.from_pretrained(model_name)
processor.save_pretrained(output_dir)
logger.info(f"Conversion complete. Model saved to: {output_dir}")
return output_dir
except Exception as e:
logger.error(f"Conversion failed: {str(e)}")
return None
|