Как я могу экспортировать модель Opus-MT-FR-EN Pytorch в один файл ONNX? ⇐ Python

[Anonymous]

Я преобразовал модель Pytorch Helsinki-NLP/Opus-MT-FR-EN (HuggingFace) в ONNX. Я использую этот скрипт: < /p>

Код: Выделить всё

import os
from optimum.onnxruntime import ORTModelForSeq2SeqLM
from transformers import AutoTokenizer, AutoConfig

hf_model_id = "Helsinki-NLP/opus-mt-fr-en"
onnx_save_directory = "./onnx_model_fr_en"

os.makedirs(onnx_save_directory, exist_ok=True)

print(f"Starting conversion for model: {hf_model_id}")
print(f"ONNX model will be saved to: {onnx_save_directory}")

print("Loading tokenizer and config...")
tokenizer = AutoTokenizer.from_pretrained(hf_model_id)
config = AutoConfig.from_pretrained(hf_model_id)

model = ORTModelForSeq2SeqLM.from_pretrained(
hf_model_id,
export=True,
from_transformers=True,
# Pass the loaded config explicitly during export
config=config
)

print("Saving ONNX model components, tokenizer and configuration...")
model.save_pretrained(onnx_save_directory)
tokenizer.save_pretrained(onnx_save_directory)

print("-" * 30)
print(f"Successfully converted '{hf_model_id}' to ONNX.")
print(f"Files saved in: {onnx_save_directory}")
if os.path.exists(onnx_save_directory):
print("Generated files:", os.listdir(onnx_save_directory))
else:
print("Warning: Save directory not found after saving.")
print("-" * 30)

print("Loading ONNX model and tokenizer for testing...")
onnx_tokenizer = AutoTokenizer.from_pretrained(onnx_save_directory)

onnx_model = ORTModelForSeq2SeqLM.from_pretrained(onnx_save_directory)

french_text= "je regarde la tele"
print(f"Input (French): {french_text}")
inputs = onnx_tokenizer(french_text, return_tensors="pt") # Use PyTorch tensors

print("Generating translation using the ONNX model...")
generated_ids = onnx_model.generate(**inputs)
english_translation = onnx_tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

print(f"Output (English): {english_translation}")
print("--- Test complete ---")
< /code>
Выходная папка, содержащая файлы ONNX: < /p>
franck@server:~/tests/onnx_model_fr_en$ ls -la
total 860968
drwxr-xr-x 2 franck users      4096 Apr 16 17:29 .
drwxr-xr-x 5 franck users      4096 Apr 17 23:54 ..
-rw-r--r-- 1 franck users      1360 Apr 17 04:38 config.json
-rw-r--r-- 1 franck users 346250804 Apr 17 04:38 decoder_model.onnx
-rw-r--r-- 1 franck users 333594274 Apr 17 04:38 decoder_with_past_model.onnx
-rw-r--r-- 1 franck users 198711098 Apr 17 04:38 encoder_model.onnx
-rw-r--r-- 1 franck users       288 Apr 17 04:38 generation_config.json
-rw-r--r-- 1 franck users    802397 Apr 17 04:38 source.spm
-rw-r--r-- 1 franck users        74 Apr 17 04:38 special_tokens_map.json
-rw-r--r-- 1 franck users    778395 Apr 17 04:38 target.spm
-rw-r--r-- 1 franck users       847 Apr 17 04:38 tokenizer_config.json
-rw-r--r-- 1 franck users   1458196 Apr 17 04:38 vocab.json
< /code>
Как я могу экспортировать модель Pytorch Opus-mt-fr-en pytorch в один файл ONNX? encoder_model.onnx 

и decoder_model.onnx , которая является проблемой, поскольку встраивающий слой большой (представляет ~ 40% от размера модели Pytorch).


Подробнее здесь: https://stackoverflow.com/questions/795 ... -onnx-file