cde-small-v1 是一个用于文本嵌入的模型，在 MTEB 基准测试中表现优异。该模型支持多种文本分类和嵌入任务，包括 Amazon 评论分类、情感分析等。
cde-small-v1-ascend/
├── inference.py # 推理测试脚本
├── test.log # 测试日志
├── README.md # 本文档cde-small-v1 模型依赖于 nomic-ai/nomic-bert-2048 作为其 embedder。由于网络限制，需要预先下载到本地。
当前实现通过手动加载本地 safetensors 权重并重映射 bert. 前缀的键来正确加载模型。
cd /opt/atomgit/mxy/cde-small-v1-ascend/
python3 inference.py --mode inference --device npu:0python3 inference.py --mode precision_test --device npu:0| 参数 | 说明 | 默认值 |
|---|---|---|
--mode | 测试模式: inference 或 precision_test | inference |
--device | 运行设备: npu:0, cuda:0, cpu, auto (默认auto) | auto |
| 组件 | 说明 |
|---|---|
| first_stage_model | 第一阶段嵌入模型 |
| second_stage_model | 第二阶段 Transformer 模型 |
| output_projection | 输出投影层 |
从 config.json 提取:
{
"architecture": "transductive",
"embedder": "nomic-ai/nomic-bert-2048",
"embedder_rerank": "sentence-transformers/gtr-t5-base",
"max_seq_length": 512,
"logit_scale": 50.0
}nomic-bert-2048 embedder，需要预先下载到本地目录bert. 前缀的键以正确加载 safetensors 格式的权重成功在 NPU 上运行文本嵌入提取，输出 768 维归一化嵌入向量。
测试输出示例：
Input: Hello, this is a test sentence.
Input shape: torch.Size([1, 10])
Inference time: 0.368s
Embedding shape: torch.Size([1, 768])
Embedding norm: 1.0000
Input: This is another example for embedding.
Input shape: torch.Size([1, 11])
Inference time: 0.032s
Embedding shape: torch.Size([1, 768])
Embedding norm: 1.0000
Input: CDE model for text embedding extraction
Input shape: torch.Size([1, 11])
Inference time: 0.026s
Embedding shape: torch.Size([1, 768])
Embedding norm: 1.0000
Inference Summary
Total sentences processed: 3
Total inference time: 0.426s
Average time per sentence: 0.142s完整测试日志保存在 test.log。
本项目遵循 Apache-2.0 许可证