使用示例

import onnxruntime as ort
import numpy as np
from transformers import MobileBertTokenizer

# 设置模型和分词器路径
model_path = r'C:\NEW_tinybert\AI\tinybert_model.onnx'  # ONNX 模型路径
tokenizer_path = r'C:\NEW_distilbert\AI'  # 本地分词器路径

# 初始化 ONNX 模型会话
ort_session = ort.InferenceSession(model_path)

# 加载 MobileBERT 分词器
tokenizer = MobileBertTokenizer.from_pretrained(tokenizer_path)

# 文本分类函数
def test_model(text):
    \"\"\"
    使用 ONNX 模型对输入文本进行分类的函数
    Args:
        text (str): 要分析的文本
    Returns:
        str: 预测结果消息
    \"\"\"
    # 将输入文本分词并转换为 ONNX 模型输入格式
    inputs = tokenizer(
        text,
        padding=\"max_length\",  # 将输入长度固定为 128
        truncation=True,       # 截断过长的文本
        max_length=128,        # 最大 token 长度
        return_tensors=\"np\"    # 以 NumPy 数组格式返回
    )
    
    # 将 NumPy 数组转换为 int64
    input_ids = inputs[\"input_ids\"].astype(np.int64)
    attention_mask = inputs[\"attention_mask\"].astype(np.int64)
    
    # 准备 ONNX 模型输入
    ort_inputs = {
        \"input_ids\": input_ids,
        \"attention_mask\": attention_mask
    }
    
    # 执行 ONNX 模型推理
    outputs = ort_session.run(None, ort_inputs)
    logits = outputs[0]  # 模型输出（logit 值）

    # 将 logit 值转换为概率并预测类别
    predicted_class = np.argmax(logits, axis=1).item()

    # 返回结果
    return \"可能是浪漫骗局\" if predicted_class == 1 else \"不是浪漫骗局\"

# 要测试的对话内容
test_texts = [
    \"你没有妈妈吗？\",
    \"我是介绍金融产品的人。投资的话收益很大。\",
    \"我的下面发烫了\",
    \"想摸我的胸部吗？？\"
]

# 输出每个测试文本的结果
for text in test_texts:
    result = test_model(text)
    print(f\"输入: {text} => 结果: {result}\")