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

# 加载模型
model_path = "baby2008/Dmeta-embedding-zh-onnx"
tokenizer = AutoTokenizer.from_pretrained(model_path)
session = ort.InferenceSession(
    f"{model_path}/model_int8.onnx",
    providers=["CPUExecutionProvider"]
)

# 编码文本
text = "这是一个测试句子。"
inputs = tokenizer(
    text,
    max_length=512,
    padding=True,
    truncation=True,
    return_tensors="np"
)

# 推理
input_ids = inputs["input_ids"].astype(np.int64)
attention_mask = inputs["attention_mask"].astype(np.int64)

result = session.run(None, {
    "input_ids": input_ids,
    "attention_mask": attention_mask
})

embedding = result[0]  # shape: (1, seq_len, 768)

# 获取句子嵌入（平均池化）
sentence_embedding = embedding.mean(axis=1)

# 归一化
normalized = sentence_embedding / np.linalg.norm(sentence_embedding)

相似度计算

python

def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

# 编码两个句子
text1 = "今天天气很好"
text2 = "阳光明媚"

# ... (获取嵌入)

# 计算相似度
similarity = cosine_similarity(embedding1, embedding2)
print(f"Similarity: {similarity:.4f}")

模型输出

输出形状: (batch_size, sequence_length, 768)
输出类型: float32
推荐池化: 平均池化 (mean pooling) 或 CLS token

使用场景

✅ 语义检索
✅ 文本相似度计算
✅ 文本聚类
✅ 推荐系统
✅ 问答系统

与原模型对比

指标	原模型	INT8 量化
模型大小	~400 MB	98.7 MB
内存占用	较高	低
推理速度	基准	相似
精度损失	-	< 1%

文件说明

.
├── config.json              # 模型配置
├── model_int8.onnx          # INT8 量化模型 (98.7 MB)
├── special_tokens_map.json  # 特殊 token 映射
├── tokenizer.json           # 分词器
├── tokenizer_config.json    # 分词器配置
└── vocab.txt                # 词汇表

baby2008/
Dmeta-embedding-zh-onnx

DMeta-Embedding-ZH ONNX

模型特点

模型信息

性能基准

安装依赖

使用方法

Python

相似度计算

模型输出

使用场景

与原模型对比

文件说明

注意事项

许可证

参考资料