初探RAG

源码

核心工作流程

1. 读取文件的内容
2. 将内容保存在向量数据库
3. 检索向量数据库用户的问题
4. 用户问题 + 上下文【向量数据】 => LLM

读取文件内容【pdf为例】

from pdfminer.high_level import extract_pages
from pdfminer.layout import LTTextContainer

class PDFFileLoader():
    def __init__(self, file) -> None:
        self.paragraphs = self.extract_text_from_pdf(file)
        # i = 1
        # for para in self.paragraphs[:10]:
        #     print(f"========= 第{i}段 ==========")
        #     print(para+"\n")
        #     i += 1
    
    def getParagraphs(self):
        return self.paragraphs

    ################################# 文档的加载与切割 ############################
    def extract_text_from_pdf(self, filename, page_numbers=None):
        '''从 PDF 文件中（按指定页码）提取文字'''
        paragraphs = []
        buffer = ''
        full_text = ''
        # 提取全部文本
        for i, page_layout in enumerate(extract_pages(filename)):
            # 如果指定了页码范围，跳过范围外的页
            if page_numbers is not None and i not in page_numbers:
                continue
            for element in page_layout:
                if isinstance(element, LTTextContainer):
                    full_text += element.get_text() + '\n'
        
        # 段落分割
        lines = full_text.split('。\n')
        for text in lines:
            buffer = text.replace('\n', ' ')
            
            if buffer:
                paragraphs.append(buffer)
                buffer = ''
                row_count = 0
                
        if buffer:
            paragraphs.append(buffer)
        return paragraphs

向量数据

在保存到向量数据库的时候，需要对文件的内容处理。这可以使用OpenAI、以及开源的模型做处理。本文是用开源的模型。

# from openai import OpenAI
from langchain.embeddings import HuggingFaceBgeEmbeddings
import os
# 加载环境变量
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())  # 读取本地 .env 文件，里面定义了 OPENAI_API_KEY

# client = OpenAI()


# 加载中文embedding模型 开源模型
bge_embeddings = HuggingFaceBgeEmbeddings(model_name="BAAI/bge-large-zh-v1.5")


def get_embeddings(texts, model="text-embedding-3-small"):
    '''封装 OpenAI 的 Embedding 模型接口'''
    # print(texts)
    # print(model)
    '''open ai 提供的embedding 接口'''
    # data = client.embeddings.create(input=texts, model=model).data
    # return [x.embedding for x in data]
    '''huggingface 提供的embedding 接口'''
    data = bge_embeddings.embed_documents(texts)
    # print(data)
    return data

这样处理会得到一堆向量的数组。就可以入库了

import chromadb
from chromadb.config import Settings

class MyVectorDBConnector:
    def __init__(self, collection_name, embedding_fn):
        chroma_client = chromadb.Client(Settings(allow_reset=True))

        # 为了演示，实际不需要每次 reset()
        chroma_client.reset()

        # 创建一个 collection
        self.collection = chroma_client.get_or_create_collection(name=collection_name)
        self.embedding_fn = embedding_fn

    def add_documents(self, documents):
        '''向 collection 中添加文档与向量'''
        self.collection.add(
            embeddings=self.embedding_fn(documents),  # 每个文档的向量
            documents=documents,  # 文档的原文
            ids=[f"id{i}" for i in range(len(documents))]  # 每个文档的 id
        )

    def search(self, query, top_n):
        '''检索向量数据库'''
        results = self.collection.query(
            query_embeddings=self.embedding_fn([query]),
            n_results=top_n
        )
        return results

结合LLM + RAG

import chromadb
from chromadb.config import Settings

class MyVectorDBConnector:
    def __init__(self, collection_name, embedding_fn):
        chroma_client = chromadb.Client(Settings(allow_reset=True))

        # 为了演示，实际不需要每次 reset()
        chroma_client.reset()

        # 创建一个 collection
        self.collection = chroma_client.get_or_create_collection(name=collection_name)
        self.embedding_fn = embedding_fn

    def add_documents(self, documents):
        '''向 collection 中添加文档与向量'''
        self.collection.add(
            embeddings=self.embedding_fn(documents),  # 每个文档的向量
            documents=documents,  # 文档的原文
            ids=[f"id{i}" for i in range(len(documents))]  # 每个文档的 id
        )

    def search(self, query, top_n):
        '''检索向量数据库'''
        results = self.collection.query(
            query_embeddings=self.embedding_fn([query]),
            n_results=top_n
        )
        return results

以上就是初探RAG的代码。
参考：
Embedding模型的选择
RAG教程