检索模型bi-encoder笔记

文章目录

- - - 基础示例
    - 使用python和transform从零实现

bi-encoder的核心是独立编码(向量)，然后计算相似度。

基础示例

python 复制代码

from sentence_transformers import SentenceTransformer, util

# 1. 加载预训练的 Bi-Encoder 模型
# 'all-MiniLM-L6-v2' 是一个速度快、效果好的轻量级模型
model = SentenceTransformer('all-MiniLM-L6-v2')

# 2. 定义两个句子
sentence1 = "深度学习模型在自然语言处理中的应用"
sentence2 = "NLP领域如何使用深度神经网络技术"

# 3. 分别将句子编码为向量
# Bi-Encoder 的核心：两个句子是独立编码的
embedding1 = model.encode(sentence1, convert_to_tensor=True)
embedding2 = model.encode(sentence2, convert_to_tensor=True)

# 4. 计算两个向量的余弦相似度
cosine_score = util.cos_sim(embedding1, embedding2)

print(f"句子1: '{sentence1}'")
print(f"句子2: '{sentence2}'")
print(f"语义相似度得分: {cosine_score.item():.4f}")

输出结果：

python 复制代码

句子1: '深度学习模型在自然语言处理中的应用'
句子2: 'NLP领域如何使用深度神经网络技术'
语义相似度得分: 0.3555

解读：

可以看到，过程为将语句1和语句2分别编码，然后计算相似度。

实际业务中，大概是如下方式：

1、语句2对应数据库数据，会提前向量化并保存。

2、语句1对应查询数据，会实时向量化。

3、然后拿查询的向量和数据库向量比对，这个是非常快的。

使用python和transform从零实现

python 复制代码

import torch
from transformers import AutoModel, AutoTokenizer

class SimpleBiEncoder(torch.nn.Module):
    def __init__(self, model_name="bert-base-chinese"):
        super().__init__()
        # 加载一个预训练的 Transformer 模型作为编码器
        self.encoder = AutoModel.from_pretrained(model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)

    def encode(self, text):
        """将文本编码为向量"""
        # 对文本进行分词
        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        with torch.no_grad():
            # 获取模型输出
            outputs = self.encoder(**inputs)
            # 使用平均池化（Mean Pooling）将序列转换为一个句向量
            # 这是获取句子表示的常用方法之一
            return torch.mean(outputs.last_hidden_state, dim=1)

    def forward(self, text_a, text_b):
        """计算两个文本的相似度"""
        # 分别独立编码两个文本
        emb_a = self.encode(text_a)
        emb_b = self.encode(text_b)
        # 计算余弦相似度
        return torch.nn.functional.cosine_similarity(emb_a, emb_b)

# --- 使用示例 ---
if __name__ == "__main__":
    bi_encoder = SimpleBiEncoder()
    text1 = "深度学习模型架构"
    text2 = "神经网络结构设计"
    
    similarity = bi_encoder(text1, text2)
    print(f"文本相似度: {similarity.item():.4f}")