从0掌握 RAG（检索增强生成）系统

第一步：创建一个知识库类（KB）

目标：将长文档分块并生成嵌入向量，便于后续检索。

1.1 定义 KB 类及初始化方法

class KB:
    """
        知识库类，用于管理文档分块和生成嵌入向量。
    """
    def __init__(self, filepath):
        try:
            with open(filepath, 'r', encoding='UTF-8') as f:
                content = f.read()
        except FileNotFoundError:
            raise ValueError(f"File {filepath} not found.")

        self.docs = list(self.split_content(content))
        self.embeds = self.encode(self.docs)

1.2 文本分块方法

@staticmethod
def split_content(content, max_length=128):
    """
        将文档按标点符号分块，块大小不超过 max_length。
    """
    sentences = re.split(r'(?<=[.!?。，！])\s+', content)
    current_chunk = []
    current_length = 0

    for sentence in sentences:
        sentence_length = len(sentence)
        if current_length + sentence_length <= max_length:
            current_chunk.append(sentence)
            current_length += sentence_length
        else:
            yield ' '.join(current_chunk).strip()
            current_chunk = [sentence]
            current_length = sentence_length
    if current_chunk:
        yield ' '.join(current_chunk).strip()

1.3 嵌入向量生成方法

@staticmethod
def encode(texts):
    """
    Encode a list of texts into embeddings.
    """
    embeds = []
    for text in texts:
        try:
            response = ollama.embeddings(model='nomic-embed-text', prompt=text)
            embeds.append(response['embedding'])
        except Exception as e:
            print(f"Error encoding text: {text[:30]}... -> {e}")

    return np.array(embeds)

---

第二步：实现检索功能

目标：通过余弦相似度，从知识库中找到与输入文本最相关的文档块。

2.1 定义相似度计算方法

@staticmethod
def similarity(vec1, vec2):
    """
    计算余弦相似度
    """
    dot_product = np.dot(vec1, vec2)
    norm1, norm2 = np.linalg.norm(vec1), np.linalg.norm(vec2)
    return dot_product / (norm1 * norm2)

2.2 实现检索逻辑

def search(self, text):
    """
    根据输入文本匹配知识库中相似度最高的片段
    """
    query_embed = self.encode([text])[0]
    similarities = np.dot(self.embeds, query_embed) / (
            np.linalg.norm(self.embeds, axis=1) * np.linalg.norm(query_embed)
    )
    max_index = np.argmax(similarities)
    return self.docs[max_index]

---

第三步：构建 RAG 问答类

目标：结合知识库和生成模型，生成基于上下文的回答。

3.1 定义 Rag 类及初始化方法

class Rag:
    """
    用于结合知识库和语言模型的检索增强生成（RAG）类。
    """

    def __init__(self, model, kb: KB):
        self.model = model
        self.kb = kb
        self.prompt_template = """
        任务描述：使用中文回答用户的问题，请根据提供的上下文信息生成精准回答。
        检索上下文：{retrieved_context}
        问题：{user_query}
        请基于检索上下文和问题生成答案。如果检索上下文中没有明确的信息，请礼貌地告诉用户无法回答。
        """

3.2 实现生成回答方法

def chat(self, text):
    """
    根据检索到的上下文和输入问题创建响应
    """
    context = self.kb.search(text)
    prompt = self.prompt_template.format(retrieved_context=context, user_query=text)
    print(f"Generated Prompt:\n{prompt}")

    try:
        response = ollama.chat(self.model, [ollama.Message(role='user', content=prompt),
                                            ollama.Message(role='system', content='你是怀光长纵的智能助理')])
        return response['message']
    except Exception as e:
        print(f"Error during chat generation: {e}")
        return {"content": "抱歉，我无法生成答案。"}

---

第四步：整合并运行

4.1 主程序逻辑

if __name__ == '__main__':
    try:
        kb = KB('../file/test.txt')  # 替换为你的“知识库”路径
        rag = Rag('qwen2.5:0.5b-instruct', kb)
        
        question = '明吉集团是怎么创办的'
        result = rag.chat(question)
        print("回答:", result.get('content', '未生成答案。'))
    except ValueError as e:
        print(e)

---

测试与输出

问题：明吉集团是怎么创办的？

控制台日志：

生成的提示语:
根据以下提供的内容回答问题：
内容：明吉集团成立于2001年，致力于科技创新。
问题：明吉集团是怎么创办的
回答需简洁明了。

回答: 明吉集团成立于2001年，致力于科技创新。

---

实现过程：

1. 构建知识库并生成嵌入向量。
2. 实现基于余弦相似度的检索功能。
3. 将检索结果结合语言模型生成答案。

可扩展方向：

1. 多文档支持：支持从多个知识库中检索答案。
2. 模型优化：使用更强大的语言模型。
3. 性能改进：通过向量索引（如 FAISS）加速检索。