answerai-colbert-small-v1 是 Answer.AI 开发的高效多向量检索模型，基于 JaColBERTv2.5 训练方案，仅用 3300 万参数即可超越大多数同等规模的单向量模型，甚至优于 e5-large-v2 和 bge-base-en-v1.5 等更大规模的流行模型。该模型特别适合 RAG 场景下的语义检索和重排序任务。
answerai-colbert-small-v1-ascend/
├── inference.py # 推理测试脚本
├── test.log # 测试日志
├── README.md # 本文档source /usr/local/Ascend/ascend-toolkit/set_env.sh模型文件位于 /opt/atomgit/mxy/answerai-colbert-small-v1/ 目录下：
运行推理脚本进行句子嵌入提取：
cd /opt/atomgit/mxy/answerai-colbert-small-v1-ascend/
python3 inference.py --mode inference --device npu:0运行精度对比测试，验证 NPU 计算结果与 CPU 一致性：
cd /opt/atomgit/mxy/answerai-colbert-small-v1-ascend/
python3 inference.py --mode precision_test| 参数 | 说明 | 默认值 |
|---|---|---|
--mode | 测试模式: inference 或 precision_test | inference |
--device | 运行设备 | npu:0 (自动检测) |
| 指标 | 实测值 | 阈值 | 状态 |
|---|---|---|---|
| Hidden states 相对误差 | 0.0460% | < 1.00% | PASS |
| Hidden states Cosine 相似度 | 0.999999 | > 0.99 | PASS |
| Mean pooling Cosine 相似度 | 1.000000 | > 0.99 | PASS |
| 操作 | 耗时 |
|---|---|
| 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 |
import torch
from transformers import AutoTokenizer, AutoModel
MODEL_DIR = "/opt/atomgit/mxy/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}")| 组件 | 说明 |
|---|---|
| embeddings | BERT 词嵌入 (vocab_size=30522) |
| encoder | 12 层 Transformer 编码器 |
| pooler | Mean Pooling over token embeddings |
从 config.json 提取的关键参数:
{
"hidden_size": 384,
"intermediate_size": 1536,
"num_attention_heads": 12,
"num_hidden_layers": 12,
"vocab_size": 30522,
"max_position_embeddings": 512
}A: ColBERT 使用"Late Interaction"机制，每个 token 都产生一个向量表示，适合高效的向量检索场景。普通 BERT 通常使用 [CLS] token 或 mean pooling 作为句子级别嵌入。
A: 使用批处理可以提高吞吐量。ColBERT 的向量可以预先计算并存储，检索时只需计算查询向量然后进行快速的最大内积搜索（MIPS）。
A: 这是正常的加载报告，表示模型权重中有一个 linear.weight 不在标准 BERT 模型中。这是 ColBERT 模型特有的权重，不影响推理功能。
本项目遵循 Apache-2.0 许可证