Spaces:
Running
Running
File size: 2,744 Bytes
55e66ae | 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 66 67 68 69 70 71 72 73 74 75 76 77 | """
SNAP BERT ONNX Export Utility
=============================
Converts a HuggingFace BERT model to ONNX for use with SNAP.
SNAP explicitly requires the 10th layer outputs (hidden_states[-3]) to maintain
compatibility with the original TTS architectures it was trained on.
Standard optimum-cli export will NOT work correctly.
Usage:
pip install transformers torch onnx
python scripts/export_bert.py --model kykim/bert-kor-base --output snap/weights/ko/bert.onnx
"""
import os
import argparse
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForMaskedLM
class BERTFeatureExtractor(nn.Module):
"""Wraps BERT to specifically output the 10th layer hidden state."""
def __init__(self, bert_model):
super().__init__()
self.bert = bert_model
def forward(self, input_ids, attention_mask, token_type_ids):
outputs = self.bert(
input_ids=input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
output_hidden_states=True,
)
return outputs.hidden_states[-3]
def export_bert(model_id, output_path):
print(f"Loading {model_id}...")
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForMaskedLM.from_pretrained(model_id)
wrapper = BERTFeatureExtractor(model).eval()
# Dummy inputs
text = "Hello, world!"
inputs = tokenizer(text, return_tensors="pt")
input_ids = inputs["input_ids"]
attention_mask = inputs["attention_mask"]
token_type_ids = inputs.get("token_type_ids", torch.zeros_like(input_ids))
os.makedirs(os.path.dirname(os.path.abspath(output_path)), exist_ok=True)
print(f"Exporting to {output_path}...")
torch.onnx.export(
wrapper,
(input_ids, attention_mask, token_type_ids),
output_path,
input_names=["input_ids", "attention_mask", "token_type_ids"],
output_names=["hidden_states"],
dynamic_axes={
"input_ids": {0: "batch", 1: "seq_len"},
"attention_mask": {0: "batch", 1: "seq_len"},
"token_type_ids": {0: "batch", 1: "seq_len"},
"hidden_states": {0: "batch", 1: "seq_len"},
},
opset_version=17,
do_constant_folding=True,
)
tok_dir = os.path.dirname(os.path.abspath(output_path))
tokenizer.save_pretrained(tok_dir)
print(f"Done! Tokenizer saved to {tok_dir}")
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--model", required=True, help="HuggingFace model ID")
parser.add_argument("--output", required=True, help="Output ONNX path")
args = parser.parse_args()
export_bert(args.model, args.output)
|