187 lines
5.9 KiB
Python
187 lines
5.9 KiB
Python
import os
|
|
from langchain_community.vectorstores import FAISS
|
|
from langchain.retrievers import ContextualCompressionRetriever
|
|
from langchain.retrievers.document_compressors import LLMChainExtractor
|
|
from langchain_community.embeddings import HuggingFaceBgeEmbeddings
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
from langchain_community.document_loaders import TextLoader
|
|
from langchain_deepseek import ChatDeepSeek
|
|
|
|
# 导入ColBERT重排器需要的模块
|
|
from langchain.retrievers.document_compressors.base import BaseDocumentCompressor
|
|
from langchain.retrievers.document_compressors import DocumentCompressorPipeline
|
|
from langchain_core.documents import Document
|
|
from typing import Sequence
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModel
|
|
import torch.nn.functional as F
|
|
|
|
class ColBERTReranker(BaseDocumentCompressor):
|
|
"""ColBERT重排器"""
|
|
|
|
def __init__(self, **kwargs):
|
|
super().__init__(**kwargs)
|
|
|
|
model_name = "bert-base-uncased"
|
|
|
|
# 加载模型和分词器
|
|
object.__setattr__(self, 'tokenizer', AutoTokenizer.from_pretrained(model_name))
|
|
object.__setattr__(self, 'model', AutoModel.from_pretrained(model_name))
|
|
self.model.eval()
|
|
print(f"ColBERT模型加载完成")
|
|
|
|
def encode_text(self, texts):
|
|
"""ColBERT文本编码"""
|
|
inputs = self.tokenizer(
|
|
texts,
|
|
return_tensors="pt",
|
|
padding=True,
|
|
truncation=True,
|
|
max_length=128
|
|
)
|
|
|
|
with torch.no_grad():
|
|
outputs = self.model(**inputs)
|
|
|
|
embeddings = outputs.last_hidden_state
|
|
embeddings = F.normalize(embeddings, p=2, dim=-1)
|
|
|
|
return embeddings
|
|
|
|
def calculate_colbert_similarity(self, query_emb, doc_embs, query_mask, doc_masks):
|
|
"""ColBERT相似度计算（MaxSim操作）"""
|
|
scores = []
|
|
|
|
for i, doc_emb in enumerate(doc_embs):
|
|
doc_mask = doc_masks[i:i+1]
|
|
|
|
# 计算相似度矩阵
|
|
similarity_matrix = torch.matmul(query_emb, doc_emb.unsqueeze(0).transpose(-2, -1))
|
|
|
|
# 应用文档mask
|
|
doc_mask_expanded = doc_mask.unsqueeze(1)
|
|
similarity_matrix = similarity_matrix.masked_fill(~doc_mask_expanded.bool(), -1e9)
|
|
|
|
# MaxSim操作
|
|
max_sim_per_query_token = similarity_matrix.max(dim=-1)[0]
|
|
|
|
# 应用查询mask
|
|
query_mask_expanded = query_mask.unsqueeze(0)
|
|
max_sim_per_query_token = max_sim_per_query_token.masked_fill(~query_mask_expanded.bool(), 0)
|
|
|
|
# 求和得到最终分数
|
|
colbert_score = max_sim_per_query_token.sum(dim=-1).item()
|
|
scores.append(colbert_score)
|
|
|
|
return scores
|
|
|
|
def compress_documents(
|
|
self,
|
|
documents: Sequence[Document],
|
|
query: str,
|
|
callbacks=None,
|
|
) -> Sequence[Document]:
|
|
"""对文档进行ColBERT重排序"""
|
|
if len(documents) == 0:
|
|
return documents
|
|
|
|
# 编码查询
|
|
query_inputs = self.tokenizer(
|
|
[query],
|
|
return_tensors="pt",
|
|
padding=True,
|
|
truncation=True,
|
|
max_length=128
|
|
)
|
|
|
|
with torch.no_grad():
|
|
query_outputs = self.model(**query_inputs)
|
|
query_embeddings = F.normalize(query_outputs.last_hidden_state, p=2, dim=-1)
|
|
|
|
# 编码文档
|
|
doc_texts = [doc.page_content for doc in documents]
|
|
doc_inputs = self.tokenizer(
|
|
doc_texts,
|
|
return_tensors="pt",
|
|
padding=True,
|
|
truncation=True,
|
|
max_length=128
|
|
)
|
|
|
|
with torch.no_grad():
|
|
doc_outputs = self.model(**doc_inputs)
|
|
doc_embeddings = F.normalize(doc_outputs.last_hidden_state, p=2, dim=-1)
|
|
|
|
# 计算ColBERT相似度
|
|
scores = self.calculate_colbert_similarity(
|
|
query_embeddings,
|
|
doc_embeddings,
|
|
query_inputs['attention_mask'],
|
|
doc_inputs['attention_mask']
|
|
)
|
|
|
|
# 排序并返回前5个
|
|
scored_docs = list(zip(documents, scores))
|
|
scored_docs.sort(key=lambda x: x[1], reverse=True)
|
|
reranked_docs = [doc for doc, _ in scored_docs[:5]]
|
|
|
|
return reranked_docs
|
|
|
|
|
|
|
|
# 初始化配置
|
|
hf_bge_embeddings = HuggingFaceBgeEmbeddings(
|
|
model_name="BAAI/bge-large-zh-v1.5"
|
|
)
|
|
|
|
llm = ChatDeepSeek(
|
|
model="deepseek-chat",
|
|
temperature=0.1,
|
|
api_key=os.getenv("DEEPSEEK_API_KEY")
|
|
)
|
|
|
|
# 1. 加载和处理文档
|
|
loader = TextLoader("../../data/C4/txt/ai.txt", encoding="utf-8")
|
|
documents = loader.load()
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
|
|
docs = text_splitter.split_documents(documents)
|
|
|
|
# 2. 创建向量存储和基础检索器
|
|
vectorstore = FAISS.from_documents(docs, hf_bge_embeddings)
|
|
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 20})
|
|
|
|
# 3. 设置ColBERT重排序器
|
|
reranker = ColBERTReranker()
|
|
|
|
# 4. 设置LLM压缩器
|
|
compressor = LLMChainExtractor.from_llm(llm)
|
|
|
|
# 5. 使用DocumentCompressorPipeline组装压缩管道
|
|
# 流程: ColBERT重排 -> LLM压缩
|
|
pipeline_compressor = DocumentCompressorPipeline(
|
|
transformers=[reranker, compressor]
|
|
)
|
|
|
|
# 6. 创建最终的压缩检索器
|
|
final_retriever = ContextualCompressionRetriever(
|
|
base_compressor=pipeline_compressor,
|
|
base_retriever=base_retriever
|
|
)
|
|
|
|
# 7. 执行查询并展示结果
|
|
query = "AI还有哪些缺陷需要克服？"
|
|
print(f"\n{'='*20} 开始执行查询 {'='*20}")
|
|
print(f"查询: {query}\n")
|
|
|
|
# 7.1 基础检索结果
|
|
print(f"--- (1) 基础检索结果 (Top 20) ---")
|
|
base_results = base_retriever.get_relevant_documents(query)
|
|
for i, doc in enumerate(base_results):
|
|
print(f" [{i+1}] {doc.page_content[:100]}...\n")
|
|
|
|
# 7.2 使用管道压缩器的最终结果
|
|
print(f"\n--- (2) 管道压缩后结果 (ColBERT重排 + LLM压缩) ---")
|
|
final_results = final_retriever.get_relevant_documents(query)
|
|
for i, doc in enumerate(final_results):
|
|
print(f" [{i+1}] {doc.page_content}\n")
|