Tungsten123/denseretrieval-chinese-lert-small

项目简介

本项目基于hfl/chinese-lert-small预训练模型开发了一个用于中文语义检索的密集向量检索模型。与传统的语义文本相似度(STS)模型不同，本模型在标准BERT架构基础上引入了自适应权重池化机制，通过多头自注意力网络和MLP权重生成器来动态计算每个token的重要性权重，从而生成更具表达力的句子向量表示。

相比于简单的[CLS]标记池化或平均池化方法(GAP)，本模型的核心创新在于：

1. 自适应权重池化：利用隐藏状态，通过多头自注意力机制计算上下文感知的token权重
2. 动态重要性建模：通过MLP网络将注意力输出映射为标量权重，经过softmax归一化后进行加权求和
3. 端到端优化：整个权重生成和池化过程完全可微，能够与检索任务损失函数联合优化

训练配置

主要训练阶段超参数

training:
  backbone_lr: "5e-6"
  new_modules_lr: "3e-5"   
  backbone_weight_decay: "1.2e-2"
  new_modules_weight_decay: "2.5e-2" 
  max_length: 512
  batch_size: 20
  num_epochs: 18
  gradient_accumulation_steps: 8
  max_grad_norm: 0.5
scheduler:
  type: "cosine"
  warmup_ratio: 0.1

量化训练阶段超参数

training:
  batch_size: 20
  max_length: 256
  num_steps: 800
  gradient_accumulation_steps: 8
  max_grad_norm: 1.0
  learning_rate: 5e-6
scheduler:
  type: "cosine"
  warmup_ratio: 0.1

使用方法

基础模型推理

from transformers import AutoTokenizer, AutoModel
import torch
import torch.nn.functional as F

model = AutoModel.from_pretrained("Tungsten123/denseretrieval-chinese-lert-small", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("Tungsten123/denseretrieval-chinese-lert-small")

sentences = [
    "机器学习中的深度学习算法有哪些优势？",
    "深度学习是机器学习的一个重要分支，它通过构建多层神经网络来模拟人脑的学习过程。相比传统的机器学习方法，深度学习具有强大的特征提取能力，能够自动学习数据中的复杂模式和抽象表示。在图像识别、自然语言处理、语音识别等领域，深度学习算法展现出了卓越的性能，特别是在处理大规模数据时，其优势更加明显。深度学习模型可以通过反向传播算法进行端到端的训练，无需人工设计特征，大大减少了特征工程的工作量。"
]

encoded = tokenizer(sentences, padding=True, truncation=True, max_length=512, return_tensors="pt")
with torch.no_grad():
    outputs = model(**encoded)
    embeddings = F.normalize(outputs.pooled_output)
    similarity = F.cosine_similarity(embeddings[0:1], embeddings[1:2], dim=1)
    print(f"相似度: {similarity.item():.4f}")

INT8量化版本

本模型提供基于Brevitas框架的INT8量化版本，通过量化感知训练(QAT)方法。量化版本在保持模型精度的同时显著减少了内存占用和推理延迟，特别适合资源受限的生产环境部署。若需要测试该模型，请克隆本仓库，并进入仓库根目录执行下面的代码：

import numpy as np
import onnxruntime as ort
from transformers import AutoTokenizer

session = ort.InferenceSession("quantization/model.onnx")
tokenizer = AutoTokenizer.from_pretrained("Tungsten123/denseretrieval-chinese-lert-small")

sentences = [
    "人工智能技术在医疗领域的应用前景如何？",
    "人工智能在医疗健康领域展现出巨大的应用潜力和变革性影响。通过深度学习和计算机视觉技术，AI系统能够协助医生进行医学影像诊断，包括X光片、CT扫描、MRI等影像的自动分析和病变识别。在药物研发方面，人工智能可以加速新药发现过程，通过分子建模和虚拟筛选大幅缩短研发周期。此外，AI还能够实现个性化治疗方案推荐，根据患者的基因信息、病史和症状特征制定最优治疗策略。智能诊断系统、医疗机器人、远程医疗监护等应用正在逐步改变传统医疗服务模式。"
]

encoded = tokenizer(sentences, padding=True, truncation=True, max_length=512, return_tensors="np")
inputs = {"input_ids": encoded["input_ids"].astype(np.int64), "attention_mask": encoded["attention_mask"].astype(np.int64)}
outputs = session.run(None, inputs)
embeddings = outputs[0]
embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
similarity = np.dot(embeddings[0], embeddings[1])
print(f"相似度: {similarity:.4f}")

数据局限性声明

由于中文问题-上下文不对称检索任务的标准化评估基准数据集相对缺乏，本模型目前尚未在权威的中文语义检索benchmark上进行全面评估。现有的中文文本相似度数据集主要关注对称的句子对相似度任务，而针对检索场景中查询短文本与长文档片段之间语义匹配的专门数据集较为稀少。

因此，本模型的性能评估主要基于DuReader训练集的内部验证指标，可能无法全面反映其在不同领域和应用场景下的泛化能力。我们建议用户在特定应用场景中进行充分的测试和验证，并根据实际需求考虑进行领域适应性微调。

未来我们将持续关注中文检索评估数据集的发展，并及时更新模型的评估结果和性能指标。