如何导出含Tokenizer的HuggingFace PyTorch模型为ONNX并适配BigQuery ML

我正尝试将HuggingFace模型转换为ONNX格式，以便在BigQuery ML中使用（该平台支持导入ONNX模型）。但transformers的tokenizer无法被包含到模型中。

我该如何将模型与Tokenizer导出为单个ONNX文件？

以下是我已尝试的操作：

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import sys
# model_name is being passed in from the command line
model_name = sys.argv[1]

# Load tokenizer and PyTorch weights form the Hub
tokenizer = AutoTokenizer.from_pretrained(model_name)
pt_model = AutoModelForSequenceClassification.from_pretrained(model_name)

# replace slashes with underscores
model_name_underscored = model_name.replace("/", "_")
# Save to disk
tokenizer.save_pretrained(model_name_underscored)
pt_model.save_pretrained(model_name_underscored)

然后执行：

python3 -m transformers.onnx --model=./$model_name_underscored ${model_name_underscored}_onnx/

将模型导入BigQuery ML后，我发现模型需要分词后的输入，而非明文。

此外，我该如何在BigQuery ML中使用AutoTokenizer，使模型输出与Python脚本的输出一致？

一、将Tokenizer与模型打包为单个ONNX文件 #

ONNX本身只定义模型计算图，无法直接内置tokenizer的预处理逻辑，但可以通过把tokenizer的分词逻辑转化为ONNX计算图的一部分，将「分词+模型推理」的完整流程打包成一个ONNX文件。具体步骤如下：

1. 封装包含分词逻辑的组合模型
   用PyTorch把tokenizer的预处理和原模型封装成一个新模块，让整个流程可被导出为ONNX：

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import sys

class TextClassificationPipeline(torch.nn.Module):
    def __init__(self, model_name):
        super().__init__()
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model.eval()  # 固定为评估模式

    def forward(self, text):
        # 处理单条/批量文本，自动应用padding、truncation
        inputs = self.tokenizer(
            text,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.tokenizer.model_max_length
        )
        return self.model(**inputs)

if __name__ == "__main__":
    model_name = sys.argv[1]
    model_name_underscored = model_name.replace("/", "_")
    
    # 初始化组合模型
    pipeline_model = TextClassificationPipeline(model_name)
    
    # 准备示例输入，用于ONNX导出时追踪计算图
    example_text = "This is a sample input text"
    
    # 导出完整流程的ONNX模型
    torch.onnx.export(
        pipeline_model,
        example_text,
        f"{model_name_underscored}_full_pipeline.onnx",
        opset_version=17,  # 选择BigQuery兼容的ONNX版本，建议用16+
        input_names=["text"],
        output_names=["logits"],
        dynamic_axes={
            "text": {0: "batch_size"},  # 支持批量输入
            "logits": {0: "batch_size"}
        }
    )

2. 验证导出的ONNX模型
   用onnxruntime测试模型是否能直接接收明文并输出正确结果：

import onnxruntime as ort

session = ort.InferenceSession(f"{model_name_underscored}_full_pipeline.onnx")
input_text = ["Test sentence 1", "Test sentence 2"]
outputs = session.run(None, {"text": input_text})
print(outputs[0])

二、在BigQuery ML中对齐Tokenizer输出 #

如果不想打包完整ONNX流程，也可以在BigQuery中复现AutoTokenizer的分词逻辑，确保输入格式与Python端一致：

1. 提取Tokenizer核心参数
   从本地保存的tokenizer文件（tokenizer.json、vocab.txt等）中获取：

• 词汇表（vocabulary）
• 最大序列长度（max_length）
• 特殊Token的ID（如[CLS]=101、[SEP]=102、[PAD]=0）
• 截断/填充策略

2. 在BigQuery中实现分词函数
   以BERT类模型的WordPiece分词为例，先将词汇表导入BigQuery表vocab_table（含token和token_id两列），再创建自定义分词函数：

CREATE OR REPLACE FUNCTION `your_project.your_dataset.hf_tokenize`(text STRING)
RETURNS STRUCT<
    input_ids ARRAY<INT64>,
    attention_mask ARRAY<INT64>,
    token_type_ids ARRAY<INT64>
> AS (
    WITH tokenized AS (
        SELECT
            -- 拆分文本并匹配词汇表
            ARRAY(
                SELECT token_id
                FROM UNNEST(SPLIT(REGEXP_REPLACE(text, r"([^\w\s])", r" \1 "), " ")) word
                LEFT JOIN `your_project.your_dataset.vocab_table` vocab
                ON vocab.token = word
                WHERE token_id IS NOT NULL
            ) AS raw_tokens
    ),
    formatted_tokens AS (
        SELECT
            -- 添加[CLS]、[SEP]并截断到最大长度（预留位置给特殊Token）
            ARRAY_CONCAT(
                [101],
                ARRAY_SLICE(raw_tokens, 0, 510),
                [102]
            ) AS input_ids,
            -- 生成attention_mask
            ARRAY(SELECT 1 FROM UNNEST(ARRAY_CONCAT([101], ARRAY_SLICE(raw_tokens, 0, 510), [102]))) AS attention_mask,
            -- 单句输入的token_type_ids全为0
            ARRAY(SELECT 0 FROM UNNEST(ARRAY_CONCAT([101], ARRAY_SLICE(raw_tokens, 0, 510), [102]))) AS token_type_ids
        FROM tokenized
    ),
    padded_tokens AS (
        SELECT
            -- 填充到模型要求的最大序列长度（以512为例）
            ARRAY_PAD(input_ids, 512, 0) AS input_ids,
            ARRAY_PAD(attention_mask, 512, 0) AS attention_mask,
            ARRAY_PAD(token_type_ids, 512, 0) AS token_type_ids
        FROM formatted_tokens
    )
    SELECT AS STRUCT input_ids, attention_mask, token_type_ids FROM padded_tokens
);

3. 调用模型进行预测
   在BigQuery中用自定义函数处理明文输入，再传入模型：

SELECT
    text,
    ML.PREDICT(
        MODEL `your_project.your_dataset.your_onnx_model`,
        (SELECT text, input_ids, attention_mask, token_type_ids FROM UNNEST([hf_tokenize(text)]))
    ) AS prediction
FROM `your_project.your_dataset.text_data`;

内容的提问来源于stack exchange，提问作者stkvtflw