nomic-bert-2048 是 Nomic AI 开发的高性能 BERT 风格嵌入模型，支持 2048 token 的长文本输入。该模型采用 12 层 Transformer 编码器，隐藏层维度为 768，在文本嵌入和语义理解任务上表现优异。
nomic-bert-2048-ascend/
├── inference.py # 推理测试脚本
├── log.txt # 测试日志
├── README.md # 本文档
└── test_sentences.txt # 测试句子source /usr/local/Ascend/ascend-toolkit/set_env.sh模型文件位于 /data/ysws/agentsp/5-15/nomic-bert-2048/ 目录下：
pip install transformers torch_npu safetensors运行推理脚本进行句子嵌入提取：
cd /data/ysws/agentsp/5-15/nomic-bert-2048-ascend/
# 使用默认测试句子
python3 inference.py
# 使用指定设备
python3 inference.py --device npu:0运行精度对比测试，验证 NPU 计算结果与 CPU 一致性：
cd /data/ysws/agentsp/5-15/nomic-bert-2048-ascend/
# 运行完整精度测试
python3 inference.py --mode precision_test| 参数 | 说明 | 默认值 |
|---|---|---|
--mode | 测试模式: inference 或 precision_test | inference |
--device | 运行设备 | npu:0 (自动检测) |
| 指标 | 实测值 | 阈值 | 状态 |
|---|---|---|---|
| 相对误差 | 0.0001% | < 1.00% | PASS |
| Cosine 相似度差 | 4.77e-07 | < 1.00e-02 | PASS |
| 操作 | 耗时 |
|---|---|
| CPU 推理时间 (3 句) | 0.782s |
| NPU 推理时间 (3 句) | 0.361s |
| 单句 NPU 推理 | ~0.12s |
| 输入句子 | 嵌入维度 | Cosine 相似度 |
|---|---|---|
| "This is a test sentence..." vs "The model maps sentences..." | 768 | 0.8686 |
结果: CPU 和 NPU 输出的 cosine 相似度差异仅 4.77e-07，几乎完全一致
完整测试日志保存在 log.txt
import torch
from transformers import AutoTokenizer, AutoModel
MODEL_DIR = "/data/ysws/agentsp/5-15/nomic-bert-2048"
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModel.from_pretrained(MODEL_DIR, trust_remote_code=True)
model = model.to("npu:0")
model.eval()
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)
sentences = ["This is a test sentence", "Each sentence is converted"]
inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
inputs = {k: v.to("npu:0") for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
embeddings = mean_pooling(outputs, inputs['attention_mask'])
print(f"Embeddings shape: {embeddings.shape}") # torch.Size([2, 768])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}")sentences = [
"First sentence for embedding",
"Second sentence for comparison",
"Third sentence in the batch"
]
inputs = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
inputs = {k: v.to("npu:0") for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
embeddings = mean_pooling(outputs, inputs['attention_mask'])
print(f"Batch embeddings shape: {embeddings.shape}") # torch.Size([3, 768])| 组件 | 说明 |
|---|---|
| embeddings | BERT 词嵌入 (vocab_size=30528) |
| encoder | 12 层 Transformer 编码器 |
| LayerNorm | 层级归一化 |
从 config.json 提取的关键参数:
{
"hidden_size": 768,
"intermediate_size": 3072,
"num_attention_heads": 12,
"num_hidden_layers": 12,
"n_positions": 2048,
"vocab_size": 30528,
"activation_function": "swiglu"
}该模型的 safetensors 文件中的键使用 bert. 前缀（如 bert.embeddings.word_embeddings.weight），而 transformers 模型期望的键不带此前缀。推理脚本中包含自动重映射逻辑来处理此问题：
def remap_state_dict(state_dict):
new_state_dict = {}
for key, value in state_dict.items():
if key.startswith('bert.'):
new_key = key[5:]
else:
new_key = key
new_state_dict[new_key] = value
return new_state_dictA: 检查 NPU 驱动是否正确安装，确保 CANN 环境变量已 source。0.1-0.2% 的数值误差是正常的，因为 NPU 和 CPU 使用不同的计算精度。
A: 使用批处理可以显著提高吞吐量。另外，首次推理会有编译开销，后续推理会更快。
A: 可用于:
A: 正常。警告显示 safetensors 中的 bert. 前缀键无法完全匹配模型期望的键名，但通过重映射后权重可以正确加载。剩余的 cls.predictions.* 键是预训练头部的权重，不影响基础嵌入功能。
本项目遵循 Apache-2.0 许可证