194 lines
6.1 KiB
Python
194 lines
6.1 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(
|
||
separators=["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""],
|
||
chunk_size=300,
|
||
chunk_overlap=20
|
||
)
|
||
|
||
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")
|