那接下來，讓我們梳理一下 RAG 的流程是什么樣的呢？

那也就是下圖所示的流程。

2. 向量化

首先讓我們來動手實現一個向量化的類，這是RAG架構的基礎。向量化的類主要是用來將文檔片段向量化，將一段文本映射為一個向量。

那首先我們要設置一個 Embedding 基類，這樣我們再用其他的模型的時候，只需要繼承這個基類，然后在此基礎上進行修改即可，方便代碼擴展。

class BaseEmbeddings:

    """

    Base class for embeddings

    """

    def __init__(self, path: str, is_api: bool) -> None:

        self.path = path

        self.is_api = is_api



    def get_embedding(self, text: str, model: str) -> List[float]:

        raise NotImplementedError



    @classmethod

    def cosine_similarity(cls, vector1: List[float], vector2: List[float]) -> float:

        """

        calculate cosine similarity between two vectors

        """

        dot_product = np.dot(vector1, vector2)

        magnitude = np.linalg.norm(vector1) * np.linalg.norm(vector2)

        if not magnitude:

            return 0

        return dot_product / magnitude

觀察一下BaseEmbeddings基類都有什么方法，首先有一個get_embedding方法，這個方法是用來獲取文本的向量表示的，然后有一個cosine_similarity方法，這個方法是用來計算兩個向量之間的余弦相似度的。其次在初始化類的時候設置了，模型的路徑或者是否是API模型。比如使用OpenAI的Embedding API的話就需要設置self.is_api=Ture。

繼承BaseEmbeddings類的話，就只需要編寫get_embedding方法即可，cosine_similarity方法會被繼承下來，直接用就行。這就是編寫基類的好處。

class OpenAIEmbedding(BaseEmbeddings):

    """

    class for OpenAI embeddings

    """

    def __init__(self, path: str = '', is_api: bool = True) -> None:

        super().__init__(path, is_api)

        if self.is_api:

            from openai import OpenAI

            self.client = OpenAI()

            self.client.api_key = os.getenv("OPENAI_API_KEY")

            self.client.base_url = os.getenv("OPENAI_BASE_URL")



    def get_embedding(self, text: str, model: str = "text-embedding-3-large") -> List[float]:

        if self.is_api:

            text = text.replace("\n", " ")

            return self.client.embeddings.create(input=[text], model=model).data[0].embedding

        else:

            raise NotImplementedError

3. 文檔加載和切分

接下來我們來實現一個文檔加載和切分的類，這個類主要是用來加載文檔并切分成文檔片段。

那我們都需要切分什么文檔呢？這個文檔可以是一篇文章，一本書，一段對話，一段代碼等等。這個文檔的內容可以是任何的，只要是文本就行。比如：pdf文件、md文件、txt文件等等。

這里只展示一部分內容了，完整的代碼可以在 RAG/utils.py 文件中找到。在這個代碼中可以看到，能加載的文件類型有：pdf、md、txt，只需要編寫對應的函數即可。

def read_file_content(cls, file_path: str):

    # 根據文件擴展名選擇讀取方法

    if file_path.endswith('.pdf'):

        return cls.read_pdf(file_path)

    elif file_path.endswith('.md'):

        return cls.read_markdown(file_path)

    elif file_path.endswith('.txt'):

        return cls.read_text(file_path)

    else:

        raise ValueError("Unsupported file type")

那我們把文件內容都讀取之后，還需要切分呀！那怎么切分呢，OK，接下來咱們就按 Token 的長度來切分文檔。我們可以設置一個最大的 Token 長度，然后根據這個最大的 Token 長度來切分文檔。這樣切分出來的文檔片段就是一個一個的差不多相同長度的文檔片段了。

不過在切分的時候要注意，片段與片段之間最好要有一些重疊的內容，這樣才能保證檢索的時候能夠檢索到相關的文檔片段。還有就是切分文檔的時候最好以句子為單位，也就是按 \n 進行粗切分，這樣可以基本保證句子內容是完整的。

def get_chunk(cls, text: str, max_token_len: int = 600, cover_content: int = 150):

    chunk_text = []



    curr_len = 0

    curr_chunk = ''



    lines = text.split('\n')  # 假設以換行符分割文本為行



    for line in lines:

        line = line.replace(' ', '')

        line_len = len(enc.encode(line))

        if line_len > max_token_len:

            print('warning line_len = ', line_len)

        if curr_len + line_len <= max_token_len:

            curr_chunk += line

            curr_chunk += '\n'

            curr_len += line_len

            curr_len += 1

        else:

            chunk_text.append(curr_chunk)

            curr_chunk = curr_chunk[-cover_content:]+line

            curr_len = line_len + cover_content



    if curr_chunk:

        chunk_text.append(curr_chunk)



    return chunk_text

4. 數據庫 && 向量檢索

上面，我們做好了文檔切分，也做好了 Embedding 模型的加載。那接下來就得設計一個向量數據庫用來存放文檔片段和對應的向量表示了。

還有就是也要設計一個檢索模塊，用來根據 Query （問題）檢索相關的文檔片段。OK，我們沖沖沖！

一個數據庫對于最小RAG架構來說，需要實現幾個功能呢？

嗯嗯，以上四個模塊就是一個最小的RAG結構數據庫需要實現的功能，具體代碼可以在 RAG/VectorBase.py 文件中找到。

class VectorStore:

    def __init__(self, document: List[str] = ['']) -> None:

        self.document = document



    def get_vector(self, EmbeddingModel: BaseEmbeddings) -> List[List[float]]:

        # 獲得文檔的向量表示

        pass



    def persist(self, path: str = 'storage'):

        # 數據庫持久化，本地保存

        pass



    def load_vector(self, path: str = 'storage'):

        # 從本地加載數據庫

        pass



    def query(self, query: str, EmbeddingModel: BaseEmbeddings, k: int = 1) -> List[str]:

        # 根據問題檢索相關的文檔片段

        pass

那讓我們來看一下， query 方法具體是怎么實現的呢？

首先先把用戶提出的問題向量化，然后去數據庫中檢索相關的文檔片段，最后返回檢索到的文檔片段。可以看到咱們在向量檢索的時候僅使用 Numpy 進行加速，代碼非常容易理解和修改。

主要是方便改寫和大家理解，并沒有使用成熟的數據庫，這樣可以更好地理解RAG的原理。

def query(self, query: str, EmbeddingModel: BaseEmbeddings, k: int = 1) -> List[str]:

    query_vector = EmbeddingModel.get_embedding(query)

    result = np.array([self.get_similarity(query_vector, vector)

                        for vector in self.vectors])

    return np.array(self.document)[result.argsort()[-k:][::-1]].tolist()

5. 大模型模塊

那就來到了最后一個模塊了，大模型模塊。這個模塊主要是用來根據檢索出來的文檔回答用戶的問題。

一樣的，我們還是先實現一個基類，這樣我們在遇到其他的自己感興趣的模型就可以快速的擴展了。

class BaseModel:

    def __init__(self, path: str = '') -> None:

        self.path = path



    def chat(self, prompt: str, history: List[dict], content: str) -> str:

        pass



    def load_model(self):

        pass

BaseModel 包含了兩個方法，chat和load_model，如果使用API模型，比如OpenAI的話，那就不需要load_model方法，如果你要本地化運行的話，那還是會選擇使用開源模型，那就需要load_model方法啦。

這里咱們以 InternLM2-chat-7B 模型為例

class InternLMChat(BaseModel):

    def __init__(self, path: str = '') -> None:

        super().__init__(path)

        self.load_model()



    def chat(self, prompt: str, history: List = [], content: str='') -> str:

        prompt = PROMPT_TEMPLATE['InternLM_PROMPT_TEMPALTE'].format(question=prompt, context=content)

        response, history = self.model.chat(self.tokenizer, prompt, history)

        return response



    def load_model(self):

        import torch

        from transformers import AutoTokenizer, AutoModelForCausalLM

        self.tokenizer = AutoTokenizer.from_pretrained(self.path, trust_remote_code=True)

        self.model = AutoModelForCausalLM.from_pretrained(self.path, torch_dtype=torch.float16, trust_remote_code=True).cuda()

可以用一個字典來保存所有的prompt，這樣比較好維護。

PROMPT_TEMPLATE = dict(

    InternLM_PROMPT_TEMPALTE="""先對上下文進行內容總結,再使用上下文來回答用戶的問題。如果你不知道答案，就說你不知道。總是使用中文回答。

        問題: {question}

        可參考的上下文：

        ···

        {context}

        ···

        如果給定的上下文無法讓你做出回答，請回答數據庫中沒有這個內容，你不知道。

        有用的回答:"""

)

那這樣的話，我們就可以利用InternLM2模型來做RAG啦！

6. LLM Tiny-RAG Demo

那接下來，我們就來看一下Tiny-RAG的Demo吧！

from RAG.VectorBase import VectorStore

from RAG.utils import ReadFiles

from RAG.LLM import OpenAIChat, InternLMChat

from RAG.Embeddings import JinaEmbedding, ZhipuEmbedding



沒有保存數據庫

docs = ReadFiles('./data').get_content(max_token_len=600, cover_content=150) # 獲得data目錄下的所有文件內容并分割

vector = VectorStore(docs)

embedding = ZhipuEmbedding() # 創建EmbeddingModel

vector.get_vector(EmbeddingModel=embedding)

vector.persist(path='storage') # 將向量和文檔內容保存到storage目錄下，下次再用就可以直接加載本地的數據庫



question = 'git的原理是什么？'



content = vector.query(question, model='zhipu', k=1)[0]

chat = InternLMChat(path='model_path')

print(chat.chat(question, [], content))

當然我們也可以從本地加載已經處理好的數據庫，畢竟我們在上面的數據庫環節已經寫過這個功能啦。

文章轉自微信公眾號@Datawhale

上一篇：

為 GraphQL API 構建身份驗證和授權的步驟

下一篇：

如何構建 Node.js 中間件來記錄 HTTP API 請求和響應
#你可能也喜歡這些API文章!

我們有何不同？

API服務商零注冊

多API并行試用

數據驅動選型，提升決策效率

查看全部API→
??

熱門場景實測，選對API

#AI文本生成大模型API

對比大模型API的內容創意新穎性、情感共鳴力、商業轉化潛力

25個渠道
一鍵對比試用API 限時免費

#AI深度推理大模型API

對比大模型API的邏輯推理準確性、分析深度、可視化建議合理性

10個渠道
一鍵對比試用API 限時免費