检索模型bi-encoder笔记

基础示例

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

输出结果：

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

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

实际业务中，大概是如下方式：
1、语句2对应数据库数据，会提前向量化并保存。
2、语句1对应查询数据，会实时向量化。
3、然后拿查询的向量和数据库向量比对，这个是非常快的。

使用python和transform从零实现

importtorchfromtransformersimportAutoModel,AutoTokenizerclassSimpleBiEncoder(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)defencode(self,text):"""将文本编码为向量"""# 对文本进行分词inputs=self.tokenizer(text,return_tensors="pt",padding=True,truncation=True)withtorch.no_grad():# 获取模型输出outputs=self.encoder(**inputs)# 使用平均池化（Mean Pooling）将序列转换为一个句向量# 这是获取句子表示的常用方法之一returntorch.mean(outputs.last_hidden_state,dim=1)defforward(self,text_a,text_b):"""计算两个文本的相似度"""# 分别独立编码两个文本emb_a=self.encode(text_a)emb_b=self.encode(text_b)# 计算余弦相似度returntorch.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}")