bge-small-en 是基于 BERT 的句子嵌入模型，专为语义相似度任务优化。该模型将文本句子映射到 384 维稠密向量空间，可用于语义搜索、文本匹配、信息检索等任务。
bge-small-en-ascend/
├── inference.py # 推理测试脚本
├── log.txt # 测试日志
├── README.md # 本文档
├── test_sentences.txt # 测试句子
├── inference_result.json # 推理结果
└── precision_result.json # 精度测试结果docker exec -it test-modelagent bashsource /usr/local/Ascend/ascend-toolkit/set_env.sh模型文件位于 /data/ysws/agentsp/5-16/bge-small-en/BAAI/bge-small-en/ 目录下：
pip install transformers torch_npu -i https://pypi.huaweicloud.com/repository/pypi/simple/Run the inference script to extract sentence embeddings:
cd /data/ysws/agentsp/5-16/bge-small-en-ascend/
python3 inference.py
python3 inference.py --mode inference运行精度对比测试：
cd /data/ysws/agentsp/5-16/bge-small-en-ascend/
python3 inference.py --mode precision_test| 参数 | 说明 | 默认值 |
|---|---|---|
--mode | 测试模式: all, inference 或 precision_test | all |
| 指标 | 实测值 | 阈值 | 状态 |
|---|---|---|---|
| 最大相对误差 | 0.0936% | < 1.00% | PASS |
| CPU 推理时间 | 0.106s | - | - |
| NPU 推理时间 | 0.016s | - | - |
| 加速比 | 6.80x | > 1x | PASS |
| Cosine 相似度差 | 0.0001 | < 0.01 | PASS |
输入句子:
输出:
bge-small-en NPU Test
Model: BAAI/bge-small-en (sentence embeddings)
Output: /data/ysws/agentsp/5-16/bge-small-en-ascend
============================================================
Inference Test (NPU)
============================================================
Device: npu:0
Loading model and tokenizer...
Model loaded successfully
Input sentences: 2
Input shape: torch.Size([2, 9])
Inference time: 0.241s
Embeddings shape: torch.Size([2, 384])
Cosine similarity: 0.8975
============================================================
Precision Test (CPU vs NPU)
============================================================
NPU Device: npu:0
Loading model...
Input sentences: 2
Running on CPU...
Running on NPU...
CPU inference time: 0.106s
NPU inference time: 0.016s
Speedup: 6.80x
Max absolute error: 2.577543e-03
Max relative error: 0.0936% (threshold: 1.0%)
Cosine similarity (CPU): 0.8976
Cosine similarity (NPU): 0.8975
Status: PASS
============================================================
Precision Test Result: PASS
============================================================
============================================================
Test Complete!
============================================================import torch
from transformers import AutoTokenizer, AutoModel
MODEL_DIR = "/data/ysws/agentsp/5-16/bge-small-en/BAAI/bge-small-en"
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModel.from_pretrained(MODEL_DIR)
model = model.to("npu:0").eval()
sentences = ["Query sentence", "Passage to match"]
inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
inputs = {k: v.to("npu:0") for k, v in inputs.items()}
def mean_pooling(model_output, attention_mask):
token_embeddings = model_output.last_hidden_state
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
with torch.no_grad():
outputs = model(**inputs)
embeddings = mean_pooling(outputs, inputs['attention_mask'])
print(f"Embeddings shape: {embeddings.shape}")from sklearn.metrics.pairwise import cosine_similarity
emb1 = embeddings[0].cpu().numpy()
emb2 = embeddings[1].cpu().numpy()
similarity = cosine_similarity([emb1], [emb2])[0][0]
print(f"Cosine similarity: {similarity:.4f}")| 组件 | 说明 |
|---|---|
| embeddings | BERT 词嵌入 (vocab_size=30522) |
| encoder | 12 层 Transformer 编码器 |
| 1_Pooling | 对 token 嵌入进行均值池化 |
从 config.json 提取的关键参数:
{
"hidden_size": 384,
"intermediate_size": 1536,
"num_attention_heads": 12,
"num_hidden_layers": 12,
"vocab_size": 30522,
"max_position_embeddings": 512
}A: 检查 NPU 驱动是否正确安装。BGE 模型在 CPU 和 NPU 上的数值误差极小（< 0.1%），远低于 1% 阈值。
A: 使用批处理可以显著提高吞吐量。NPU 相比 CPU 有显著加速（6.8x）。
A: 可用于:
本项目遵循 Apache-2.0 许可证