【新手入门】检索篇 - RAG技术深度实战与优化

import osimport reimport jsonfrom typing import List, Dict, Anyfrom openai import OpenAIimport numpy as npfrom sklearn.metrics.pairwise import cosine_similarity

client = OpenAI(
    api_key=os.environ.get("*******************************"),
    base_url="https://aistudio.baidu.com/llm/lmapi/v3")

class FixedLengthSplitter:
    def __init__(self, chunk_size=500, chunk_overlap=50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap    
    def split_text(self, text: str) -> List[str]:
        chunks = []
        start = 0
        
        while start < len(text):
            end = start + self.chunk_size
            chunk = text[start:end]
            chunks.append(chunk)
            start = end - self.chunk_overlap            
        return chunks

sample_text = """
人工智能（Artificial Intelligence，AI）是计算机科学的一个分支，它企图了解智能的实质，
并生产出一种新的能以人类智能相似的方式做出反应的智能机器。该领域的研究包括机器人、
语言识别、图像识别、自然语言处理和专家系统等。人工智能从诞生以来，理论和技术日益成熟，
应用领域也不断扩大，可以设想，未来人工智能带来的科技产品，将会是人类智慧的"容器"。
人工智能可以对人的意识、思维的信息过程的模拟。人工智能不是人的智能，但能像人那样思考、
也可能超过人的智能。
"""splitter = FixedLengthSplitter(chunk_size=100, chunk_overlap=20)
chunks = splitter.split_text(sample_text)for i, chunk in enumerate(chunks):    print(f"块 {i+1}: {chunk[:50]}...")

class SemanticSplitter:
    def __init__(self, max_chunk_size=500):
        self.max_chunk_size = max_chunk_size    
    def split_by_sentences(self, text: str) -> List[str]:
        sentences = re.split(r'[。！？]', text)
        sentences = [s.strip() for s in sentences if s.strip()]
        
        chunks = []
        current_chunk = ""
        
        for sentence in sentences:            if len(current_chunk) + len(sentence) <= self.max_chunk_size:
                current_chunk += sentence + "。"
            else:                if current_chunk:
                    chunks.append(current_chunk.strip())
                current_chunk = sentence + "。"
        
        if current_chunk:
            chunks.append(current_chunk.strip())            
        return chunks

semantic_splitter = SemanticSplitter(max_chunk_size=150)
semantic_chunks = semantic_splitter.split_by_sentences(sample_text)for i, chunk in enumerate(semantic_chunks):    print(f"语义块 {i+1}: {chunk}")

class RecursiveSplitter:
    def __init__(self, chunk_size=500, chunk_overlap=50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.separators = ["\n\n", "\n", "。", "，", " "]    
    def split_text(self, text: str) -> List[str]:
        return self._split_text_recursive(text, self.separators)    
    def _split_text_recursive(self, text: str, separators: List[str]) -> List[str]:
        if len(text) <= self.chunk_size:            return [text]        
        if not separators:            return self._split_by_length(text)
        
        separator = separators[0]
        remaining_separators = separators[1:]
        
        splits = text.split(separator)
        chunks = []
        current_chunk = ""
        
        for split in splits:            if len(current_chunk) + len(split) + len(separator) <= self.chunk_size:
                current_chunk += split + separator            else:                if current_chunk:
                    chunks.extend(self._split_text_recursive(current_chunk, remaining_separators))
                current_chunk = split + separator        
        if current_chunk:
            chunks.extend(self._split_text_recursive(current_chunk, remaining_separators))        
        return chunks    
    def _split_by_length(self, text: str) -> List[str]:
        chunks = []
        start = 0
        while start < len(text):
            end = start + self.chunk_size
            chunks.append(text[start:end])
            start = end - self.chunk_overlap        return chunks

class TextEmbedding:
    def __init__(self, client):
        self.client = client
        self.model = "embedding-v1"
    
    def get_embedding(self, text: str) -> List[float]:
        response = self.client.embeddings.create(
            model=self.model,            input=[text]
        )        return response.data[0].embedding    
    def get_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
        response = self.client.embeddings.create(
            model=self.model,            input=texts
        )        return [item.embedding for item in response.data]

embedder = TextEmbedding(client)

test_texts = [    "人工智能是计算机科学的分支",    "机器学习是AI的重要组成部分",    "今天天气很好"]

embeddings = embedder.get_embeddings_batch(test_texts)print(f"向量维度: {len(embeddings[0])}")print(f"第一个文本的向量前10个元素: {embeddings[0][:10]}")

def cosine_similarity_manual(vec1: List[float], vec2: List[float]) -> float:
    vec1 = np.array(vec1)
    vec2 = np.array(vec2)
    
    dot_product = np.dot(vec1, vec2)
    norm1 = np.linalg.norm(vec1)
    norm2 = np.linalg.norm(vec2)    
    return dot_product / (norm1 * norm2)

similarity_01 = cosine_similarity_manual(embeddings[0], embeddings[1])
similarity_02 = cosine_similarity_manual(embeddings[0], embeddings[2])print(f"'人工智能是计算机科学的分支' 与 '机器学习是AI的重要组成部分' 的相似度: {similarity_01:.4f}")print(f"'人工智能是计算机科学的分支' 与 '今天天气很好' 的相似度: {similarity_02:.4f}")

class SimpleVectorStore:
    def __init__(self, embedder):
        self.embedder = embedder
        self.documents = []
        self.embeddings = []    
    def add_documents(self, texts: List[str]):
        embeddings = self.embedder.get_embeddings_batch(texts)
        self.documents.extend(texts)
        self.embeddings.extend(embeddings)    
    def similarity_search(self, query: str, k: int = 3) -> List[Dict]:
        query_embedding = self.embedder.get_embedding(query)
        
        similarities = []        for i, doc_embedding in enumerate(self.embeddings):
            similarity = cosine_similarity_manual(query_embedding, doc_embedding)
            similarities.append({                'index': i,                'document': self.documents[i],                'similarity': similarity
            })
        
        similarities.sort(key=lambda x: x['similarity'], reverse=True)        return similarities[:k]

knowledge_base = [    "人工智能（AI）是计算机科学的一个分支，致力于创建能够执行通常需要人类智能的任务的系统。",    "机器学习是人工智能的一个子集，它使计算机能够从数据中学习而无需明确编程。",    "深度学习是机器学习的一个分支，使用多层神经网络来模拟人脑的学习过程。",    "自然语言处理（NLP）是AI的一个领域，专注于计算机与人类语言之间的交互。",    "计算机视觉是AI的一个分支，旨在让计算机能够识别和理解视觉信息。",    "强化学习是机器学习的一种类型，智能体通过与环境交互来学习最优行为。",    "神经网络是受生物神经系统启发的计算模型，是深度学习的基础。",    "大语言模型（LLM）是基于Transformer架构的深度学习模型，能够理解和生成人类语言。"]

vector_store = SimpleVectorStore(embedder)
vector_store.add_documents(knowledge_base)

query = "什么是深度学习？"search_results = vector_store.similarity_search(query, k=3)print(f"查询: {query}\n")for i, result in enumerate(search_results):    print(f"结果 {i+1} (相似度: {result['similarity']:.4f}):")    print(f"  {result['document']}\n")

class RetrievalEvaluator:
    def __init__(self, vector_store):
        self.vector_store = vector_store    
    def evaluate_retrieval(self, test_cases: List[Dict]) -> Dict:
        total_precision = 0
        total_recall = 0
        total_cases = len(test_cases)        
        for case in test_cases:
            query = case['query']
            relevant_docs = set(case['relevant_docs'])
            
            retrieved_results = self.vector_store.similarity_search(query, k=5)
            retrieved_docs = set(result['index'] for result in retrieved_results)            
            if retrieved_docs:
                precision = len(relevant_docs & retrieved_docs) / len(retrieved_docs)
                recall = len(relevant_docs & retrieved_docs) / len(relevant_docs)            else:
                precision = recall = 0
            
            total_precision += precision
            total_recall += recall
        
        avg_precision = total_precision / total_cases
        avg_recall = total_recall / total_cases
        f1_score = 2 * (avg_precision * avg_recall) / (avg_precision + avg_recall) if (avg_precision + avg_recall) > 0 else 0
        
        return {            'precision': avg_precision,            'recall': avg_recall,            'f1_score': f1_score
        }

class SimpleReranker:
    def __init__(self, embedder):
        self.embedder = embedder    
    def rerank(self, query: str, documents: List[str], top_k: int = 3) -> List[Dict]:
        query_embedding = self.embedder.get_embedding(query)
        
        scored_docs = []        for i, doc in enumerate(documents):
            doc_embedding = self.embedder.get_embedding(doc)
            similarity = cosine_similarity_manual(query_embedding, doc_embedding)
            
            keyword_score = self._keyword_overlap_score(query, doc)
            
            combined_score = 0.7 * similarity + 0.3 * keyword_score
            
            scored_docs.append({                'index': i,                'document': doc,                'similarity': similarity,                'keyword_score': keyword_score,                'combined_score': combined_score
            })
        
        scored_docs.sort(key=lambda x: x['combined_score'], reverse=True)        return scored_docs[:top_k]    
    def _keyword_overlap_score(self, query: str, document: str) -> float:
        query_words = set(query.split())
        doc_words = set(document.split())        
        if not query_words:            return 0
        
        overlap = len(query_words & doc_words)        return overlap / len(query_words)

reranker = SimpleReranker(embedder)

query = "机器学习和深度学习的关系"candidates = [    "机器学习是人工智能的一个子集，它使计算机能够从数据中学习而无需明确编程。",    "深度学习是机器学习的一个分支，使用多层神经网络来模拟人脑的学习过程。",    "自然语言处理（NLP）是AI的一个领域，专注于计算机与人类语言之间的交互。",
]

reranked_results = reranker.rerank(query, candidates)print(f"查询: {query}\n")for i, result in enumerate(reranked_results):    print(f"排名 {i+1}:")    print(f"  文档: {result['document']}")    print(f"  语义相似度: {result['similarity']:.4f}")    print(f"  关键词匹配度: {result['keyword_score']:.4f}")    print(f"  综合得分: {result['combined_score']:.4f}\n")

class RAGSystem:
    def __init__(self, client, embedder, vector_store):
        self.client = client
        self.embedder = embedder
        self.vector_store = vector_store
        self.reranker = SimpleReranker(embedder)    
    def query(self, question: str, model: str = "ernie-4.5-turbo-128k-preview", top_k: int = 3) -> Dict:
        retrieved_docs = self.vector_store.similarity_search(question, k=top_k)        
        if not retrieved_docs:            return {                'answer': '抱歉，我没有找到相关信息来回答您的问题。',                'sources': []
            }
        
        context = "\n".join([doc['document'] for doc in retrieved_docs])
        
        prompt = f"""基于以下背景信息回答问题：

背景信息：{context}问题：{question}请基于上述背景信息回答问题。如果背景信息不足以回答问题，请说明。"""

        response = self.client.chat.completions.create(
            model=model,
            messages=[
                {"role": "user", "content": prompt}
            ],
            temperature=0.1
        )        
        return {            'answer': response.choices[0].message.content,            'sources': retrieved_docs,            'context': context
        }

rag_system = RAGSystem(client, embedder, vector_store)

test_questions = [    "什么是深度学习？",    "机器学习和人工智能有什么区别？",    "强化学习是如何工作的？"]for question in test_questions:    print(f"问题: {question}")    print("-" * 50)
    
    result = rag_system.query(question)    
    print("回答:")    print(result['answer'])    print("\n参考来源:")    for i, source in enumerate(result['sources']):        print(f"{i+1}. {source['document']} (相似度: {source['similarity']:.4f})")    print("\n" + "="*80 + "\n")

class ConversationalRAG(RAGSystem):
    def __init__(self, client, embedder, vector_store):
        super().__init__(client, embedder, vector_store)
        self.conversation_history = []    
    def chat(self, question: str, model: str = "ernie-4.5-turbo-128k-preview") -> Dict:
        retrieved_docs = self.vector_store.similarity_search(question, k=3)
        context = "\n".join([doc['document'] for doc in retrieved_docs])
        
        messages = [
            {"role": "system", "content": "你是一个AI助手，请基于提供的背景信息回答问题。"}
        ]        
        for history in self.conversation_history[-5:]:
            messages.append({"role": "user", "content": history['question']})
            messages.append({"role": "assistant", "content": history['answer']})
        
        current_prompt = f"""背景信息：{context}问题：{question}"""
        
        messages.append({"role": "user", "content": current_prompt})
        
        response = self.client.chat.completions.create(
            model=model,
            messages=messages,
            temperature=0.1
        )
        
        answer = response.choices[0].message.content
        
        self.conversation_history.append({            'question': question,            'answer': answer,            'context': context
        })        
        return {            'answer': answer,            'sources': retrieved_docs,            'context': context
        }    
    def clear_history(self):
        self.conversation_history = []

class QueryExpander:
    def __init__(self, client):
        self.client = client    
    def expand_query(self, original_query: str) -> List[str]:
        prompt = f"""请为以下查询生成3个相关的扩展查询，用于信息检索：

原始查询：{original_query}扩展查询应该：
1. 保持与原始查询的语义相关性
2. 使用不同的关键词和表达方式
3. 可能涵盖更具体或更一般的方面

请以JSON格式返回扩展查询列表：
{{"expanded_queries": ["查询1", "查询2", "查询3"]}}"""

        response = self.client.chat.completions.create(
            model="ernie-4.5-turbo-128k",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3,
            response_format={"type": "json_object"}
        )        
        try:
            result = json.loads(response.choices[0].message.content)            return result.get("expanded_queries", [original_query])        except:            return [original_query]

query_expander = QueryExpander(client)

original_query = "深度学习的应用"expanded_queries = query_expander.expand_query(original_query)print(f"原始查询: {original_query}")print("扩展查询:")for i, query in enumerate(expanded_queries, 1):    print(f"  {i}. {query}")

class CachedRAGSystem(RAGSystem):
    def __init__(self, client, embedder, vector_store, cache_size=100):
        super().__init__(client, embedder, vector_store)
        self.query_cache = {}
        self.cache_size = cache_size    
    def query(self, question: str, model: str = "ernie-4.5-turbo-128k-preview", top_k: int = 3) -> Dict:
        cache_key = f"{question}_{model}_{top_k}"
        
        if cache_key in self.query_cache:            print("从缓存中获取结果")            return self.query_cache[cache_key]
        
        result = super().query(question, model, top_k)        
        if len(self.query_cache) >= self.cache_size:
            oldest_key = next(iter(self.query_cache))            del self.query_cache[oldest_key]
        
        self.query_cache[cache_key] = result        return result

class BatchRAGProcessor:
    def __init__(self, rag_system):
        self.rag_system = rag_system    
    def process_batch(self, questions: List[str], batch_size: int = 5) -> List[Dict]:
        results = []        
        for i in range(0, len(questions), batch_size):
            batch = questions[i:i + batch_size]            print(f"处理批次 {i//batch_size + 1}: {len(batch)} 个问题")
            
            batch_results = []            for question in batch:
                result = self.rag_system.query(question)
                batch_results.append(result)
            
            results.extend(batch_results)        
        return results

tech_docs = [    "Python是一种高级编程语言，以其简洁的语法和强大的库生态系统而闻名。",    "Django是一个高级Python Web框架，鼓励快速开发和干净、实用的设计。",    "Flask是一个微型Web框架，用Python编写，设计简单易用。",    "NumPy是Python中科学计算的基础包，提供了强大的N维数组对象。",    "Pandas是一个数据分析和操作库，提供了数据结构和数据分析工具。",    "Scikit-learn是一个机器学习库，提供了简单高效的数据挖掘和数据分析工具。",    "TensorFlow是一个开源机器学习框架，由Google开发。",    "PyTorch是一个开源机器学习库，基于Torch库开发。"]

tech_vector_store = SimpleVectorStore(embedder)
tech_vector_store.add_documents(tech_docs)
tech_rag = RAGSystem(client, embedder, tech_vector_store)

tech_questions = [    "Python有哪些特点？",    "Django和Flask有什么区别？", 
    "推荐一些Python机器学习库"]print("=== 技术文档问答系统演示 ===\n")for question in tech_questions:
    result = tech_rag.query(question)    print(f"Q: {question}")    print(f"A: {result['answer']}\n")

class MultiModalRAGSystem:
    def __init__(self, client, text_embedder):
        self.client = client
        self.text_embedder = text_embedder
        self.text_store = SimpleVectorStore(text_embedder)        
    def add_text_documents(self, texts: List[str]):
        self.text_store.add_documents(texts)    
    def query_with_context(self, question: str, context_type: str = "text") -> Dict:
        if context_type == "text":            return self._query_text(question)        else:            return {"answer": "暂不支持此类型的查询", "sources": []}    
    def _query_text(self, question: str) -> Dict:
        retrieved_docs = self.text_store.similarity_search(question, k=3)
        context = "\n".join([doc['document'] for doc in retrieved_docs])
        
        prompt = f"""基于以下信息回答问题：

信息：{context}问题：{question}请提供准确、有用的回答。"""

        response = self.client.chat.completions.create(
            model="ernie-4.5-turbo-vl-preview",
            messages=[{"role": "user", "content": prompt}]
        )        
        return {            'answer': response.choices[0].message.content,            'sources': retrieved_docs
        }