冬
gcw_IDzXRVNw/answerai-colbert-small-v1-ascend
模型介绍文件和版本Pull Requests讨论分析
下载使用量0

answerai-colbert-small-v1 Ascend NPU 部署指南

项目简介

answerai-colbert-small-v1 是 Answer.AI 开发的高效多向量检索模型，基于 JaColBERTv2.5 训练方案，仅用 3300 万参数即可超越大多数同等规模的单向量模型，甚至优于 e5-large-v2 和 bge-base-en-v1.5 等更大规模的流行模型。该模型特别适合 RAG 场景下的语义检索和重排序任务。

特性

  • 支持 Ascend NPU 推理加速
  • ColBERT 多向量检索架构（Late Interaction）
  • CPU vs NPU 精度对比测试（< 1% 误差）
  • 384 维词袋级别嵌入输出
  • 兼容 HuggingFace transformers

环境要求

  • 硬件: 华为 Ascend 910 系列 NPU
  • CANN: 8.0.RC1 或更高版本
  • PyTorch: 2.8.0+ with torch_npu
  • Docker: 容器名称 test-modelagent
  • transformers: 4.33+

目录结构

answerai-colbert-small-v1-ascend/
├── inference.py          # 推理测试脚本
├── log_inference.txt    # 推理测试日志
├── log_precision.txt    # 精度测试日志
├── log.txt               # 测试日志（最新）
├── README.md             # 本文档

部署步骤

1. 进入容器

docker exec -it test-modelagent bash

2. 设置环境变量

source /usr/local/Ascend/ascend-toolkit/set_env.sh

3. 准备模型文件

模型文件位于 /data/ysws/agentsp/5-15/answerai-colbert-small-v1/ 目录下：

  • model.safetensors - 模型权重
  • config.json - 模型配置
  • tokenizer.json / vocab.txt - 分词器文件
  • tokenizer_config.json - 分词器配置

使用方式

方式一：普通推理模式

运行推理脚本进行句子嵌入提取：

cd /data/ysws/agentsp/5-15/answerai-colbert-small-v1-ascend/

python3 inference.py --mode inference --device npu:0

方式二：精度测试模式 (CPU vs NPU)

运行精度对比测试，验证 NPU 计算结果与 CPU 一致性：

cd /data/ysws/agentsp/5-15/answerai-colbert-small-v1-ascend/

python3 inference.py --mode precision_test

命令行参数说明

参数说明默认值
--mode测试模式: inference 或 precision_testinference
--device运行设备npu:0 (自动检测)

测试验证

精度测试结果

指标实测值阈值状态
Hidden states 相对误差0.0460%< 1.00%PASS
Hidden states Cosine 相似度0.999999> 0.99PASS
Mean pooling Cosine 相似度1.000000> 0.99PASS

性能数据

操作耗时
CPU 推理时间 (1 句)0.092s
NPU 推理时间 (1 句)0.212s

推理结果示例

输入句子嵌入维度推理时间
"This is a test sentence..."[1, 384]0.231s
"The model maps sentences..."[1, 384]0.017s
"ColBERT uses late interaction..."[1, 384]0.016s

测试日志

推理模式日志 (log_inference.txt)

2026-05-15 15:16:12,811 - INFO - ============================================================
2026-05-15 15:16:12,811 - INFO - answerai-colbert-small-v1 NPU 推理测试
2026-05-15 15:16:12,812 - INFO - ============================================================
2026-05-15 15:16:12,812 - INFO - Model dir: /data/ysws/agentsp/5-15/answerai-colbert-small-v1
2026-05-15 15:16:12,812 - INFO - Output dir: /data/ysws/agentsp/5-15/answerai-colbert-small-v1-ascend
2026-05-15 15:16:12,812 - INFO - NPU available: True
2026-05-15 15:16:12,812 - INFO - NPU device count: 8
2026-05-15 15:16:14,449 - INFO - NPU 0: Ascend910B3, total_memory=61.0GB
2026-05-15 15:16:14,450 - INFO - NPU 1: Ascend910B3, total_memory=61.0GB
2026-05-15 15:16:14,451 - INFO - ============================================================
2026-05-15 15:16:14,451 - INFO - Inference Test on npu:0
2026-05-15 15:16:14,451 - INFO - ============================================================
2026-05-15 15:16:19,107 - INFO - Device: npu:0
2026-05-15 15:16:19,107 - INFO - Loading tokenizer...
2026-05-15 15:16:20,252 - INFO - Model loaded successfully
2026-05-15 15:16:20,252 - INFO - Processing 3 sentences...
2026-05-15 15:16:20,275 - INFO - Sentence 1: "This is a test sentence for embedding extraction."
2026-05-15 15:16:20,276 - INFO - Input IDs shape: torch.Size([1, 13])
2026-05-15 15:16:20,533 - INFO - Embeddings shape: torch.Size([1, 384])
2026-05-15 15:16:20,533 - INFO - Inference time: 0.231s
2026-05-15 15:16:20,774 - INFO - Sample embedding[0][:5]: [-0.6593183279037476, -0.07383503019809723, 0.11856944859027863, -0.10477980226278305, 0.0947643592953682]
2026-05-15 15:16:20,776 - INFO - Sentence 2: "The model maps sentences to multi-vector representations."
2026-05-15 15:16:20,776 - INFO - Input IDs shape: torch.Size([1, 12])
2026-05-15 15:16:20,793 - INFO - Embeddings shape: torch.Size([1, 384])
2026-05-15 15:16:20,793 - INFO - Inference time: 0.017s
2026-05-15 15:16:20,793 - INFO - Sample embedding[0][:5]: [-0.26567330956459045, -0.15754197537899017, 0.07123415917158127, -0.17516182363033295, -0.017370836809277534]
2026-05-15 15:16:20,794 - INFO - Sentence 3: "ColBERT uses late interaction for efficient retrieval."
2026-05-15 15:16:20,794 - INFO - Input IDs shape: torch.Size([1, 10])
2026-05-15 15:16:20,810 - INFO - Embeddings shape: torch.Size([1, 384])
2026-05-15 15:16:20,811 - INFO - Inference time: 0.016s
2026-05-15 15:16:20,811 - INFO - Sample embedding[0][:5]: [-0.3067658841609955, 0.12002492696046829, -0.03725152090191841, -0.18808825314044952, -0.10717644542455673]
2026-05-15 15:16:20,815 - INFO - ============================================================
2026-05-15 15:16:20,815 - INFO - INFERENCE RESULT
2026-05-15 15:16:20,815 - INFO - ============================================================
2026-05-15 15:16:20,815 - INFO - Output shape: torch.Size([1, 10, 384])
2026-05-15 15:16:20,815 - INFO - Inference time: 0.016s
2026-05-15 15:16:20,815 - INFO - ============================================================
2026-05-15 15:16:20,815 - INFO - Test Complete!
2026-05-15 15:16:20,815 - INFO - ============================================================

精度测试模式日志 (log_precision.txt)

2026-05-15 15:16:47,141 - INFO - ============================================================
2026-05-15 15:16:47,141 - INFO - answerai-colbert-small-v1 NPU 推理测试
2026-05-15 15:16:47,141 - INFO - ============================================================
2026-05-15 15:16:47,141 - INFO - Model dir: /data/ysws/agentsp/5-15/answerai-colbert-small-v1
2026-05-15 15:16:47,142 - INFO - Output dir: /data/ysws/agentsp/5-15/answerai-colbert-small-v1-ascend
2026-05-15 15:16:47,142 - INFO - NPU available: True
2026-05-15 15:16:47,142 - INFO - NPU device count: 8
2026-05-15 15:16:48,758 - INFO - NPU 0: Ascend910B3, total_memory=61.0GB
2026-05-15 15:16:48,759 - INFO - NPU 1: Ascend910B3, total_memory=61.0GB
2026-05-15 15:16:48,759 - INFO - ============================================================
2026-05-15 15:16:48,760 - INFO - Precision Test: CPU vs NPU (threshold: 1.0%)
2026-05-15 15:16:48,760 - INFO - ============================================================
2026-05-15 15:16:53,461 - INFO - Loading tokenizer...
2026-05-15 15:16:54,279 - INFO - Loading model for CPU...
2026-05-15 15:16:54,692 - INFO - Loading model for NPU...
2026-05-15 15:16:54,692 - INFO - Input text: "This is a test sentence for embedding extraction."
2026-05-15 15:16:54,692 - INFO - Input IDs shape: torch.Size([1, 13])
2026-05-15 15:16:54,692 - INFO - Running inference on CPU...
2026-05-15 15:16:54,785 - INFO - Running inference on NPU...
2026-05-15 15:16:55,305 - INFO - Hidden states CPU shape: (1, 13, 384)
2026-05-15 15:16:55,305 - INFO - Hidden states NPU shape: (1, 13, 384)
2026-05-15 15:16:55,305 - INFO - Embeddings CPU shape: (1, 384)
2026-05-15 15:16:55,305 - INFO - Embeddings NPU shape: (1, 384)
2026-05-15 15:16:55,306 - INFO - CPU inference time: 0.092s
2026-05-15 15:16:55,306 - INFO - NPU inference time: 0.212s
2026-05-15 15:16:55,306 - INFO - === Hidden States Precision ===
2026-05-15 15:16:55,306 - INFO - Max relative error: 4.599434e-04 (0.0460%)
2026-05-15 15:16:55,307 - INFO - Mean relative error: 1.223363e-03 (0.1223%)
2026-05-15 15:16:55,307 - INFO - Cosine similarity: 0.999999 (0.0001% angular error)
2026-05-15 15:16:55,307 - INFO - === Mean Pooling Embeddings Precision ===
2026-05-15 15:16:55,307 - INFO - Cosine similarity: 1.000000
2026-05-15 15:16:55,307 - INFO - PASS: True (threshold: 1.0%)
2026-05-15 15:16:55,318 - INFO - ============================================================
2026-05-15 15:16:55,318 - INFO - PRECISION TEST RESULT
2026-05-15 15:16:55,318 - INFO - ============================================================
2026-05-15 15:16:55,318 - INFO - Relative error: 4.599434e-04
2026-05-15 15:16:55,318 - INFO - CPU time: 0.092s
2026-05-15 15:16:55,318 - INFO - NPU time: 0.212s
2026-05-15 15:16:55,318 - INFO - PASS: True
2026-05-15 15:16:55,318 - INFO - ============================================================
2026-05-15 15:16:55,318 - INFO - Test Complete!
2026-05-15 15:16:55,318 - INFO - ============================================================

Python API 使用示例

基本推理

import torch
from transformers import AutoTokenizer, AutoModel

MODEL_DIR = "/data/ysws/agentsp/5-15/answerai-colbert-small-v1"

tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModel.from_pretrained(MODEL_DIR)
model = model.to("npu:0")
model.eval()

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") if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}

with torch.no_grad():
    outputs = model(**inputs)

embeddings = outputs.last_hidden_state
print(f"Embeddings shape: {embeddings.shape}")  # [batch, seq_len, 384]

均值池化嵌入

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)

embeddings = mean_pooling(outputs, inputs["attention_mask"])
print(f"Pooled embeddings shape: {embeddings.shape}")  # [batch, 384]

语义相似度计算

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}")

模型结构

  • 架构类型: HF_ColBERT (基于 BERT 的多向量检索)
  • 编码器: 12 层 Transformer
  • 隐藏层维度: 384
  • 注意力头数: 12
  • 参数量: ~33M
  • 输出: 每 token 384 维向量
组件说明
embeddingsBERT 词嵌入 (vocab_size=30522)
encoder12 层 Transformer 编码器
pooler对 token 嵌入进行 Mean Pooling 操作

推理参数配置

从 config.json 提取的关键参数:

{
  "hidden_size": 384,
  "intermediate_size": 1536,
  "num_attention_heads": 12,
  "num_hidden_layers": 12,
  "vocab_size": 30522,
  "max_position_embeddings": 512
}

常见问题

Q: ColBERT 与普通 BERT 嵌入有什么区别？

A: ColBERT 使用"Late Interaction"机制，每个 token 都产生一个向量表示，适合高效的向量检索场景。普通 BERT 通常使用 [CLS] token 或 mean pooling 作为句子级别嵌入。

Q: 如何提高检索速度？

A: 使用批处理可以提高吞吐量。ColBERT 的向量可以预先计算并存储，检索时只需计算查询向量然后进行快速的最大内积搜索（MIPS）。

Q: linear.weight UNEXPECTED 警告是什么意思？

A: 这是正常的加载报告，表示模型权重中有一个 linear.weight 不在标准 BERT 模型中。这是 ColBERT 模型特有的权重，不影响推理功能。

参考链接

  • 原始模型: https://huggingface.co/answerdotai/answerai-colbert-small-v1
  • Answer.AI 博客: https://www.answer.ai/posts/2024-08-13-small-but-mighty-colbert.html
  • JaColBERTv2.5 论文: https://arxiv.org/abs/2407.20750
  • HuggingFace Transformers: https://huggingface.co/transformers

许可证

本项目遵循 Apache-2.0 许可证