Initial commit

2026-05-12 09:41:56 +08:00
commit 572283e101
936 changed files with 133949 additions and 0 deletions
@@ -0,0 +1,186 @@
+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")